chore: refactor export context behavior

Author: ibuildthecloud

pkg/parser/parser.go

@@ -100,13 +100,13 @@ func isParam(line string, tool *types.Tool) (_ bool, err error) {
 			return false, err
 		}
 		tool.Parameters.Chat = v
-	case "export":
+	case "export", "exporttool", "exports", "exporttools":
 		tool.Parameters.Export = append(tool.Parameters.Export, csv(value)...)
 	case "tool", "tools":
 		tool.Parameters.Tools = append(tool.Parameters.Tools, csv(value)...)
 	case "globaltool", "globaltools":
 		tool.Parameters.GlobalTools = append(tool.Parameters.GlobalTools, csv(value)...)
-	case "exportcontext":
+	case "exportcontext", "exportcontexts":
 		tool.Parameters.ExportContext = append(tool.Parameters.ExportContext, csv(value)...)
 	case "context":
 		tool.Parameters.Context = append(tool.Parameters.Context, csv(value)...)

pkg/tests/testdata/TestExportContext/call1.golden

@@ -4,7 +4,7 @@
   "Tools": [
     {
       "function": {
-        "toolID": "testdata/TestExportContext/test.gpt:21",
+        "toolID": "testdata/TestExportContext/test.gpt:22",
         "name": "subtool",
         "parameters": {
           "properties": {
@@ -22,7 +22,7 @@
     },
     {
       "function": {
-        "toolID": "testdata/TestExportContext/test.gpt:14",
+        "toolID": "testdata/TestExportContext/test.gpt:15",
         "name": "sampletool",
         "description": "sample",
         "parameters": {

pkg/tests/testdata/TestExportContext/test.gpt

@@ -6,6 +6,7 @@ This is from tool
 ---
 name: fromcontext
 export: sampletool
+export context: fromexportcontext
 
 #!/bin/bash
 echo this is from context
@@ -19,7 +20,6 @@ Dummy body
 
 ---
 name: subtool
-export context: fromexportcontext
 
 Dummy body
 

pkg/types/set.go

@@ -0,0 +1,45 @@
+package types
+
+type toolRefKey struct {
+	name   string
+	toolID string
+	arg    string
+}
+
+type toolRefSet struct {
+	set   map[toolRefKey]ToolReference
+	order []toolRefKey
+	err   error
+}
+
+func (t *toolRefSet) List() (result []ToolReference, err error) {
+	for _, k := range t.order {
+		result = append(result, t.set[k])
+	}
+	return result, t.err
+}
+
+func (t *toolRefSet) AddAll(values []ToolReference, err error) {
+	if t.err != nil {
+		t.err = err
+	}
+	for _, v := range values {
+		t.Add(v)
+	}
+}
+
+func (t *toolRefSet) Add(value ToolReference) {
+	key := toolRefKey{
+		name:   value.Named,
+		toolID: value.ToolID,
+		arg:    value.Arg,
+	}
+
+	if _, ok := t.set[key]; !ok {
+		if t.set == nil {
+			t.set = map[toolRefKey]ToolReference{}
+		}
+		t.set[key] = value
+		t.order = append(t.order, key)
+	}
+}

pkg/types/tool.go

@@ -63,59 +63,8 @@ type ToolReference struct {
 	ToolID    string
 }
 
-func (p Program) GetContextToolRefs(toolID string) (result []ToolReference, _ error) {
-	seen := map[struct {
-		toolID string
-		arg    string
-	}]struct{}{}
-	tool := p.ToolSet[toolID]
-
-	subToolRefs, err := tool.GetToolRefsFromNames(tool.Tools)
-	if err != nil {
-		return nil, err
-	}
-
-	for _, subToolRef := range subToolRefs {
-		subTool := p.ToolSet[subToolRef.ToolID]
-		exportContextToolRefs, err := subTool.GetToolRefsFromNames(subTool.ExportContext)
-		if err != nil {
-			return nil, err
-		}
-		for _, exportContextToolRef := range exportContextToolRefs {
-			key := struct {
-				toolID string
-				arg    string
-			}{
-				toolID: exportContextToolRef.ToolID,
-				arg:    exportContextToolRef.Arg,
-			}
-			if _, ok := seen[key]; !ok {
-				seen[key] = struct{}{}
-				result = append(result, exportContextToolRef)
-			}
-		}
-	}
-
-	contextToolRefs, err := p.ToolSet[toolID].GetToolRefsFromNames(p.ToolSet[toolID].Context)
-	if err != nil {
-		return nil, err
-	}
-
-	for _, contextToolRef := range contextToolRefs {
-		key := struct {
-			toolID string
-			arg    string
-		}{
-			toolID: contextToolRef.ToolID,
-			arg:    contextToolRef.Arg,
-		}
-		if _, ok := seen[key]; !ok {
-			seen[key] = struct{}{}
-			result = append(result, contextToolRef)
-		}
-	}
-
-	return
+func (p Program) GetContextToolRefs(toolID string) ([]ToolReference, error) {
+	return p.ToolSet[toolID].GetContextTools(p)
 }
 
 func (p Program) GetCompletionTools() (result []CompletionTool, err error) {
@@ -295,105 +244,139 @@ func (t Tool) String() string {
 	return buf.String()
 }
 
-func (t Tool) GetCompletionTools(prg Program) (result []CompletionTool, err error) {
-	toolNames := map[string]struct{}{}
+func (t Tool) GetExportedContext(prg Program) ([]ToolReference, error) {
+	result := &toolRefSet{}
 
-	subToolRefs, err := t.GetToolRefsFromNames(t.Parameters.Tools)
+	exportRefs, err := t.GetToolRefsFromNames(t.ExportContext)
 	if err != nil {
 		return nil, err
 	}
 
-	for _, subToolRef := range subToolRefs {
-		result, err = appendTool(result, prg, t, subToolRef.Reference, toolNames, subToolRef.Named)
-		if err != nil {
-			return nil, err
-		}
+	for _, exportRef := range exportRefs {
+		result.Add(exportRef)
+
+		tool := prg.ToolSet[exportRef.ToolID]
+		result.AddAll(tool.GetExportedContext(prg))
 	}
 
-	for _, subToolName := range t.Parameters.Context {
-		result, err = appendExports(result, prg, t, subToolName, toolNames)
-		if err != nil {
-			return nil, err
-		}
+	return result.List()
+}
+
+func (t Tool) GetExportedTools(prg Program) ([]ToolReference, error) {
+	result := &toolRefSet{}
+
+	exportRefs, err := t.GetToolRefsFromNames(t.Export)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, exportRef := range exportRefs {
+		result.Add(exportRef)
+		result.AddAll(prg.ToolSet[exportRef.ToolID].GetExportedTools(prg))
 	}
 
-	return result, nil
+	return result.List()
 }
 
-func getTool(prg Program, parent Tool, name string) (Tool, error) {
-	toolID, ok := parent.ToolMapping[name]
-	if !ok {
-		return Tool{}, &ErrToolNotFound{
-			ToolName: name,
-		}
+func (t Tool) GetContextTools(prg Program) ([]ToolReference, error) {
+	result := &toolRefSet{}
+
+	contextRefs, err := t.GetToolRefsFromNames(t.Context)
+	if err != nil {
+		return nil, err
 	}
-	tool, ok := prg.ToolSet[toolID]
-	if !ok {
-		return Tool{}, &ErrToolNotFound{
-			ToolName: name,
-		}
+
+	for _, contextRef := range contextRefs {
+		result.AddAll(prg.ToolSet[contextRef.ToolID].GetExportedContext(prg))
+		result.Add(contextRef)
 	}
-	return tool, nil
+
+	return result.List()
 }
 
-func appendExports(completionTools []CompletionTool, prg Program, parentTool Tool, subToolName string, toolNames map[string]struct{}) ([]CompletionTool, error) {
-	subTool, err := getTool(prg, parentTool, subToolName)
+func (t Tool) GetCompletionTools(prg Program) (result []CompletionTool, err error) {
+	refs, err := t.getCompletionToolRefs(prg)
 	if err != nil {
 		return nil, err
 	}
+	return toolRefsToCompletionTools(refs, prg), nil
+}
+
+func (t Tool) addReferencedTools(prg Program, result *toolRefSet) error {
+	subToolRefs, err := t.GetToolRefsFromNames(t.Parameters.Tools)
+	if err != nil {
+		return err
+	}
 
-	for _, export := range subTool.Export {
-		completionTools, err = appendTool(completionTools, prg, subTool, export, toolNames, "")
-		if err != nil {
-			return nil, err
-		}
+	for _, subToolRef := range subToolRefs {
+		// Add the tool
+		result.Add(subToolRef)
+
+		// Get all tools exports
+		result.AddAll(prg.ToolSet[subToolRef.ToolID].GetExportedTools(prg))
 	}
 
-	return completionTools, nil
+	return nil
 }
 
-func appendTool(completionTools []CompletionTool, prg Program, parentTool Tool, subToolName string, toolNames map[string]struct{}, asName string) ([]CompletionTool, error) {
-	subTool, err := getTool(prg, parentTool, subToolName)
+func (t Tool) addContextExportedTools(prg Program, result *toolRefSet) error {
+	contextTools, err := t.GetContextTools(prg)
 	if err != nil {
+		return err
+	}
+
+	for _, contextTool := range contextTools {
+		result.AddAll(prg.ToolSet[contextTool.ToolID].GetExportedTools(prg))
+	}
+
+	return nil
+}
+
+func (t Tool) getCompletionToolRefs(prg Program) ([]ToolReference, error) {
+	result := toolRefSet{}
+
+	if err := t.addReferencedTools(prg, &result); err != nil {
 		return nil, err
 	}
 
-	args := subTool.Parameters.Arguments
-	if args == nil && !subTool.IsCommand() && !subTool.Chat {
-		args = &system.DefaultToolSchema
+	if err := t.addContextExportedTools(prg, &result); err != nil {
+		return nil, err
 	}
 
-	for _, existingTool := range completionTools {
-		if existingTool.Function.ToolID == subTool.ID {
-			return completionTools, nil
+	return result.List()
+}
+
+func toolRefsToCompletionTools(completionTools []ToolReference, prg Program) (result []CompletionTool) {
+	toolNames := map[string]struct{}{}
+
+	for _, subToolRef := range completionTools {
+		subTool := prg.ToolSet[subToolRef.ToolID]
+
+		subToolName := subToolRef.Reference
+		if subToolRef.Named != "" {
+			subToolName = subToolRef.Named
 		}
-	}
 
-	if subTool.Instructions == "" {
-		log.Debugf("Skipping zero instruction tool %s (%s)", subToolName, subTool.ID)
-	} else {
-		name := subToolName
-		if asName != "" {
-			name = asName
+		args := subTool.Parameters.Arguments
+		if args == nil && !subTool.IsCommand() && !subTool.Chat {
+			args = &system.DefaultToolSchema
 		}
-		completionTools = append(completionTools, CompletionTool{
-			Function: CompletionFunctionDefinition{
-				ToolID:      subTool.ID,
-				Name:        PickToolName(name, toolNames),
-				Description: subTool.Parameters.Description,
-				Parameters:  args,
-			},
-		})
-	}
 
-	for _, export := range subTool.Export {
-		completionTools, err = appendTool(completionTools, prg, subTool, export, toolNames, "")
-		if err != nil {
-			return nil, err
+		if subTool.Instructions == "" {
+			log.Debugf("Skipping zero instruction tool %s (%s)", subToolName, subTool.ID)
+		} else {
+			result = append(result, CompletionTool{
+				Function: CompletionFunctionDefinition{
+					ToolID:      subTool.ID,
+					Name:        PickToolName(subToolName, toolNames),
+					Description: subTool.Parameters.Description,
+					Parameters:  args,
+				},
+			})
 		}
 	}
 
-	return completionTools, nil
+	return
 }
 
 type Repo struct {