LocalAI/cogito
1
0
Fork 0
mirror of https://github.com/mudler/cogito.git synced 2026-07-23 18:35:23 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat: declare tool args via struct (#14)

Browse Source 
* feat: declare tool args via struct

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* simplify even more the API

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto
2025-10-29 19:20:17 +01:00
committed by GitHub
parent d94e0bd20e
commit 8859641f81
18 changed files with 621 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+197 -8
View File
@@ -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)
```
examples/chat/main.go
+9 -3
View File
@@ -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)
examples/internal/search/search.go
+12 -28
View File
@@ -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)
}
fragment_e2e_test.go
+11 -22
View File
@@ -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())
goal_e2e_test.go
+14 -8
View File
@@ -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())
guidelines.go
+1 -1
View File
@@ -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...)
guidelines_e2e_test.go
+12 -2
View File
@@ -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",
),
},
},
}
mcp.go
+18 -18
View File
@@ -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 {
options.go
+4 -2
View File
@@ -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...)
}
plan_e2e_test.go
+17 -3
View File
@@ -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.")
plan_test.go
+2 -2
View File
@@ -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{
reviewer_e2e_test.go
+14 -2
View File
@@ -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))
reviewer_test.go
+2 -2
View File
@@ -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{
tests/mock/tool.go
+41 -13
View File
@@ -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)
}
}
tool_intention.go
+47
View File
@@ -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"},
},
}
}
tools.go
+90 -58
View File
@@ -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
tools_e2e_test.go
+35 -31
View File
@@ -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())
tools_test.go
+95 -8
View File
@@ -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{