diff --git a/README.md b/README.md index ba48ed4..5bca30f 100644 --- a/README.md +++ b/README.md @@ -1389,6 +1389,36 @@ make example-chat This starts an interactive chat session with tool support including web search capabilities. +### Strict Tool Schemas + +A tool's schema tells the model which arguments to send, but on most +backends nothing makes the model follow it: a call with a misspelled or +unknown argument reaches the tool, and the tool rejects it. With +`EnableStrictToolSchemas`, cogito sends each tool with `"strict": true`, so a +backend that supports it constrains the arguments while the model generates +them. On LocalAI, a strict tool switches on its tool grammar. OpenAI enforces +the schema. + +```go +result, err := cogito.ExecuteTools(llm, fragment, + cogito.WithMCPs(session), + cogito.EnableStrictToolSchemas) +``` + +Strict mode has rules of its own, so cogito adjusts each schema it sends: + +- Every object gets `"additionalProperties": false`. +- Every property is required. A property that was optional becomes nullable, + and the model sends `null` to leave it out. cogito removes those `null` + values before it calls the tool, so the tool sees an omitted argument. + +A tool whose schema strict mode cannot express (an object with schema-valued +or `true` `additionalProperties`, or a `oneOf`) is sent unchanged and not +strict. + +MCP tools are advertised with the `additionalProperties` their server +declares, with or without this option. + ### Custom Tool Implementation See `examples/internal/search/search.go` for a complete example of implementing a DuckDuckGo search tool. diff --git a/guidelines.go b/guidelines.go index 543b8ee..eb09cc9 100644 --- a/guidelines.go +++ b/guidelines.go @@ -207,5 +207,9 @@ func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines, } } + if o.strictToolSchemas { + tools = strictTools(tools) + } + return tools, guidelines, prompts, nil } diff --git a/mcp.go b/mcp.go index 16a7ce6..2b95887 100644 --- a/mcp.go +++ b/mcp.go @@ -26,10 +26,11 @@ func (t *mcpTool) Tool() openai.Tool { Function: &openai.FunctionDefinition{ Name: t.name, Description: t.description, - Parameters: jsonschema.Definition{ - Type: jsonschema.Object, - Properties: t.props, - Required: t.inputSchema.Required, + Parameters: mcpToolParameters{ + Type: jsonschema.Object, + Properties: t.props, + Required: t.inputSchema.Required, + AdditionalProperties: t.inputSchema.AdditionalProperties, }, }, } @@ -101,6 +102,22 @@ type toolInputSchema struct { Type string `json:"type"` Properties map[string]interface{} `json:"properties,omitempty"` Required []string `json:"required,omitempty"` + // AdditionalProperties is kept as the server declared it. A server whose + // schema forbids unknown arguments (the go-sdk infers false for a Go + // struct) rejects a call that has one, so the model should see the same + // rule, and a strict or grammar-constrained backend can enforce it. + AdditionalProperties any `json:"additionalProperties,omitempty"` +} + +// mcpToolParameters is the parameters schema an MCP tool is advertised +// with. It is not a jsonschema.Definition because that type has no +// additionalProperties field, and embedding it would not help: its value +// MarshalJSON would replace this struct's encoding. +type mcpToolParameters struct { + Type jsonschema.DataType `json:"type"` + Properties map[string]jsonschema.Definition `json:"properties"` + Required []string `json:"required,omitempty"` + AdditionalProperties any `json:"additionalProperties,omitempty"` } // CoerceNullableTypes is an exported alias for the same workaround so diff --git a/options.go b/options.go index 8340e73..4eda8d1 100644 --- a/options.go +++ b/options.go @@ -37,6 +37,7 @@ type Options struct { maxAdjustmentAttempts int toolCallResultCallback func(ToolStatus) strictGuidelines bool + strictToolSchemas bool mcpSessions []*mcp.ClientSession guidelines Guidelines mcpPrompts bool @@ -158,6 +159,16 @@ var ( o.strictGuidelines = true } + // EnableStrictToolSchemas sends every tool that can be expressed in strict + // mode with "strict": true and a schema adjusted to the rules strict mode + // needs (see strictTool). A backend that honors it constrains the model's + // arguments to the schema while it generates them, instead of the tool + // rejecting a malformed call afterwards. LocalAI switches its tool grammar + // on for a request with a strict tool; OpenAI enforces the schema. + EnableStrictToolSchemas Option = func(o *Options) { + o.strictToolSchemas = true + } + // EnableAutoPlan enables cogito to automatically use planning if needed EnableAutoPlan Option = func(o *Options) { o.autoPlan = true diff --git a/strict.go b/strict.go new file mode 100644 index 0000000..c0486da --- /dev/null +++ b/strict.go @@ -0,0 +1,194 @@ +package cogito + +import ( + "encoding/json" + "sort" + + "github.com/sashabaranov/go-openai" +) + +// strictTool advertises inner with "strict": true and a schema that follows +// the rules strict mode needs: +// +// - every object sets additionalProperties to false; +// - every property is listed in required, so a property that was optional +// becomes nullable instead, and the model sends null to leave it out. +// +// Execute removes those nulls again before it calls inner, so inner sees an +// omitted argument, as it would without strict mode. A schema that strict +// mode cannot express (an object whose additionalProperties is a schema or +// true, oneOf, a non-object top level) is sent unchanged and not strict. +type strictTool struct { + inner ToolDefinitionInterface +} + +// strictTools wraps each tool in strictTool, once. +func strictTools(tools Tools) Tools { + out := make(Tools, len(tools)) + for i, t := range tools { + if _, done := t.(strictTool); done { + out[i] = t + continue + } + out[i] = strictTool{inner: t} + } + return out +} + +func (t strictTool) Tool() openai.Tool { + tool := t.inner.Tool() + if tool.Function == nil || tool.Function.Parameters == nil { + return tool + } + raw, err := json.Marshal(tool.Function.Parameters) + if err != nil { + return tool + } + var schema map[string]any + if err := json.Unmarshal(raw, &schema); err != nil { + return tool + } + if schema["type"] != "object" || !strictSchema(schema) { + return tool + } + fn := *tool.Function + fn.Parameters = schema + fn.Strict = true + tool.Function = &fn + return tool +} + +func (t strictTool) Execute(args map[string]any) (string, any, error) { + return t.inner.Execute(dropNulls(args)) +} + +// strictSchema rewrites schema in place to follow strict mode's rules and +// reports whether it could. On false the caller must discard schema. +func strictSchema(schema map[string]any) bool { + if _, ok := schema["oneOf"]; ok { + return false + } + if isObjectSchema(schema) { + switch ap := schema["additionalProperties"].(type) { + case nil: + case bool: + if ap { + return false + } + default: + return false // a map type: strict mode cannot describe it + } + schema["additionalProperties"] = false + + props, _ := schema["properties"].(map[string]any) + if props == nil { + props = map[string]any{} + schema["properties"] = props + } + required := map[string]bool{} + if rs, ok := schema["required"].([]any); ok { + for _, r := range rs { + if s, ok := r.(string); ok { + required[s] = true + } + } + } + names := make([]string, 0, len(props)) + for name, p := range props { + ps, ok := p.(map[string]any) + if !ok { + return false // a boolean schema + } + if !strictSchema(ps) { + return false + } + if !required[name] { + props[name] = nullable(ps) + } + names = append(names, name) + } + sort.Strings(names) + req := make([]any, len(names)) + for i, n := range names { + req[i] = n + } + schema["required"] = req + } + if items, ok := schema["items"].(map[string]any); ok && !strictSchema(items) { + return false + } + for _, key := range []string{"anyOf", "allOf"} { + if list, ok := schema[key].([]any); ok { + for _, s := range list { + if m, ok := s.(map[string]any); ok && !strictSchema(m) { + return false + } + } + } + } + for _, key := range []string{"$defs", "definitions"} { + if defs, ok := schema[key].(map[string]any); ok { + for _, d := range defs { + if m, ok := d.(map[string]any); ok && !strictSchema(m) { + return false + } + } + } + } + return true +} + +func isObjectSchema(schema map[string]any) bool { + if schema["type"] == "object" { + return true + } + if ts, ok := schema["type"].([]any); ok { + for _, t := range ts { + if t == "object" { + return true + } + } + } + return false +} + +// nullable returns schema extended to also accept null. +func nullable(schema map[string]any) map[string]any { + switch t := schema["type"].(type) { + case string: + schema["type"] = []any{t, "null"} + case []any: + for _, x := range t { + if x == "null" { + return schema + } + } + schema["type"] = append(t, "null") + default: + return map[string]any{"anyOf": []any{schema, map[string]any{"type": "null"}}} + } + if enum, ok := schema["enum"].([]any); ok { + schema["enum"] = append(enum, nil) + } + return schema +} + +// dropNulls returns args without the null values strict mode makes the +// model send for an argument it leaves out, in nested objects too. +func dropNulls(args map[string]any) map[string]any { + if args == nil { + return nil + } + out := make(map[string]any, len(args)) + for k, v := range args { + switch x := v.(type) { + case nil: + continue + case map[string]any: + out[k] = dropNulls(x) + default: + out[k] = v + } + } + return out +} diff --git a/strict_test.go b/strict_test.go new file mode 100644 index 0000000..240784b --- /dev/null +++ b/strict_test.go @@ -0,0 +1,166 @@ +package cogito + +import ( + "context" + "encoding/json" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/sashabaranov/go-openai" +) + +type shellArgs struct { + Script string `json:"script" jsonschema:"the script to run"` + Timeout int `json:"timeout,omitempty" jsonschema:"seconds"` +} + +// startTypedMCP serves one tool whose schema the go-sdk infers from +// shellArgs, as nib's bash tool does: script required, timeout optional, +// additionalProperties false. The server validates every call against it. +func startTypedMCP() (*mcpsdk.ClientSession, func()) { + impl := &mcpsdk.Implementation{Name: "typed", Version: "0.0.1"} + srv := mcpsdk.NewServer(impl, nil) + mcpsdk.AddTool(srv, &mcpsdk.Tool{Name: "bash", Description: "run a script"}, + func(_ context.Context, _ *mcpsdk.CallToolRequest, in shellArgs) (*mcpsdk.CallToolResult, any, error) { + return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "ran " + in.Script}}}, nil, nil + }) + srvT, clientT := mcpsdk.NewInMemoryTransports() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + go func() { _ = srv.Run(ctx, srvT) }() + sess, err := mcpsdk.NewClient(impl, nil).Connect(ctx, clientT, nil) + Expect(err).ToNot(HaveOccurred()) + return sess, func() { _ = sess.Close(); cancel() } +} + +func paramsOf(t openai.Tool) map[string]any { + raw, err := json.Marshal(t.Function.Parameters) + Expect(err).ToNot(HaveOccurred()) + var m map[string]any + Expect(json.Unmarshal(raw, &m)).To(Succeed()) + return m +} + +var _ = Describe("strict tool schemas", func() { + var ( + sess *mcpsdk.ClientSession + teardown func() + ) + AfterEach(func() { + if teardown != nil { + teardown() + teardown = nil + } + }) + + It("advertises an MCP tool with the additionalProperties its server declares", func() { + sess, teardown = startTypedMCP() + tools, err := mcpToolsFromTransport(context.Background(), sess, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(tools).To(HaveLen(1)) + + p := paramsOf(tools[0].Tool()) + Expect(p).To(HaveKeyWithValue("additionalProperties", false)) + Expect(p["required"]).To(ConsistOf("script")) + Expect(tools[0].Tool().Function.Strict).To(BeFalse(), "strict is opt-in") + }) + + It("sends a strict schema and drops the nulls it makes the model send", func() { + sess, teardown = startTypedMCP() + tools, err := mcpToolsFromTransport(context.Background(), sess, nil) + Expect(err).ToNot(HaveOccurred()) + tool := strictTools(tools)[0] + + def := tool.Tool() + Expect(def.Function.Strict).To(BeTrue()) + p := paramsOf(def) + Expect(p).To(HaveKeyWithValue("additionalProperties", false)) + Expect(p["required"]).To(ConsistOf("script", "timeout")) + timeout := p["properties"].(map[string]any)["timeout"].(map[string]any) + Expect(timeout["type"]).To(ConsistOf("integer", "null")) + + // The server rejects "timeout": null (it wants an integer), so this + // only succeeds if the wrapper leaves the null out. + out, _, err := tool.Execute(map[string]any{"script": "ls", "timeout": nil}) + Expect(err).ToNot(HaveOccurred()) + Expect(out).To(ContainSubstring("ran ls")) + }) + + It("wraps the tools usableTools returns only when the option is set", func() { + sess, teardown = startTypedMCP() + plain, _, _, err := usableTools(nil, NewEmptyFragment(), WithMCPs(sess)) + Expect(err).ToNot(HaveOccurred()) + Expect(plain[0].Tool().Function.Strict).To(BeFalse()) + + strict, _, _, err := usableTools(nil, NewEmptyFragment(), WithMCPs(sess), EnableStrictToolSchemas) + Expect(err).ToNot(HaveOccurred()) + Expect(strict[0].Tool().Function.Strict).To(BeTrue()) + Expect(strictTools(strict)[0]).To(Equal(strict[0]), "wrapping twice is a no-op") + }) +}) + +var _ = Describe("strictSchema", func() { + parse := func(s string) map[string]any { + var m map[string]any + Expect(json.Unmarshal([]byte(s), &m)).To(Succeed()) + return m + } + + It("makes nested optional properties nullable and closes every object", func() { + m := parse(`{"type":"object","properties":{ + "mode":{"type":"string","enum":["a","b"]}, + "opts":{"type":"object","properties":{"deep":{"type":"boolean"}}}, + "list":{"type":"array","items":{"type":"object","properties":{"k":{"type":"string"}},"required":["k"]}} + },"required":["opts"]}`) + Expect(strictSchema(m)).To(BeTrue()) + Expect(m["required"]).To(ConsistOf("mode", "opts", "list")) + + props := m["properties"].(map[string]any) + mode := props["mode"].(map[string]any) + Expect(mode["type"]).To(ConsistOf("string", "null")) + Expect(mode["enum"]).To(ContainElement(BeNil())) + + opts := props["opts"].(map[string]any) + Expect(opts["type"]).To(Equal("object"), "a required property stays non-null") + Expect(opts).To(HaveKeyWithValue("additionalProperties", false)) + deep := opts["properties"].(map[string]any)["deep"].(map[string]any) + Expect(deep["type"]).To(ConsistOf("boolean", "null")) + + item := props["list"].(map[string]any)["items"].(map[string]any) + Expect(item).To(HaveKeyWithValue("additionalProperties", false)) + Expect(item["required"]).To(ConsistOf("k")) + }) + + It("wraps an optional property without a type in anyOf with null", func() { + m := parse(`{"type":"object","properties":{"v":{"anyOf":[{"type":"string"},{"type":"integer"}]}}}`) + Expect(strictSchema(m)).To(BeTrue()) + v := m["properties"].(map[string]any)["v"].(map[string]any) + Expect(v["anyOf"]).To(HaveLen(2)) + Expect(v["anyOf"].([]any)[1]).To(Equal(map[string]any{"type": "null"})) + }) + + DescribeTable("refuses what strict mode cannot express", + func(schema string) { + Expect(strictSchema(parse(schema))).To(BeFalse()) + }, + Entry("a map-typed object", `{"type":"object","additionalProperties":{"type":"string"}}`), + Entry("open objects", `{"type":"object","additionalProperties":true}`), + Entry("oneOf", `{"type":"object","properties":{"x":{"oneOf":[{"type":"string"}]}}}`), + ) + + It("leaves such a tool unchanged and not strict", func() { + inner := &fakeParamsTool{params: map[string]any{"type": "object", "additionalProperties": map[string]any{"type": "string"}}} + def := strictTool{inner: inner}.Tool() + Expect(def.Function.Strict).To(BeFalse()) + Expect(paramsOf(def)["additionalProperties"]).To(Equal(map[string]any{"type": "string"})) + }) +}) + +type fakeParamsTool struct{ params any } + +func (f *fakeParamsTool) Tool() openai.Tool { + return openai.Tool{Type: openai.ToolTypeFunction, Function: &openai.FunctionDefinition{Name: "f", Parameters: f.params}} +} +func (f *fakeParamsTool) Execute(map[string]any) (string, any, error) { return "", nil, nil }