diff --git a/toolschema.go b/toolschema.go new file mode 100644 index 00000000..275d94b9 --- /dev/null +++ b/toolschema.go @@ -0,0 +1,266 @@ +package tok + +import ( + "encoding/json" + "sort" + "strings" +) + +// Tool-catalog compression ("shrink"). +// +// A large MCP/OpenAI tool catalog is paid on every request: names, parameter +// schemas, and — dominantly — free-text descriptions. ShrinkToolCatalog +// reduces that cost while preserving, byte-for-byte, every token the model +// needs to SELECT a tool and CONSTRUCT valid arguments: +// +// kept exactly: tool names, parameter names, "type", "enum", "required", +// "default", "const", "$ref" target strings, "format" +// dropped: "$schema"/"$id"/"title"/"examples"/"$comment"/"x-*" +// reduced (lossy): long descriptions -> lead sentence plus every sentence +// carrying a constraint marker (must/required/max/format/…) +// +// Correctness boundary: a structural-profile equality test proves the +// selection surface survives; behavioral selection equivalence is model- +// visible and stays inferred. Over-keep is the rule — short descriptions are +// never touched, so a constraint phrased without a recognized marker still +// survives whole. Fail-open: any parse problem returns the input unchanged. + +const ( + // toolschemaShortDesc keeps descriptions up to this many bytes whole. + toolschemaShortDesc = 600 + // toolschemaLeadSentences caps how many lead sentences survive. +) + +// constraintMarkers are substrings whose presence in a description sentence +// signals argument-construction information that must survive compression. +var constraintMarkers = []string{ + "must", "cannot", "required", "rejected", "invalid", "not allowed", + "exactly one", "at least", "at most", "only ", "never", + "max", "maximum", "min", "minimum", "limit", "between", "range", + "greater than", "less than", "over ", "under ", "above", "below", + "unique", "mutually exclusive", "either", "neither", + "format", "iso", "rfc", "absolute", "relative path", "url", "utf-8", + "default", "defaults to", "optional", "ignored", "deprecated", +} + +// ToolShrinkStats reports one tool's reduction. +type ToolShrinkStats struct { + Name string `json:"name"` + Before int `json:"bytes_before"` + After int `json:"bytes_after"` + DescBefore int `json:"desc_bytes_before"` + DescAfter int `json:"desc_bytes_after"` +} + +// LintToolCatalog reports per-tool reductions without committing to them. +func LintToolCatalog(catalog string) ([]ToolShrinkStats, bool) { + out, ok := shrinkCatalog([]byte(catalog), true) + if !ok { + return nil, false + } + stats, ok := out.([]ToolShrinkStats) + if !ok { + return nil, false + } + return stats, true +} + +// ShrinkToolCatalog compresses an OpenAI-style function-tool catalog +// ([{"type":"function","function":{...}}, ...]). Fail-open: on any problem +// the input is returned unchanged with ok=false. +func ShrinkToolCatalog(catalog string) (string, bool) { + out, ok := shrinkCatalog([]byte(catalog), false) + if !ok { + return catalog, false + } + b, err := json.Marshal(out) + if err != nil { + return catalog, false + } + if len(b) >= len(catalog) { + return catalog, false // not smaller: pass through, claim nothing + } + return string(b), true +} + +func shrinkCatalog(raw []byte, lintOnly bool) (interface{}, bool) { + var tools []map[string]json.RawMessage + if json.Unmarshal(raw, &tools) != nil || len(tools) == 0 { + return nil, false + } + stats := make([]ToolShrinkStats, 0, len(tools)) + out := make([]map[string]json.RawMessage, 0, len(tools)) + changed := false + + for _, tool := range tools { + fnRaw, has := tool["function"] + if !has { + out = append(out, tool) + continue + } + before := len(fnRaw) + fn, err := shrinkFunctionDef(fnRaw) + if err != nil { + out = append(out, tool) // fail open per-tool + continue + } + after := len(fn) + var name string + var fnMap map[string]json.RawMessage + if err := json.Unmarshal(fn, &fnMap); err == nil { + if n, ok := fnMap["name"]; ok { + if uerr := json.Unmarshal(n, &name); uerr != nil { + name = "" + } + } + } + descB, descA := countDescBytes(fnRaw), len(valueOf(fnMap["description"])) + stats = append(stats, ToolShrinkStats{Name: name, Before: before, After: after, DescBefore: descB, DescAfter: descA}) + if after < before { + changed = true + } + if !lintOnly { + newTool := map[string]json.RawMessage{} + for k, v := range tool { + if k == "function" { + v = fn + } + newTool[k] = v + } + out = append(out, newTool) + } + } + if lintOnly { + return stats, true + } + if !changed { + return nil, false + } + return out, true +} + +func shrinkFunctionDef(raw json.RawMessage) (json.RawMessage, error) { + var fn map[string]interface{} + if err := json.Unmarshal(raw, &fn); err != nil { + return nil, err + } + if desc, ok := fn["description"].(string); ok && len(desc) > toolschemaShortDesc { + fn["description"] = shrinkDescription(desc) + } + params, ok := fn["parameters"].(map[string]interface{}) + if !ok { + b, err := json.Marshal(fn) + return b, err + } + cleanSchema(params) + b, err := json.Marshal(fn) + return b, err +} + +// cleanSchema recursively drops annotation-only keys from a JSON-schema tree. +func cleanSchema(node map[string]interface{}) { + for k := range node { + switch k { + case "$schema", "$id", "$comment", "title", "examples", "exclusiveMaximumDoc", "exclusiveMinimumDoc": + delete(node, k) + default: + if strings.HasPrefix(k, "x-") { + delete(node, k) + } + } + } + for _, k := range sortedKeys(node) { + v := node[k] + switch tv := v.(type) { + case map[string]interface{}: + cleanSchema(tv) + case []interface{}: + for _, item := range tv { + if m, ok := item.(map[string]interface{}); ok { + cleanSchema(m) + } + } + case string: + _ = tv // keep: types/enums/refs/formats/descriptions-of-leaf survive here + } + } + if props, ok := node["properties"].(map[string]interface{}); ok { + for _, pk := range sortedKeys(props) { + if pv, ok := props[pk].(map[string]interface{}); ok { + // Leaf property description: shorten when long. + if d, ok := pv["description"].(string); ok && len(d) > toolschemaShortDesc { + pv["description"] = shrinkDescription(d) + } + if items, ok := pv["items"].(map[string]interface{}); ok { + if d, ok := items["description"].(string); ok && len(d) > toolschemaShortDesc { + items["description"] = shrinkDescription(d) + } + } + } + } + } +} + +func shrinkDescription(desc string) string { + sentences := splitSentences(desc) + if len(sentences) <= 1 { + return desc + } + keep := []string{sentences[0]} // lead sentence always survives + for _, s := range sentences[1:] { + low := strings.ToLower(s) + for _, marker := range constraintMarkers { + if strings.Contains(low, marker) { + keep = append(keep, s) + break + } + } + } + out := strings.Join(keep, " ") + if len(out) >= len(desc) { + return desc // reduction would be zero or negative: keep original + } + return out +} + +func splitSentences(s string) []string { + var out []string + cur := strings.Builder{} + for _, r := range s { + cur.WriteRune(r) + if r == '.' || r == '!' || r == '?' { + out = append(out, strings.TrimSpace(cur.String())) + cur.Reset() + } + } + if t := strings.TrimSpace(cur.String()); t != "" { + out = append(out, t) + } + return out +} + +func sortedKeys(m map[string]interface{}) []string { + ks := make([]string, 0, len(m)) + for k := range m { + ks = append(ks, k) + } + sort.Strings(ks) + return ks +} + +func countDescBytes(fnRaw json.RawMessage) int { + var fn struct { + Description string `json:"description"` + } + if json.Unmarshal(fnRaw, &fn) == nil { + return len(fn.Description) + } + return 0 +} + +func valueOf(raw json.RawMessage) string { + if raw == nil { + return "" + } + return string(raw) +} diff --git a/toolschema_test.go b/toolschema_test.go new file mode 100644 index 00000000..b34ce781 --- /dev/null +++ b/toolschema_test.go @@ -0,0 +1,187 @@ +package tok + +import ( + "encoding/json" + "reflect" + "strings" + "testing" +) + +const bigDesc = "Reads a file from the local filesystem. " + + "The path parameter must be an absolute path. " + + "Lines longer than 2000 characters are truncated. " + + "You cannot read binary files with this tool. " + + "At most 100 lines are returned by default. " + + "The result includes line numbers as a gutter prefix on every line. " + + "Performance remains fast even on very large inputs of many megabytes. " + + "Files are read from the workspace sandbox when one is configured. " + + "Symlinks are followed unless they escape the allowed directories. " + + "Encoding must be valid UTF-8; ISO-8859-1 input is rejected with a clear error message. " + + "When the same file is requested repeatedly results are served from an internal cache. " + + "This tool is great for exploring code quickly." + +func sampleCatalog() string { + catalog := []map[string]interface{}{{ + "type": "function", + "function": map[string]interface{}{ + "name": "read_file", + "description": bigDesc, + "parameters": map[string]interface{}{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "Read file parameters", + "examples": []interface{}{map[string]string{"path": "/x"}}, + "type": "object", + "required": []string{"path"}, + "additionalProperties": false, + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "Absolute path to the file. Must exist on disk. Max length 4096 bytes.", + }, + "limit": map[string]interface{}{ + "type": "integer", + "default": 100, + "maximum": 500, + }, + }, + }, + }, + }} + b, _ := json.Marshal(catalog) + return string(b) +} + +// selectionProfile extracts the structural surface that must survive +// byte-for-byte: names, param names, types, enums, required, defaults. +func selectionProfile(t *testing.T, catalog string) interface{} { + t.Helper() + var tools []struct { + Function struct { + Name string `json:"name"` + Parameters map[string]interface{} `json:"parameters"` + } `json:"function"` + } + if err := json.Unmarshal([]byte(catalog), &tools); err != nil { + t.Fatalf("parse catalog: %v", err) + } + var profile []interface{} + for _, tool := range tools { + props := map[string]interface{}{} + if params, ok := tool.Function.Parameters["properties"].(map[string]interface{}); ok { + for name, p := range params { + pm := p.(map[string]interface{}) + entry := map[string]interface{}{"type": pm["type"]} + for _, k := range []string{"enum", "required", "default", "maximum", "minimum", "format"} { + if v, ok := pm[k]; ok { + entry[k] = v + } + } + props[name] = entry + } + } + req, _ := tool.Function.Parameters["required"].([]interface{}) + profile = append(profile, map[string]interface{}{ + "name": tool.Function.Name, + "required": req, + "props": props, + "addProps": tool.Function.Parameters["additionalProperties"], + "paramType": tool.Function.Parameters["type"], + }) + } + return profile +} + +func TestShrinkKeepsSelectionSurfaceIdentical(t *testing.T) { + in := sampleCatalog() + out, ok := ShrinkToolCatalog(in) + if !ok { + t.Fatal("expected shrink to change a bloated catalog") + } + before := selectionProfile(t, in) + after := selectionProfile(t, out) + if !reflect.DeepEqual(before, after) { + t.Fatalf("selection surface changed:\nbefore=%#v\nafter=%#v", before, after) + } +} + +func TestShrinkReducesLongDescriptions(t *testing.T) { + out, ok := ShrinkToolCatalog(sampleCatalog()) + if !ok { + t.Fatal("no reduction") + } + var tools []struct { + Function struct { + Description string `json:"description"` + Parameters struct { + Properties struct { + Path struct { + Description string `json:"description"` + } `json:"path"` + } `json:"properties"` + } `json:"parameters"` + } `json:"function"` + } + if err := json.Unmarshal([]byte(out), &tools); err != nil { + t.Fatal(err) + } + d := tools[0].Function.Description + if len(d) >= len(bigDesc) { + t.Fatalf("description not reduced: %q", d) + } + if !strings.HasPrefix(d, "Reads a file") { + t.Fatalf("lead sentence lost: %q", d) + } + for _, want := range []string{"must be an absolute path", "cannot read binary", "At most 100"} { + if !strings.Contains(d, want) { + t.Fatalf("constraint sentence lost (%q): %q", want, d) + } + } + if strings.Contains(d, "great for exploring") { + t.Fatalf("flattery should have been dropped: %q", d) + } +} + +func TestShrinkDropsAnnotationsKeepsDefaults(t *testing.T) { + out, _ := ShrinkToolCatalog(sampleCatalog()) + if strings.Contains(out, "$schema") || strings.Contains(out, `"title"`) || strings.Contains(out, "examples") { + t.Fatalf("annotation metadata survived: %s", out) + } + if !strings.Contains(out, `"default":100`) && !strings.Contains(out, `"default": 100`) { + t.Fatalf("default dropped: %s", out) + } +} + +func TestShrinkFailOpenOnGarbage(t *testing.T) { + for _, in := range []string{"", "not json", `{"a":1}`, `[]`} { + if out, changed := ShrinkToolCatalog(in); changed || out != in { + t.Fatalf("input %q must pass through unchanged", in) + } + } +} + +func TestShrinkPassThroughWhenNotSmaller(t *testing.T) { + in := `[{"type":"function","function":{"name":"tiny","description":"ok","parameters":{"type":"object"}}}]` + out, changed := ShrinkToolCatalog(in) + if changed || out != in { + t.Fatalf("not-smaller catalogs must be returned unchanged (changed=%v)", changed) + } +} + +func TestShrinkShortDescriptionUntouched(t *testing.T) { + in := `[{"type":"function","function":{"name":"f","description":"Does a thing. Optional: set mode.","parameters":{"type":"object"}}}]` + out, changed := ShrinkToolCatalog(in) + _ = out + if changed { + t.Fatal("short descriptions must be kept whole (fail-open over-keep)") + } +} + +func TestLintToolCatalogReports(t *testing.T) { + stats, ok := LintToolCatalog(sampleCatalog()) + if !ok || len(stats) != 1 { + t.Fatalf("lint = %v %v", stats, ok) + } + if stats[0].Name != "read_file" || stats[0].After >= stats[0].Before { + t.Fatalf("stats = %+v", stats[0]) + } +}