Merge pull request #831 from ibuildthecloud/share-context-bug

chore: refactor logic for tool sharing
gptscript-ai · Aug 29, 2024 · bfe96cf · bfe96cf
2 parents fbb8f5d + 33741b1
commit bfe96cf
Show file tree

Hide file tree

Showing 17 changed files with 211 additions and 313 deletions.
diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
@@ -204,7 +204,7 @@ func NewContext(ctx context.Context, prg *types.Program, input string) (Context,
 		Input:   input,
 	}
 
-	agentGroup, err := callCtx.Tool.GetAgents(*prg)
+	agentGroup, err := callCtx.Tool.GetToolsByType(prg, types.ToolTypeAgent)
 	if err != nil {
 		return callCtx, err
 	}
@@ -225,7 +225,7 @@ func (c *Context) SubCallContext(ctx context.Context, input, toolID, callID stri
 		callID = counter.Next()
 	}
 
-	agentGroup, err := c.Tool.GetNextAgentGroup(*c.Program, c.AgentGroup, toolID)
+	agentGroup, err := c.Tool.GetNextAgentGroup(c.Program, c.AgentGroup, toolID)
 	if err != nil {
 		return Context{}, err
 	}
@@ -272,7 +272,7 @@ func populateMessageParams(ctx Context, completion *types.CompletionRequest, too
 	}
 
 	var err error
-	completion.Tools, err = tool.GetCompletionTools(*ctx.Program, ctx.AgentGroup...)
+	completion.Tools, err = tool.GetChatCompletionTools(*ctx.Program, ctx.AgentGroup...)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/runner/input.go b/pkg/runner/input.go
@@ -5,10 +5,11 @@ import (
 	"fmt"
 
 	"github.com/gptscript-ai/gptscript/pkg/engine"
+	"github.com/gptscript-ai/gptscript/pkg/types"
 )
 
 func (r *Runner) handleInput(callCtx engine.Context, monitor Monitor, env []string, input string) (string, error) {
-	inputToolRefs, err := callCtx.Tool.GetInputFilterTools(*callCtx.Program)
+	inputToolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeInput)
 	if err != nil {
 		return "", err
 	}

diff --git a/pkg/runner/output.go b/pkg/runner/output.go
@@ -6,10 +6,11 @@ import (
 	"fmt"
 
 	"github.com/gptscript-ai/gptscript/pkg/engine"
+	"github.com/gptscript-ai/gptscript/pkg/types"
 )
 
 func (r *Runner) handleOutput(callCtx engine.Context, monitor Monitor, env []string, state *State, retErr error) (*State, error) {
-	outputToolRefs, err := callCtx.Tool.GetOutputFilterTools(*callCtx.Program)
+	outputToolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeOutput)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
@@ -330,7 +330,7 @@ func getToolRefInput(prg *types.Program, ref types.ToolReference, input string)
 }
 
 func (r *Runner) getContext(callCtx engine.Context, state *State, monitor Monitor, env []string, input string) (result []engine.InputContext, _ error) {
-	toolRefs, err := callCtx.Tool.GetContextTools(*callCtx.Program)
+	toolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeContext)
 	if err != nil {
 		return nil, err
 	}
@@ -387,7 +387,7 @@ func (r *Runner) start(callCtx engine.Context, state *State, monitor Monitor, en
 		return nil, err
 	}
 
-	credTools, err := callCtx.Tool.GetCredentialTools(*callCtx.Program, callCtx.AgentGroup)
+	credTools, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeCredential)
 	if err != nil {
 		return nil, err
 	}
@@ -503,7 +503,7 @@ func (r *Runner) resume(callCtx engine.Context, monitor Monitor, env []string, s
 	progress, progressClose := streamProgress(&callCtx, monitor)
 	defer progressClose()
 
-	credTools, err := callCtx.Tool.GetCredentialTools(*callCtx.Program, callCtx.AgentGroup)
+	credTools, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeCredential)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/tests/testdata/TestAgentOnly/call2.golden b/pkg/tests/testdata/TestAgentOnly/call2.golden
@@ -4,8 +4,8 @@
   "tools": [
     {
       "function": {
-        "toolID": "testdata/TestAgentOnly/test.gpt:agent1",
-        "name": "agent1",
+        "toolID": "testdata/TestAgentOnly/test.gpt:agent3",
+        "name": "agent3",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {
@@ -19,8 +19,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestAgentOnly/test.gpt:agent3",
-        "name": "agent3",
+        "toolID": "testdata/TestAgentOnly/test.gpt:agent1",
+        "name": "agent1",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {

diff --git a/pkg/tests/testdata/TestAgentOnly/step1.golden b/pkg/tests/testdata/TestAgentOnly/step1.golden
@@ -96,8 +96,8 @@
                 "tools": [
                   {
                     "function": {
-                      "toolID": "testdata/TestAgentOnly/test.gpt:agent1",
-                      "name": "agent1",
+                      "toolID": "testdata/TestAgentOnly/test.gpt:agent3",
+                      "name": "agent3",
                       "parameters": {
                         "properties": {
                           "defaultPromptParameter": {
@@ -111,8 +111,8 @@
                   },
                   {
                     "function": {
-                      "toolID": "testdata/TestAgentOnly/test.gpt:agent3",
-                      "name": "agent3",
+                      "toolID": "testdata/TestAgentOnly/test.gpt:agent1",
+                      "name": "agent1",
                       "parameters": {
                         "properties": {
                           "defaultPromptParameter": {

diff --git a/pkg/tests/testdata/TestAgents/call3-resp.golden b/pkg/tests/testdata/TestAgents/call3-resp.golden
@@ -3,7 +3,7 @@
   "content": [
     {
       "toolCall": {
-        "index": 1,
+        "index": 0,
         "id": "call_3",
         "function": {
           "name": "agent3"

diff --git a/pkg/tests/testdata/TestAgents/call3.golden b/pkg/tests/testdata/TestAgents/call3.golden
@@ -4,8 +4,8 @@
   "tools": [
     {
       "function": {
-        "toolID": "testdata/TestAgents/test.gpt:agent1",
-        "name": "agent1",
+        "toolID": "testdata/TestAgents/test.gpt:agent3",
+        "name": "agent3",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {
@@ -19,8 +19,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestAgents/test.gpt:agent3",
-        "name": "agent3",
+        "toolID": "testdata/TestAgents/test.gpt:agent1",
+        "name": "agent1",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {

diff --git a/pkg/tests/testdata/TestAgents/step1.golden b/pkg/tests/testdata/TestAgents/step1.golden
@@ -178,8 +178,8 @@
                       "tools": [
                         {
                           "function": {
-                            "toolID": "testdata/TestAgents/test.gpt:agent1",
-                            "name": "agent1",
+                            "toolID": "testdata/TestAgents/test.gpt:agent3",
+                            "name": "agent3",
                             "parameters": {
                               "properties": {
                                 "defaultPromptParameter": {
@@ -193,8 +193,8 @@
                         },
                         {
                           "function": {
-                            "toolID": "testdata/TestAgents/test.gpt:agent3",
-                            "name": "agent3",
+                            "toolID": "testdata/TestAgents/test.gpt:agent1",
+                            "name": "agent1",
                             "parameters": {
                               "properties": {
                                 "defaultPromptParameter": {
@@ -222,7 +222,7 @@
                           "content": [
                             {
                               "toolCall": {
-                                "index": 1,
+                                "index": 0,
                                 "id": "call_3",
                                 "function": {
                                   "name": "agent3"
@@ -237,7 +237,7 @@
                     },
                     "pending": {
                       "call_3": {
-                        "index": 1,
+                        "index": 0,
                         "id": "call_3",
                         "function": {
                           "name": "agent3"

diff --git a/pkg/tests/testdata/TestExport/call1-resp.golden b/pkg/tests/testdata/TestExport/call1-resp.golden
@@ -3,7 +3,7 @@
   "content": [
     {
       "toolCall": {
-        "index": 2,
+        "index": 1,
         "id": "call_1",
         "function": {
           "name": "transient"

diff --git a/pkg/tests/testdata/TestExport/call1.golden b/pkg/tests/testdata/TestExport/call1.golden
@@ -18,8 +18,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestExport/parent.gpt:parent-local",
-        "name": "parentLocal",
+        "toolID": "testdata/TestExport/sub/child.gpt:transient",
+        "name": "transient",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {
@@ -33,8 +33,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestExport/sub/child.gpt:transient",
-        "name": "transient",
+        "toolID": "testdata/TestExport/parent.gpt:parent-local",
+        "name": "parentLocal",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {

diff --git a/pkg/tests/testdata/TestExport/call3.golden b/pkg/tests/testdata/TestExport/call3.golden
@@ -18,8 +18,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestExport/parent.gpt:parent-local",
-        "name": "parentLocal",
+        "toolID": "testdata/TestExport/sub/child.gpt:transient",
+        "name": "transient",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {
@@ -33,8 +33,8 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestExport/sub/child.gpt:transient",
-        "name": "transient",
+        "toolID": "testdata/TestExport/parent.gpt:parent-local",
+        "name": "parentLocal",
         "parameters": {
           "properties": {
             "defaultPromptParameter": {
@@ -62,7 +62,7 @@
       "content": [
         {
           "toolCall": {
-            "index": 2,
+            "index": 1,
             "id": "call_1",
             "function": {
               "name": "transient"
@@ -80,7 +80,7 @@
         }
       ],
       "toolCall": {
-        "index": 2,
+        "index": 1,
         "id": "call_1",
         "function": {
           "name": "transient"

diff --git a/pkg/tests/testdata/TestExportContext/call1.golden b/pkg/tests/testdata/TestExportContext/call1.golden
@@ -38,7 +38,7 @@
       "role": "system",
       "content": [
         {
-          "text": "this is from external context\nthis is from context\nThis is from tool"
+          "text": "this is from context\nthis is from external context\nThis is from tool"
         }
       ],
       "usage": {}

diff --git a/pkg/tests/testdata/TestToolRefAll/call1.golden b/pkg/tests/testdata/TestToolRefAll/call1.golden
@@ -18,12 +18,12 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestToolRefAll/test.gpt:none",
-        "name": "none",
+        "toolID": "testdata/TestToolRefAll/test.gpt:agentAssistant",
+        "name": "agentAssistant",
         "parameters": {
           "properties": {
-            "noneArg": {
-              "description": "stuff",
+            "defaultPromptParameter": {
+              "description": "Prompt to send to the tool. This may be an instruction or question.",
               "type": "string"
             }
           },
@@ -33,12 +33,12 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestToolRefAll/test.gpt:agentAssistant",
-        "name": "agent",
+        "toolID": "testdata/TestToolRefAll/test.gpt:none",
+        "name": "none",
         "parameters": {
           "properties": {
-            "defaultPromptParameter": {
-              "description": "Prompt to send to the tool. This may be an instruction or question.",
+            "noneArg": {
+              "description": "stuff",
               "type": "string"
             }
           },
@@ -52,7 +52,7 @@
       "role": "system",
       "content": [
         {
-          "text": "\nShared context\n\nContext Body\nMain tool"
+          "text": "\nContext Body\n\nShared context\nMain tool"
         }
       ],
       "usage": {}

diff --git a/pkg/types/completion.go b/pkg/types/completion.go
@@ -9,15 +9,15 @@ import (
 )
 
 type CompletionRequest struct {
-	Model                string              `json:"model,omitempty"`
-	InternalSystemPrompt *bool               `json:"internalSystemPrompt,omitempty"`
-	Tools                []CompletionTool    `json:"tools,omitempty"`
-	Messages             []CompletionMessage `json:"messages,omitempty"`
-	MaxTokens            int                 `json:"maxTokens,omitempty"`
-	Chat                 bool                `json:"chat,omitempty"`
-	Temperature          *float32            `json:"temperature,omitempty"`
-	JSONResponse         bool                `json:"jsonResponse,omitempty"`
-	Cache                *bool               `json:"cache,omitempty"`
+	Model                string               `json:"model,omitempty"`
+	InternalSystemPrompt *bool                `json:"internalSystemPrompt,omitempty"`
+	Tools                []ChatCompletionTool `json:"tools,omitempty"`
+	Messages             []CompletionMessage  `json:"messages,omitempty"`
+	MaxTokens            int                  `json:"maxTokens,omitempty"`
+	Chat                 bool                 `json:"chat,omitempty"`
+	Temperature          *float32             `json:"temperature,omitempty"`
+	JSONResponse         bool                 `json:"jsonResponse,omitempty"`
+	Cache                *bool                `json:"cache,omitempty"`
 }
 
 func (r *CompletionRequest) GetCache() bool {
@@ -27,7 +27,7 @@ func (r *CompletionRequest) GetCache() bool {
 	return *r.Cache
 }
 
-type CompletionTool struct {
+type ChatCompletionTool struct {
 	Function CompletionFunctionDefinition `json:"function,omitempty"`
 }
 

diff --git a/pkg/types/set.go b/pkg/types/set.go
@@ -19,6 +19,17 @@ func (t *toolRefSet) List() (result []ToolReference, err error) {
 	return result, t.err
 }
 
+func (t *toolRefSet) Contains(value ToolReference) bool {
+	key := toolRefKey{
+		name:   value.Named,
+		toolID: value.ToolID,
+		arg:    value.Arg,
+	}
+
+	_, ok := t.set[key]
+	return ok
+}
+
 func (t *toolRefSet) HasTool(toolID string) bool {
 	for _, ref := range t.set {
 		if ref.ToolID == toolID {