chore: refactor logic for tool sharing

ibuildthecloud · ibuildthecloud · commit c348a29ccbf2 · 2024-08-28T23:22:29.000-07:00
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 {
diff --git a/pkg/types/tool.go b/pkg/types/tool.go

Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ func NewContext(ctx context.Context, prg *types.Program, input string) (Context,`
`204`	`204`	`Input: input,`
`205`	`205`	`}`
`206`	`206`
`207`		`- agentGroup, err := callCtx.Tool.GetAgents(*prg)`
	`207`	`+ agentGroup, err := callCtx.Tool.GetToolsByType(prg, types.ToolTypeAgent)`
`208`	`208`	`if err != nil {`
`209`	`209`	`return callCtx, err`
`210`	`210`	`}`
`@@ -225,7 +225,7 @@ func (c *Context) SubCallContext(ctx context.Context, input, toolID, callID stri`
`225`	`225`	`callID = counter.Next()`
`226`	`226`	`}`
`227`	`227`
`228`		`- agentGroup, err := c.Tool.GetNextAgentGroup(*c.Program, c.AgentGroup, toolID)`
	`228`	`+ agentGroup, err := c.Tool.GetNextAgentGroup(c.Program, c.AgentGroup, toolID)`
`229`	`229`	`if err != nil {`
`230`	`230`	`return Context{}, err`
`231`	`231`	`}`
`@@ -272,7 +272,7 @@ func populateMessageParams(ctx Context, completion *types.CompletionRequest, too`
`272`	`272`	`}`
`273`	`273`
`274`	`274`	`var err error`
`275`		`- completion.Tools, err = tool.GetCompletionTools(*ctx.Program, ctx.AgentGroup...)`
	`275`	`+ completion.Tools, err = tool.GetChatCompletionTools(*ctx.Program, ctx.AgentGroup...)`
`276`	`276`	`if err != nil {`
`277`	`277`	`return err`
`278`	`278`	`}`
Original file line number	Diff line number	Diff line change
`@@ -5,10 +5,11 @@ import (`
`5`	`5`	`"fmt"`
`6`	`6`
`7`	`7`	`"github.com/gptscript-ai/gptscript/pkg/engine"`
	`8`	`+ "github.com/gptscript-ai/gptscript/pkg/types"`
`8`	`9`	`)`
`9`	`10`
`10`	`11`	`func (r *Runner) handleInput(callCtx engine.Context, monitor Monitor, env []string, input string) (string, error) {`
`11`		`- inputToolRefs, err := callCtx.Tool.GetInputFilterTools(*callCtx.Program)`
	`12`	`+ inputToolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeInput)`
`12`	`13`	`if err != nil {`
`13`	`14`	`return "", err`
`14`	`15`	`}`
Original file line number	Diff line number	Diff line change
`@@ -6,10 +6,11 @@ import (`
`6`	`6`	`"fmt"`
`7`	`7`
`8`	`8`	`"github.com/gptscript-ai/gptscript/pkg/engine"`
	`9`	`+ "github.com/gptscript-ai/gptscript/pkg/types"`
`9`	`10`	`)`
`10`	`11`
`11`	`12`	`func (r Runner) handleOutput(callCtx engine.Context, monitor Monitor, env []string, state State, retErr error) (*State, error) {`
`12`		`- outputToolRefs, err := callCtx.Tool.GetOutputFilterTools(*callCtx.Program)`
	`13`	`+ outputToolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeOutput)`
`13`	`14`	`if err != nil {`
`14`	`15`	`return nil, err`
`15`	`16`	`}`
Original file line number	Diff line number	Diff line change
`@@ -330,7 +330,7 @@ func getToolRefInput(prg *types.Program, ref types.ToolReference, input string)`
`330`	`330`	`}`
`331`	`331`
`332`	`332`	`func (r Runner) getContext(callCtx engine.Context, state State, monitor Monitor, env []string, input string) (result []engine.InputContext, _ error) {`
`333`		`- toolRefs, err := callCtx.Tool.GetContextTools(*callCtx.Program)`
	`333`	`+ toolRefs, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeContext)`
`334`	`334`	`if err != nil {`
`335`	`335`	`return nil, err`
`336`	`336`	`}`
`@@ -387,7 +387,7 @@ func (r Runner) start(callCtx engine.Context, state State, monitor Monitor, en`
`387`	`387`	`return nil, err`
`388`	`388`	`}`
`389`	`389`
`390`		`- credTools, err := callCtx.Tool.GetCredentialTools(*callCtx.Program, callCtx.AgentGroup)`
	`390`	`+ credTools, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeCredential)`
`391`	`391`	`if err != nil {`
`392`	`392`	`return nil, err`
`393`	`393`	`}`
`@@ -503,7 +503,7 @@ func (r *Runner) resume(callCtx engine.Context, monitor Monitor, env []string, s`
`503`	`503`	`progress, progressClose := streamProgress(&callCtx, monitor)`
`504`	`504`	`defer progressClose()`
`505`	`505`
`506`		`- credTools, err := callCtx.Tool.GetCredentialTools(*callCtx.Program, callCtx.AgentGroup)`
	`506`	`+ credTools, err := callCtx.Tool.GetToolsByType(callCtx.Program, types.ToolTypeCredential)`
`507`	`507`	`if err != nil {`
`508`	`508`	`return nil, err`
`509`	`509`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,7 @@`
`3`	`3`	`"content": [`
`4`	`4`	`{`
`5`	`5`	`"toolCall": {`
`6`		`- "index": 1,`
	`6`	`+ "index": 0,`
`7`	`7`	`"id": "call_3",`
`8`	`8`	`"function": {`
`9`	`9`	`"name": "agent3"`