diff --git a/pkg/agent/tool.go b/pkg/agent/tool.go index d1d637596..09cf0eb8c 100644 --- a/pkg/agent/tool.go +++ b/pkg/agent/tool.go @@ -18,6 +18,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "go.probo.inc/probo/pkg/llm" ) @@ -39,10 +40,11 @@ type ( } functionTool[P any] struct { - name string - description string - fn func(ctx context.Context, params P) (ToolResult, error) - schema json.RawMessage + name string + description string + fn func(ctx context.Context, params P) (ToolResult, error) + schema json.RawMessage + requiredFields []string } ) @@ -53,11 +55,17 @@ func FunctionTool[P any]( ) Tool { schema := jsonSchemaFor[P]() + var parsed struct { + Required []string `json:"required"` + } + _ = json.Unmarshal(schema, &parsed) + return &functionTool[P]{ - name: name, - description: description, - fn: fn, - schema: schema, + name: name, + description: description, + fn: fn, + schema: schema, + requiredFields: parsed.Required, } } @@ -72,6 +80,33 @@ func (t *functionTool[P]) Definition() llm.Tool { } func (t *functionTool[P]) Execute(ctx context.Context, arguments string) (ToolResult, error) { + if len(t.requiredFields) > 0 { + var fields map[string]json.RawMessage + if err := json.Unmarshal([]byte(arguments), &fields); err != nil { + return ToolResult{ + Content: fmt.Sprintf("Invalid parameters: %s", err.Error()), + IsError: true, + }, nil + } + + var missing []string + for _, f := range t.requiredFields { + if _, ok := fields[f]; !ok { + missing = append(missing, f) + } + } + + if len(missing) > 0 { + return ToolResult{ + Content: fmt.Sprintf( + "Missing required parameters: %s", + strings.Join(missing, ", "), + ), + IsError: true, + }, nil + } + } + var params P if err := json.Unmarshal([]byte(arguments), ¶ms); err != nil { return ToolResult{ diff --git a/pkg/agent/tool_test.go b/pkg/agent/tool_test.go index f6b5e582f..76cee47eb 100644 --- a/pkg/agent/tool_test.go +++ b/pkg/agent/tool_test.go @@ -275,6 +275,109 @@ func TestFunctionTool_Execute(t *testing.T) { }, ) + t.Run( + "missing single required field returns tool error", + func(t *testing.T) { + t.Parallel() + + type Params struct { + City string `json:"city"` + } + + tool := agent.FunctionTool( + "weather", + "Get weather", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + t.Fatal("function should not be called") + return agent.ToolResult{}, nil + }, + ) + + result, err := tool.Execute(context.Background(), `{}`) + require.NoError(t, err) + assert.True(t, result.IsError) + assert.Contains(t, result.Content, "city") + }, + ) + + t.Run( + "missing multiple required fields lists all of them", + func(t *testing.T) { + t.Parallel() + + type Params struct { + City string `json:"city"` + Country string `json:"country"` + } + + tool := agent.FunctionTool( + "weather", + "Get weather", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + t.Fatal("function should not be called") + return agent.ToolResult{}, nil + }, + ) + + result, err := tool.Execute(context.Background(), `{}`) + require.NoError(t, err) + assert.True(t, result.IsError) + assert.Contains(t, result.Content, "city") + assert.Contains(t, result.Content, "country") + }, + ) + + t.Run( + "partially missing required fields detected", + func(t *testing.T) { + t.Parallel() + + type Params struct { + City string `json:"city"` + Country string `json:"country"` + } + + tool := agent.FunctionTool( + "weather", + "Get weather", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + t.Fatal("function should not be called") + return agent.ToolResult{}, nil + }, + ) + + result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) + require.NoError(t, err) + assert.True(t, result.IsError) + assert.Contains(t, result.Content, "country") + }, + ) + + t.Run( + "optional fields can be omitted", + func(t *testing.T) { + t.Parallel() + + type Params struct { + City string `json:"city"` + Units *string `json:"units,omitempty"` + } + + tool := agent.FunctionTool( + "weather", + "Get weather", + func(_ context.Context, p Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "sunny in " + p.City}, nil + }, + ) + + result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) + require.NoError(t, err) + assert.False(t, result.IsError) + assert.Equal(t, "sunny in Paris", result.Content) + }, + ) + t.Run( "extra JSON fields are ignored", func(t *testing.T) {