diff --git a/README.md b/README.md index 8cd9d84..28b24d9 100644 --- a/README.md +++ b/README.md @@ -57,16 +57,42 @@ func main() { ### Using Tools +Tools in Cogito are defined using `NewToolDefinition`, which automatically generates `openai.Tool` via the `Tool()` method. You can mix tools with different type parameters in the same `Tools` slice. + ```go +// Define tool argument types +type WeatherArgs struct { + City string `json:"city" description:"The city to get weather for"` +} + +type SearchArgs struct { + Query string `json:"query" description:"The search query"` +} + +// Create tool definitions - these automatically generate openai.Tool +weatherTool := cogito.NewToolDefinition( + &WeatherTool{}, + WeatherArgs{}, + "get_weather", + "Get the current weather for a city", +) + +searchTool := cogito.NewToolDefinition( + &SearchTool{}, + SearchArgs{}, + "search", + "Search for information", +) + // Create a fragment with user input fragment := cogito.NewFragment(openai.ChatCompletionMessage{ Role: "user", Content: "What's the weather in San Francisco?", }) -// Execute with tools +// Execute with tools - you can pass multiple tools with different types result, err := cogito.ExecuteTools(llm, fragment, - cogito.WithTools(&WeatherTool{})) + cogito.WithTools(weatherTool, searchTool)) if err != nil { panic(err) } @@ -74,25 +100,157 @@ if err != nil { // result.Status.ToolsCalled will contain all the tools being called ``` +#### Creating Custom Tools + +To create a custom tool, implement the `Tool[T]` interface: + +```go +type MyToolArgs struct { + Param string `json:"param" description:"A parameter"` +} + +type MyTool struct{} + +// Implement the Tool interface +func (t *MyTool) Run(args MyToolArgs) (string, error) { + // Your tool logic here + return fmt.Sprintf("Processed: %s", args.Param), nil +} + +// Create a ToolDefinition using NewToolDefinition helper +myTool := cogito.NewToolDefinition( + &MyTool{}, + MyToolArgs{}, + "my_tool", + "A custom tool", +) +``` + +#### Field Annotations for Tool Arguments + +Cogito supports several struct field annotations to control how tool arguments are defined in the generated JSON schema: + +**Available Annotations:** + +- `json:"field_name"` - **Required**. Defines the JSON field name for the parameter. +- `description:"text"` - Provides a description for the field that helps the LLM understand what the parameter is for. +- `enum:"value1,value2,value3"` - Restricts the field to a specific set of allowed values (comma-separated). +- `required:"false"` - Makes the field optional. By default, all fields are required unless marked with `required:"false"`. + +**Examples:** + +```go +// Basic required field with description +type BasicArgs struct { + Query string `json:"query" description:"The search query"` +} + +// Optional field +type OptionalArgs struct { + Query string `json:"query" required:"false" description:"Optional search query"` + Limit int `json:"limit" required:"false" description:"Maximum number of results"` +} + +// Field with enum values +type EnumArgs struct { + Action string `json:"action" enum:"create,read,update,delete" description:"The action to perform"` +} + +// Field with enum and description +type WeatherArgs struct { + City string `json:"city" description:"The city name"` + Unit string `json:"unit" enum:"celsius,fahrenheit" description:"Temperature unit"` + Format string `json:"format" enum:"short,detailed" required:"false" description:"Output format"` +} + +// Complete example with multiple field types +type AdvancedSearchArgs struct { + // Required field with description + Query string `json:"query" description:"The search query"` + + // Optional field with enum + SortBy string `json:"sort_by" enum:"relevance,date,popularity" required:"false" description:"Sort order"` + + // Optional numeric field + Limit int `json:"limit" required:"false" description:"Maximum number of results"` + + // Optional boolean field + IncludeImages bool `json:"include_images" required:"false" description:"Include images in results"` +} + +// Create tool with advanced arguments +searchTool := cogito.NewToolDefinition( + &AdvancedSearchTool{}, + AdvancedSearchArgs{}, + "advanced_search", + "Advanced search with sorting and filtering options", +) +``` + +**Notes:** +- Fields without `required:"false"` are automatically marked as required in the JSON schema +- Enum values are case-sensitive and should match exactly what you expect in `Run()` +- The `json` tag is required for all fields that should be included in the tool schema +- Descriptions help the LLM understand the purpose of each parameter, leading to better tool calls + +Alternatively, you can implement `ToolDefinitionInterface` directly if you prefer more control: + +```go +type CustomTool struct{} + +func (t *CustomTool) Tool() openai.Tool { + return openai.Tool{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{ + Name: "custom_tool", + Description: "A custom tool", + Parameters: jsonschema.Definition{ + // Define your schema + }, + }, + } +} + +func (t *CustomTool) Execute(args map[string]any) (string, error) { + // Your execution logic + return "result", nil +} +``` + ### Guidelines for Intelligent Tool Selection Guidelines provide a powerful way to define conditional rules for tool usage. The LLM intelligently selects which guidelines are relevant based on the conversation context, enabling dynamic and context-aware tool selection. ```go +// Create tool definitions +searchTool := cogito.NewToolDefinition( + &SearchTool{}, + SearchArgs{}, + "search", + "Search for information", +) + +weatherTool := cogito.NewToolDefinition( + &WeatherTool{}, + WeatherArgs{}, + "get_weather", + "Get weather information", +) + // Define guidelines with conditions and associated tools guidelines := cogito.Guidelines{ cogito.Guideline{ Condition: "User asks about information or facts", Action: "Use the search tool to find information", Tools: cogito.Tools{ - &SearchTool{}, + searchTool, }, }, cogito.Guideline{ Condition: "User asks for the weather in a city", Action: "Use the weather tool to find the weather", Tools: cogito.Tools{ - &WeatherTool{}, + weatherTool, }, }, } @@ -119,16 +277,24 @@ if err != nil { panic(err) } +// Create tool definition +searchTool := cogito.NewToolDefinition( + &SearchTool{}, + SearchArgs{}, + "search", + "Search for information", +) + // Create a plan to achieve the goal plan, err := cogito.ExtractPlan(llm, fragment, goal, - cogito.WithTools(&SearchTool{})) + cogito.WithTools(searchTool)) if err != nil { panic(err) } // Execute the plan result, err := cogito.ExecutePlan(llm, fragment, plan, goal, - cogito.WithTools(&SearchTool{})) + cogito.WithTools(searchTool)) if err != nil { panic(err) } @@ -137,10 +303,18 @@ if err != nil { ### Content Refinement ```go +// Create tool definition +searchTool := cogito.NewToolDefinition( + &SearchTool{}, + SearchArgs{}, + "search", + "Search for information", +) + // Refine content through iterative improvement refined, err := cogito.ContentReview(llm, fragment, cogito.WithIterations(3), - cogito.WithTools(&SearchTool{})) + cogito.WithTools(searchTool)) if err != nil { panic(err) } @@ -162,10 +336,25 @@ initial := cogito.NewEmptyFragment(). response, _ := llm.Ask(ctx, initial) +// Create tool definitions +searchTool := cogito.NewToolDefinition( + &SearchTool{}, + SearchArgs{}, + "search", + "Search for information", +) + +factCheckTool := cogito.NewToolDefinition( + &FactCheckTool{}, + FactCheckArgs{}, + "fact_check", + "Verify facts", +) + // Iteratively improve with tool support improvedResponse, _ := cogito.ContentReview(reviewerLLM, response, cogito.WithIterations(3), - cogito.WithTools(&SearchTool{}, &FactCheckTool{}), + cogito.WithTools(searchTool, factCheckTool), cogito.EnableToolReasoner) ``` diff --git a/examples/chat/main.go b/examples/chat/main.go index ef0fcc8..15c0208 100644 --- a/examples/chat/main.go +++ b/examples/chat/main.go @@ -21,6 +21,14 @@ func main() { defaultLLM := cogito.NewOpenAILLM(model, apiKey, baseURL) + // Create tool definition - this automatically generates openai.Tool via Tool() method + searchTool := cogito.NewToolDefinition( + &search.SearchTool{}, + search.SearchArgs{}, + "search", + "A search engine to find information about a topic", + ) + f := cogito.NewEmptyFragment() for { reader := bufio.NewReader(os.Stdin) @@ -37,9 +45,7 @@ func main() { fmt.Println(s) fmt.Println("___________________ END STATUS _________________") }), - cogito.WithTools( - &search.SearchTool{}, - ), + cogito.WithTools(searchTool), ) if err != nil && !errors.Is(err, cogito.ErrNoToolSelected) { panic(err) diff --git a/examples/internal/search/search.go b/examples/internal/search/search.go index 19a2cd2..a4bc985 100644 --- a/examples/internal/search/search.go +++ b/examples/internal/search/search.go @@ -4,43 +4,27 @@ import ( "context" "fmt" - "github.com/sashabaranov/go-openai" - "github.com/tmc/langchaingo/jsonschema" "github.com/tmc/langchaingo/tools/duckduckgo" ) +// SearchArgs defines the arguments for the search tool +type SearchArgs struct { + Query string `json:"query" description:"The query to search for"` +} + +// SearchTool implements the Tool interface for searching type SearchTool struct { } -func (s *SearchTool) Run(args map[string]any) (string, error) { - - q, ok := args["query"].(string) - if !ok { - return "", fmt.Errorf("no query") +// Run executes the search with typed arguments +func (s *SearchTool) Run(args SearchArgs) (string, error) { + if args.Query == "" { + return "", fmt.Errorf("query is required") } + ddg, err := duckduckgo.New(5, "LocalAGI") if err != nil { return "", err } - return ddg.Call(context.Background(), q) -} - -func (s *SearchTool) Tool() openai.Tool { - return openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "search", - Description: "A search engine to find information about a topic", - Parameters: jsonschema.Definition{ - Type: jsonschema.Object, - Properties: map[string]jsonschema.Definition{ - "query": { - Type: jsonschema.String, - Description: "The query to search for", - }, - }, - Required: []string{"query"}, - }, - }, - } + return ddg.Call(context.Background(), args.Query) } diff --git a/fragment_e2e_test.go b/fragment_e2e_test.go index 4cdf37a..5a0fd1d 100644 --- a/fragment_e2e_test.go +++ b/fragment_e2e_test.go @@ -13,7 +13,7 @@ import ( "github.com/sashabaranov/go-openai/jsonschema" ) -type GeatWeatherTool struct { +type GetWeatherTool struct { status *ToolStatus } @@ -25,36 +25,20 @@ func (m MultimediaImage) URL() string { return m.url } -func (s *GeatWeatherTool) Status() *ToolStatus { +func (s *GetWeatherTool) Status() *ToolStatus { if s.status == nil { s.status = &ToolStatus{} } return s.status } -func (s *GeatWeatherTool) Run(args map[string]any) (string, error) { +func (s *GetWeatherTool) Run(args map[string]any) (string, error) { return "", nil } -func (s *GeatWeatherTool) Tool() openai.Tool { - return openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "get_weather", - Description: "Get the weather", - Parameters: jsonschema.Definition{ - Type: jsonschema.Object, - Properties: map[string]jsonschema.Definition{ - "city": { - Type: jsonschema.String, - Description: "The city to look-up the weather for", - }, - }, - Required: []string{"city"}, - }, - }, - } +type WeatherArgs struct { + City string `json:"city"` } var _ = Describe("Fragment test", func() { @@ -172,7 +156,12 @@ var _ = Describe("Result test", Label("e2e"), func() { }) newFragment, result, err := fragment.SelectTool(context.TODO(), defaultLLM, Tools{ - &GeatWeatherTool{}, + NewToolDefinition( + (&GetWeatherTool{}), + WeatherArgs{}, + "get_weather", + "Get weather information", + ), }, "") Expect(err).ToNot(HaveOccurred()) diff --git a/goal_e2e_test.go b/goal_e2e_test.go index d2ab1a3..ef8601d 100644 --- a/goal_e2e_test.go +++ b/goal_e2e_test.go @@ -34,15 +34,21 @@ var _ = Describe("cogito test", Label("e2e"), func() { var achieved *structures.Boolean for range 4 { // Simulate an "infinite loop" - conv, err = ExecuteTools(defaultLLM, conv, WithTools( - &SearchTool{ - results: []string{ - "India warns new US fee for H-1B visa will have 'humanitarian consequences' - bbc.com", - "Estonia seeks Nato consultation after Russian jets violate airspace - bbc.com", - "Day of delays at Heathrow after cyber-attack brings disruption - bbc.com", - "Kildunne stars as England see off France to make final - bbc.com", - }, + searchTool := &SearchTool{ + results: []string{ + "India warns new US fee for H-1B visa will have 'humanitarian consequences' - bbc.com", + "Estonia seeks Nato consultation after Russian jets violate airspace - bbc.com", + "Day of delays at Heathrow after cyber-attack brings disruption - bbc.com", + "Kildunne stars as England see off France to make final - bbc.com", }, + } + conv, err = ExecuteTools(defaultLLM, conv, WithTools( + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), )) Expect(err).ToNot(HaveOccurred()) diff --git a/guidelines.go b/guidelines.go index ce8f116..c97afd1 100644 --- a/guidelines.go +++ b/guidelines.go @@ -130,7 +130,7 @@ func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines, if len(o.guidelines) > 0 { if o.strictGuidelines { - tools = []Tool{} + tools = Tools{} } var err error guidelines, err = GetRelevantGuidelines(llm, o.guidelines, fragment, opts...) diff --git a/guidelines_e2e_test.go b/guidelines_e2e_test.go index 901651a..9b2b198 100644 --- a/guidelines_e2e_test.go +++ b/guidelines_e2e_test.go @@ -18,14 +18,24 @@ var _ = Describe("cogito test", Label("e2e"), func() { Condition: "User asks about informations", Action: "Use the search tool to find information about Isaac Asimov.", Tools: Tools{ - &SearchTool{}, + NewToolDefinition( + (&SearchTool{}), + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), }, }, Guideline{ Condition: "User asks for the weather in a city ", Action: "Use the weather tool to find the weather in the city.", Tools: Tools{ - &GeatWeatherTool{}, + NewToolDefinition( + (&GetWeatherTool{}), + WeatherArgs{}, + "get_weather", + "Get weather information", + ), }, }, } diff --git a/mcp.go b/mcp.go index 14c58e6..4ecb0e5 100644 --- a/mcp.go +++ b/mcp.go @@ -19,7 +19,22 @@ type mcpTool struct { props map[string]jsonschema.Definition } -func (t *mcpTool) Run(args map[string]any) (string, error) { +func (t *mcpTool) Tool() openai.Tool { + return openai.Tool{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{ + Name: t.name, + Description: t.description, + Parameters: jsonschema.Definition{ + Type: jsonschema.Object, + Properties: t.props, + Required: t.inputSchema.Required, + }, + }, + } +} + +func (t *mcpTool) Execute(args map[string]any) (string, error) { // Call a tool on the server. params := &mcp.CallToolParams{ @@ -45,21 +60,6 @@ func (t *mcpTool) Run(args map[string]any) (string, error) { return result, nil } -func (t *mcpTool) Tool() openai.Tool { - return openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: t.name, - Description: t.description, - Parameters: jsonschema.Definition{ - Type: jsonschema.Object, - Properties: t.props, - Required: t.inputSchema.Required, - }, - }, - } -} - func (t *mcpTool) Close() { t.session.Close() } @@ -98,8 +98,8 @@ func mcpPromptsFromTransport(ctx context.Context, session *mcp.ClientSession, ar } // probe the MCP remote and generate tools that are compliant with cogito -func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession) ([]*mcpTool, error) { - allTools := []*mcpTool{} +func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession) ([]ToolDefinitionInterface, error) { + allTools := []ToolDefinitionInterface{} tools, err := session.ListTools(ctx, nil) if err != nil { diff --git a/options.go b/options.go index d331c07..e54d196 100644 --- a/options.go +++ b/options.go @@ -122,8 +122,10 @@ func WithPrompt(t prompt.PromptType, p prompt.StaticPrompt) func(o *Options) { } } -// WithTools allows to set the tools available to the Agent -func WithTools(tools ...Tool) func(o *Options) { +// WithTools allows to set the tools available to the Agent. +// Pass *ToolDefinition[T] instances - they will automatically generate openai.Tool via their Tool() method. +// Example: WithTools(&ToolDefinition[SearchArgs]{...}, &ToolDefinition[WeatherArgs]{...}) +func WithTools(tools ...ToolDefinitionInterface) func(o *Options) { return func(o *Options) { o.tools = append(o.tools, tools...) } diff --git a/plan_e2e_test.go b/plan_e2e_test.go index 6431a6d..0193164 100644 --- a/plan_e2e_test.go +++ b/plan_e2e_test.go @@ -24,7 +24,13 @@ var _ = Describe("cogito test", Label("e2e"), func() { Expect(strings.ToLower(goal.Goal)).To(ContainSubstring(strings.ToLower("Isaac Asimov"))) plan, err := ExtractPlan(defaultLLM, conv, goal, WithTools( - &SearchTool{})) + NewToolDefinition( + (&SearchTool{}), + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), + )) Expect(err).ToNot(HaveOccurred()) Expect(plan).ToNot(BeNil()) @@ -35,14 +41,22 @@ var _ = Describe("cogito test", Label("e2e"), func() { // This is more of an integration test It("is able to extract a plan and execute subtasks", func() { defaultLLM := NewOpenAILLM(defaultModel, "", apiEndpoint) - tools := Tools{&SearchTool{ + searchTool := &SearchTool{ results: []string{ "Isaac Asimov was a prolific science fiction writer and biochemist.", "He was born on January 2, 1920, in Petrovichi, Russia.", "Asimov is best known for his Foundation series and Robot series.", "He wrote or edited over 500 books and an estimated 90,000 letters and postcards.", }, - }} + } + tools := Tools{ + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), + } conv := NewEmptyFragment().AddMessage("user", "You need to search all informations you can about Isaac Asimov.") diff --git a/plan_test.go b/plan_test.go index ce2a57b..1dde4b9 100644 --- a/plan_test.go +++ b/plan_test.go @@ -34,7 +34,7 @@ var _ = Describe("Plannings with tools", func() { // Mock tool call (Subtask #1) - tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) - mockTool.SetRunResult("Chlorophyll is a green pigment found in plants.") + mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ @@ -54,7 +54,7 @@ var _ = Describe("Plannings with tools", func() { // Mock tool call (Subtask #2) - tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "photosynthesis"}`) - mockTool.SetRunResult("Photosynthesis is the process by which plants convert sunlight into energy.") + mock.SetRunResult(mockTool, "Photosynthesis is the process by which plants convert sunlight into energy.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ diff --git a/reviewer_e2e_test.go b/reviewer_e2e_test.go index 6f7474b..0c34061 100644 --- a/reviewer_e2e_test.go +++ b/reviewer_e2e_test.go @@ -36,7 +36,12 @@ var _ = Describe("cogito test", Label("e2e"), func() { searchTool := &SearchTool{} f, err := ContentReview(defaultLLM, result, WithMaxAttempts(1), WithIterations(2), WithTools( - searchTool, + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), )) Expect(err).ToNot(HaveOccurred()) Expect(f.String()).ToNot(BeEmpty()) @@ -52,7 +57,14 @@ var _ = Describe("cogito test", Label("e2e"), func() { conv := NewEmptyFragment().AddMessage("user", "What are the latest news today?") searchTool := &SearchTool{} - f, err := ContentReview(defaultLLM, conv, WithIterations(2), WithTools(searchTool)) + f, err := ContentReview(defaultLLM, conv, WithIterations(2), WithTools( + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), + )) Expect(err).ToNot(HaveOccurred()) Expect(f.String()).ToNot(BeEmpty()) Expect(f.Status.Iterations).To(Equal(2)) diff --git a/reviewer_test.go b/reviewer_test.go index 565976b..88f2237 100644 --- a/reviewer_test.go +++ b/reviewer_test.go @@ -27,7 +27,7 @@ var _ = Describe("ContentReview", func() { // First iteration - tool selection and execution mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) - mockTool.SetRunResult("Chlorophyll is a green pigment found in plants.") + mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ @@ -52,7 +52,7 @@ var _ = Describe("ContentReview", func() { // Second iteration - tool selection and execution mockLLM.AddCreateChatCompletionFunction("search", `{"query": "why chlorophyll is green"}`) - mockTool.SetRunResult("Chlorophyll is green because it absorbs blue and red light and reflects green light.") + mock.SetRunResult(mockTool, "Chlorophyll is green because it absorbs blue and red light and reflects green light.") // After second tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ diff --git a/tests/mock/tool.go b/tests/mock/tool.go index 5786e5d..a7abbb9 100644 --- a/tests/mock/tool.go +++ b/tests/mock/tool.go @@ -2,7 +2,6 @@ package mock import ( . "github.com/mudler/cogito" - "github.com/sashabaranov/go-openai" ) // MockTool implements the Tool interface for testing @@ -13,26 +12,26 @@ type MockTool struct { runError error runIndex int status *ToolStatus + toolDef *ToolDefinition[map[string]any] } -func NewMockTool(name, description string) *MockTool { - return &MockTool{ +func NewMockTool(name, description string) ToolDefinitionInterface { + mockTool := &MockTool{ name: name, description: description, status: &ToolStatus{}, } -} - -func (m *MockTool) Tool() openai.Tool { - return openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: m.name, - Description: m.description, - // We don't need parameters in this mock - // Will show up as null + toolDef := &ToolDefinition[map[string]any]{ + ToolRunner: mockTool, + Name: name, + Description: description, + InputArguments: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, }, } + mockTool.toolDef = toolDef + return toolDef } func (m *MockTool) Status() *ToolStatus { @@ -49,6 +48,11 @@ func (m *MockTool) Run(args map[string]any) (string, error) { return m.runResults[m.runIndex], nil } +func (m *MockTool) NewArgs() *map[string]any { + args := make(map[string]any) + return &args +} + func (m *MockTool) SetRunResult(result string) { m.runResults = append(m.runResults, result) } @@ -56,3 +60,27 @@ func (m *MockTool) SetRunResult(result string) { func (m *MockTool) SetRunError(err error) { m.runError = err } + +// GetMockTool extracts the MockTool from a ToolDef (if it contains one) +func GetMockTool(toolDef ToolDefinitionInterface) *MockTool { + if toolDefT, ok := toolDef.(*ToolDefinition[map[string]any]); ok { + if mockTool, ok := toolDefT.ToolRunner.(*MockTool); ok { + return mockTool + } + } + return nil +} + +// SetRunResult sets the result for a mock tool within a ToolDef +func SetRunResult(toolDef ToolDefinitionInterface, result string) { + if mockTool := GetMockTool(toolDef); mockTool != nil { + mockTool.SetRunResult(result) + } +} + +// SetRunError sets an error for a mock tool within a ToolDef +func SetRunError(toolDef ToolDefinitionInterface, err error) { + if mockTool := GetMockTool(toolDef); mockTool != nil { + mockTool.SetRunError(err) + } +} diff --git a/tool_intention.go b/tool_intention.go new file mode 100644 index 0000000..a6cabc0 --- /dev/null +++ b/tool_intention.go @@ -0,0 +1,47 @@ +package cogito + +import "fmt" + +// IntentionResponse is used to extract the tool choice from the intention tool +type IntentionResponse struct { + Tool string `json:"tool"` + Reasoning string `json:"reasoning"` +} + +// intentionToolWrapper wraps the intention tool to match the Tool interface +type intentionToolWrapper struct{} + +func (i *intentionToolWrapper) Run(args IntentionResponse) (string, error) { + return "", fmt.Errorf("intention tool should not be executed") +} + +func (i *intentionToolWrapper) NewArgs() *IntentionResponse { + return &IntentionResponse{} +} + +// intentionTool creates a tool that forces the LLM to pick one of the available tools +func intentionTool(toolNames []string) *ToolDefinition[IntentionResponse] { + // Build enum for the tool names + enumValues := append([]string{"reply"}, toolNames...) + + return &ToolDefinition[IntentionResponse]{ + ToolRunner: &intentionToolWrapper{}, + Name: "pick_tool", + InputArguments: map[string]interface{}{ + "description": "Pick the most appropriate tool to use based on the reasoning. Choose 'reply' if no tool is needed.", + "type": "object", + "properties": map[string]interface{}{ + "tool": map[string]interface{}{ + "type": "string", + "description": "The tool to use", + "enum": enumValues, + }, + "reasoning": map[string]interface{}{ + "type": "string", + "description": "The reasoning for the tool choice", + }, + }, + "required": []string{"tool"}, + }, + } +} diff --git a/tools.go b/tools.go index 389f8dc..de25c39 100644 --- a/tools.go +++ b/tools.go @@ -10,6 +10,7 @@ import ( "github.com/mudler/cogito/pkg/xlog" "github.com/mudler/cogito/prompt" "github.com/sashabaranov/go-openai" + "github.com/sashabaranov/go-openai/jsonschema" ) var ( @@ -24,56 +25,6 @@ type ToolStatus struct { Name string } -// IntentionResponse is used to extract the tool choice from the intention tool -type IntentionResponse struct { - Tool string `json:"tool" jsonschema:"title=Tool,description=The tool to use,enum=reply"` - Reasoning string `json:"reasoning" jsonschema:"title=Reasoning,description=The reasoning for the tool choice"` -} - -// intentionToolWrapper wraps the intention tool to match the Tool interface -type intentionToolWrapper struct { - tool openai.Tool -} - -func (i intentionToolWrapper) Tool() openai.Tool { - return i.tool -} - -func (i intentionToolWrapper) Run(args map[string]any) (string, error) { - return "", fmt.Errorf("intention tool should not be executed") -} - -// intentionTool creates a tool that forces the LLM to pick one of the available tools -func intentionTool(toolNames []string) Tool { - // Build enum for the tool names - enumValues := append([]string{"reply"}, toolNames...) - - return intentionToolWrapper{ - tool: openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "pick_tool", - Description: "Pick the most appropriate tool to use based on the reasoning. Choose 'reply' if no tool is needed.", - Parameters: map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{ - "tool": map[string]interface{}{ - "type": "string", - "description": "The tool to use", - "enum": enumValues, - }, - "reasoning": map[string]interface{}{ - "type": "string", - "description": "The reasoning for the tool choice", - }, - }, - "required": []string{"tool"}, - }, - }, - }, - } -} - // decisionResult holds the result of a tool decision from the LLM type decisionResult struct { toolChoice *ToolChoice @@ -81,14 +32,93 @@ type decisionResult struct { toolName string } -type Tool interface { +type ToolDefinitionInterface interface { Tool() openai.Tool - Run(args map[string]any) (string, error) + // Execute runs the tool with the given arguments (as JSON map) and returns the result + Execute(args map[string]any) (string, error) } -type Tools []Tool +type Tool[T any] interface { + Run(args T) (string, error) +} -func (t Tools) Find(name string) Tool { +type ToolDefinition[T any] struct { + ToolRunner Tool[T] + InputArguments any + Name, Description string +} + +func NewToolDefinition[T any](toolRunner Tool[T], inputArguments any, name, description string) ToolDefinitionInterface { + return &ToolDefinition[T]{ + ToolRunner: toolRunner, + InputArguments: inputArguments, + Name: name, + Description: description, + } +} + +var _ ToolDefinitionInterface = &ToolDefinition[any]{} + +func (t ToolDefinition[T]) Tool() openai.Tool { + var schema *jsonschema.Definition + + // Handle map[string]interface{} (JSON schema format) + if inputMap, ok := t.InputArguments.(map[string]any); ok { + dat, err := json.Marshal(inputMap) + if err != nil { + panic(err) + } + s := &jsonschema.Definition{} + err = json.Unmarshal(dat, s) + if err != nil { + panic(err) + } + schema = s + } else { + // Try to generate schema from struct type + var err error + schema, err = jsonschema.GenerateSchemaForType(t.InputArguments) + if err != nil { + panic(fmt.Errorf("unsupported InputArguments type: %T, error: %w", t.InputArguments, err)) + } + } + + return openai.Tool{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{ + Name: t.Name, + Description: t.Description, + Parameters: *schema, + }, + } +} + +// Execute implements ToolDef.Execute by marshaling the arguments map to type T and calling ToolRunner.Run +func (t *ToolDefinition[T]) Execute(args map[string]any) (string, error) { + if t.ToolRunner == nil { + return "", fmt.Errorf("tool %s has no ToolRunner", t.Name) + } + + argsPtr := new(T) + + // Marshal the map to JSON and unmarshal into the typed struct + argsBytes, err := json.Marshal(args) + if err != nil { + return "", fmt.Errorf("failed to marshal tool arguments: %w", err) + } + + err = json.Unmarshal(argsBytes, argsPtr) + if err != nil { + return "", fmt.Errorf("failed to unmarshal tool arguments: %w", err) + } + + // Call Run with the typed arguments + return t.ToolRunner.Run(*argsPtr) +} + +type Tools []ToolDefinitionInterface + +func (t Tools) Find(name string) ToolDefinitionInterface { for _, tool := range t { if tool.Tool().Function.Name == name { return tool @@ -101,9 +131,7 @@ func (t Tools) ToOpenAI() []openai.Tool { openaiTools := []openai.Tool{} for _, tool := range t { openaiTools = append(openaiTools, tool.Tool()) - } - return openaiTools } @@ -209,7 +237,7 @@ func formatToolParameters(params interface{}) string { // generateToolParameters generates parameters for a specific tool with enhanced reasoning // Similar to agent.go's generateParameters but adapted for cogito -func generateToolParameters(o *Options, llm LLM, tool Tool, conversation []openai.ChatCompletionMessage, +func generateToolParameters(o *Options, llm LLM, tool ToolDefinitionInterface, conversation []openai.ChatCompletionMessage, reasoning string) (*ToolChoice, error) { toolFunc := tool.Tool().Function @@ -663,6 +691,7 @@ func toolSelection(llm LLM, f Fragment, tools Tools, guidelines Guidelines, tool toolFunc := selectedToolObj.Tool().Function if o.forceReasoning && toolFunc != nil && toolFunc.Parameters != nil { xlog.Debug("[toolSelection] Regenerating parameters with reasoning") + enhancedChoice, err := generateToolParameters(o, llm, selectedToolObj, messages, reasoning) if err != nil { xlog.Warn("[toolSelection] Failed to regenerate parameters, using original", "error", err) @@ -856,13 +885,16 @@ func ExecuteTools(llm LLM, f Fragment, opts ...Option) (Fragment, error) { //f.Messages = append(f.Messages, selectedToolFragment.LastAssistantMessages()...) toolResult := tools.Find(selectedToolResult.Name) + if toolResult == nil { + return f, fmt.Errorf("tool %s not found", selectedToolResult.Name) + } // Execute tool attempts := 1 var result string RETRY: for range o.maxAttempts { - result, err = toolResult.Run(selectedToolResult.Arguments) + result, err = toolResult.Execute(selectedToolResult.Arguments) if err != nil { if attempts >= o.maxAttempts { // don't return error, set it as result diff --git a/tools_e2e_test.go b/tools_e2e_test.go index 9f4a625..e354a4b 100644 --- a/tools_e2e_test.go +++ b/tools_e2e_test.go @@ -10,8 +10,6 @@ import ( . "github.com/mudler/cogito" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "github.com/sashabaranov/go-openai" - "github.com/sashabaranov/go-openai/jsonschema" ) type SearchTool struct { @@ -27,13 +25,13 @@ func (s *SearchTool) Status() *ToolStatus { return s.status } -func (s *SearchTool) Run(args map[string]any) (string, error) { - q, ok := args["query"].(string) - if !ok { - return "", nil - } +type SearchArgs struct { + Query string `json:"query"` +} - s.searchedQuery = q +func (s *SearchTool) Run(args SearchArgs) (string, error) { + + s.searchedQuery = args.Query // Mocked search result searchResult := struct { Results []string `json:"results"` @@ -58,25 +56,7 @@ func (s *SearchTool) Run(args map[string]any) (string, error) { return string(b), nil } -func (s *SearchTool) Tool() openai.Tool { - return openai.Tool{ - Type: openai.ToolTypeFunction, - Function: &openai.FunctionDefinition{ - Name: "search", - Description: "A search engine to find information about a topic", - Parameters: jsonschema.Definition{ - Type: jsonschema.Object, - Properties: map[string]jsonschema.Definition{ - "query": { - Type: jsonschema.String, - Description: "The query to search for", - }, - }, - Required: []string{"query"}, - }, - }, - } -} +// ToToolDefinition converts SearchTool to ToolDefinition var _ = Describe("Tool execution", Label("e2e"), func() { Context("Using user-defined tools", func() { @@ -85,7 +65,12 @@ var _ = Describe("Tool execution", Label("e2e"), func() { conv := NewEmptyFragment().AddMessage("user", "Hi! All you are good?") searchTool := &SearchTool{} f, err := ExecuteTools(defaultLLM, conv, EnableToolReasoner, WithTools( - searchTool, + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), )) Expect(err).To(HaveOccurred()) Expect(errors.Is(err, ErrNoToolSelected)).To(BeTrue()) @@ -105,7 +90,12 @@ var _ = Describe("Tool execution", Label("e2e"), func() { conv := NewEmptyFragment().AddMessage("user", "What are the latest news today?") searchTool := &SearchTool{} f, err := ExecuteTools(defaultLLM, conv, EnableToolReasoner, WithTools( - searchTool, + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), )) Expect(err).ToNot(HaveOccurred()) Expect(f.Status.Iterations).To(Equal(1)) @@ -118,7 +108,14 @@ var _ = Describe("Tool execution", Label("e2e"), func() { defaultLLM := NewOpenAILLM(defaultModel, "", apiEndpoint) conv := NewEmptyFragment().AddMessage("user", "What are the latest news today?") searchTool := &SearchTool{} - f, err := ExecuteTools(defaultLLM, conv, WithTools(searchTool)) + f, err := ExecuteTools(defaultLLM, conv, WithTools( + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), + )) Expect(err).ToNot(HaveOccurred()) Expect(f.Status.Iterations).To(Equal(1)) Expect(f.Status.ToolsCalled).To(HaveLen(1)) @@ -163,7 +160,14 @@ var _ = Describe("Tool execution", Label("e2e"), func() { conv := NewEmptyFragment().AddMessage("user", "I need you to search for information about Isaac Asimov's life, his major works, and then his contributions to science fiction.") f, err := ExecuteTools(defaultLLM, conv, EnableAutoPlan, - WithTools(searchTool), + WithTools( + NewToolDefinition( + searchTool, + SearchArgs{}, + "search", + "A search engine to find information about a topic", + ), + ), WithMaxAttempts(1), WithIterations(1)) Expect(err).ToNot(HaveOccurred()) diff --git a/tools_test.go b/tools_test.go index 807a4cc..b5f9c60 100644 --- a/tools_test.go +++ b/tools_test.go @@ -8,6 +8,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/sashabaranov/go-openai" + "github.com/sashabaranov/go-openai/jsonschema" ) var _ = Describe("ExecuteTools", func() { @@ -21,25 +22,111 @@ var _ = Describe("ExecuteTools", func() { AddMessage("assistant", "Photosynthesis is the process by which plants convert sunlight into energy.") }) + Context("ToolDefinition", func() { + It("should create a valid ToolDefinition", func() { + mockToolDef := mock.NewMockTool("search", "Search for information") + mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) + toolDefinition := ToolDefinition[map[string]any]{ + ToolRunner: mockToolDefT.ToolRunner, + Name: "search", + Description: "Search for information", + InputArguments: &struct { + Query string `json:"query"` + }{}, + } + tool := toolDefinition.Tool() + Expect(tool.Function.Name).To(Equal("search")) + Expect(tool.Function.Description).To(Equal("Search for information")) + Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ + Type: jsonschema.Object, + AdditionalProperties: false, + Properties: map[string]jsonschema.Definition{ + "query": { + Type: jsonschema.String, + Enum: nil, + }, + }, + Required: []string{"query"}, + Defs: map[string]jsonschema.Definition{}, + })) + }) + + It("should create a valid ToolDefinition with enums and description", func() { + mockToolDef := mock.NewMockTool("search", "Search for information") + mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) + toolDefinition := ToolDefinition[map[string]any]{ + ToolRunner: mockToolDefT.ToolRunner, + Name: "search", + Description: "Search for information", + InputArguments: &struct { + Query string `json:"query" enum:"foo,bar" description:"The query to search for"` + }{}, + } + tool := toolDefinition.Tool() + Expect(tool.Function.Name).To(Equal("search")) + Expect(tool.Function.Description).To(Equal("Search for information")) + Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ + Type: jsonschema.Object, + AdditionalProperties: false, + Properties: map[string]jsonschema.Definition{ + "query": { + Type: jsonschema.String, + Enum: []string{"foo", "bar"}, + Description: "The query to search for", + }, + }, + Required: []string{"query"}, + Defs: map[string]jsonschema.Definition{}, + })) + }) + + It("should create a valid ToolDefinition which arg is not required", func() { + mockToolDef := mock.NewMockTool("search", "Search for information") + mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) + toolDefinition := ToolDefinition[map[string]any]{ + ToolRunner: mockToolDefT.ToolRunner, + Name: "search", + Description: "Search for information", + InputArguments: &struct { + Query string `json:"query" required:"false"` + }{}, + } + tool := toolDefinition.Tool() + Expect(tool.Function.Name).To(Equal("search")) + Expect(tool.Function.Description).To(Equal("Search for information")) + Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ + Type: jsonschema.Object, + AdditionalProperties: false, + Properties: map[string]jsonschema.Definition{ + "query": { + Type: jsonschema.String, + }, + }, + Required: nil, + Defs: map[string]jsonschema.Definition{}, + })) + }) + }) + Context("ExecuteTools with tools", func() { It("should execute tools when provided", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection and execution mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) - mockTool.SetRunResult("Chlorophyll is a green pigment found in plants.") + mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // After tool execution, ToolReEvaluator (toolSelection) picks next tool mockLLM.AddCreateChatCompletionFunction("search", `{"query": "grass"}`) // Second tool selection and execution // (The "grass" tool call above will be picked as nextAction) - mockTool.SetRunResult("Grass is a plant that grows on the ground.") + mock.SetRunResult(mockTool, "Grass is a plant that grows on the ground.") // After tool execution, ToolReEvaluator (toolSelection) picks next tool mockLLM.AddCreateChatCompletionFunction("search", `{"query": "baz"}`) // Third tool selection and execution // (The "baz" tool call above will be picked as nextAction) - mockTool.SetRunResult("Baz is a plant that grows on the ground.") + mock.SetRunResult(mockTool, "Baz is a plant that grows on the ground.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ @@ -85,7 +172,7 @@ var _ = Describe("ExecuteTools", func() { mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [1]}`) // 2. Tool selection (direct): mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) - mockTool.SetRunResult("Chlorophyll is a green pigment found in plants.") + mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // 3. ToolReEvaluator (toolSelection) returns no tool (text response): mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ @@ -104,7 +191,7 @@ var _ = Describe("ExecuteTools", func() { mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [1]}`) // 2. Tool selection: mockLLM.AddCreateChatCompletionFunction("search", `{"query": "grass"}`) - mockTool.SetRunResult("Grass is a plant that grows on the ground.") + mock.SetRunResult(mockTool, "Grass is a plant that grows on the ground.") // 3. ToolReEvaluator (toolSelection) returns no tool (text response): mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ @@ -123,7 +210,7 @@ var _ = Describe("ExecuteTools", func() { mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [2]}`) // 2. Tool selection: mockLLM.AddCreateChatCompletionFunction("get_weather", `{"query": "baz"}`) - mockWeatherTool.SetRunResult("Baz is a plant that grows on the ground.") + mock.SetRunResult(mockWeatherTool, "Baz is a plant that grows on the ground.") // 3. ToolReEvaluator (toolSelection) returns no tool (text response): mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ @@ -200,7 +287,7 @@ var _ = Describe("ExecuteTools", func() { // Mock first subtask execution - search mockLLM.AddCreateChatCompletionFunction("search", `{"query": "photosynthesis basics"}`) - mockTool.SetRunResult("Photosynthesis is the process by which plants convert sunlight into energy.") + mock.SetRunResult(mockTool, "Photosynthesis is the process by which plants convert sunlight into energy.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ @@ -278,7 +365,7 @@ var _ = Describe("ExecuteTools", func() { // Mock regular tool execution (since planning is not needed, it falls back to normal tool execution) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "photosynthesis"}`) - mockTool.SetRunResult("Photosynthesis is the process by which plants convert sunlight into energy.") + mock.SetRunResult(mockTool, "Photosynthesis is the process by which plants convert sunlight into energy.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{