diff --git a/components/tool/middlewares/jsonfix/example/main.go b/components/tool/middlewares/jsonfix/example/main.go new file mode 100644 index 0000000..26ef6c0 --- /dev/null +++ b/components/tool/middlewares/jsonfix/example/main.go @@ -0,0 +1,75 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// This example shows how to configure the jsonfix middleware on a ToolsNode +// to repair invalid JSON arguments before invoking a local tool. +// Run: go run ./components/tool/middlewares/jsonfix/example +package main + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + + jsonfix "github.com/cloudwego/eino-examples/components/tool/middlewares/jsonfix" +) + +type greetReq struct { + Name string `json:"name"` +} + +type greetResp struct { + Greeting string `json:"greeting"` +} + +func main() { + ctx := context.Background() + + // Define a simple local tool with JSON input/output using InferTool. + greeter, _ := utils.InferTool("greeter", "greet by name", func(ctx context.Context, in *greetReq) (*greetResp, error) { + return &greetResp{Greeting: "Hello, " + in.Name}, nil + }) + + // Create ToolsNode and register the jsonfix middleware. + tn, _ := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{greeter}, + ToolCallMiddlewares: []compose.ToolMiddleware{jsonfix.Middleware()}, + }) + + // Craft an Assistant message with an invalid JSON argument to simulate LLM output. + msg := schema.AssistantMessage("", nil) + msg.ToolCalls = []schema.ToolCall{{ + ID: "1", + Function: schema.FunctionCall{ + Name: "greeter", + Arguments: "noise <|FunctionCallBegin|>{\"name\":\"Alice\"1\"\"}<|FunctionCallEnd|>", + }, + }} + + // ToolsNode invokes the tool. Middleware repairs the argument first. + outs, err := tn.Invoke(ctx, msg) + if err != nil { + fmt.Println("error:", err) + return + } + for _, o := range outs { + fmt.Println("tool:", o.ToolName, "id:", o.ToolCallID, "content:", o.Content) + } +} diff --git a/components/tool/middlewares/jsonfix/json_fix_middleware.go b/components/tool/middlewares/jsonfix/json_fix_middleware.go new file mode 100644 index 0000000..f155646 --- /dev/null +++ b/components/tool/middlewares/jsonfix/json_fix_middleware.go @@ -0,0 +1,113 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package jsonfix provides a ToolMiddleware for Eino's ToolsNode that +// repairs malformed JSON arguments produced by LLMs before tool execution. +// +// Usage: +// +// conf := &compose.ToolsNodeConfig{ +// Tools: []tool.BaseTool{yourTool}, +// ToolCallMiddlewares: []compose.ToolMiddleware{jsonfix.Middleware()}, +// } +// +// Behavior: +// - Fast-path returns original arguments when already valid JSON. +// - Strips common LLM artifacts and isolates the first {...} region. +// - Applies robust fix using jsonrepair only when input is invalid. +// - Safe for both invokable and streamable tools. +package jsonfix + +import ( + "context" + "encoding/json" + "strings" + + "github.com/cloudwego/eino/compose" + "github.com/kaptinlin/jsonrepair" +) + +// FixJSON is a helper endpoint that returns repaired JSON for diagnostics +// or standalone use. ToolsNode should prefer Invokable/Streamable middleware. +func FixJSON(ctx context.Context, in *compose.ToolInput) (*compose.ToolOutput, error) { + fixed := repair(in.Arguments) + return &compose.ToolOutput{Result: fixed}, nil +} + +// Invokable wraps a non-stream tool endpoint to sanitize JSON arguments. +// Register via ToolCallMiddlewares to apply automatically to invokable tools. +func Invokable(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, in *compose.ToolInput) (*compose.ToolOutput, error) { + in.Arguments = repair(in.Arguments) + return next(ctx, in) + } +} + +// Streamable wraps a stream tool endpoint to sanitize JSON arguments. +// Register via ToolCallMiddlewares to apply automatically to streamable tools. +func Streamable(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, in *compose.ToolInput) (*compose.StreamToolOutput, error) { + in.Arguments = repair(in.Arguments) + return next(ctx, in) + } +} + +// Middleware bundles both invokable and streamable wrappers for convenience. +func Middleware() compose.ToolMiddleware { + return compose.ToolMiddleware{Invokable: Invokable, Streamable: Streamable} +} + +// repair attempts minimal work first (validity check, region isolation) and +// only uses jsonrepair when necessary. It trims common LLM artifacts. +func repair(input string) string { + s := strings.TrimSpace(input) + // Fast-path: valid JSON as-is + if strings.HasPrefix(s, "{") && strings.HasSuffix(s, "}") && json.Valid([]byte(s)) { + return s + } + + // Isolate JSON object region if present; strip noise if object-only is valid + i := strings.IndexByte(s, '{') + j := strings.LastIndexByte(s, '}') + if i >= 0 && j >= i { + sub := s[i : j+1] + if json.Valid([]byte(sub)) { + return sub + } + s = sub + } + + // Remove common LLM artifacts + s = strings.TrimPrefix(s, "<|FunctionCallBegin|>") + s = strings.TrimSuffix(s, "<|FunctionCallEnd|>") + s = strings.TrimPrefix(s, "") + + // Attempt robust repair only when invalid + if json.Valid([]byte(s)) { + return s + } + // Heuristic: add missing leading/trailing brace if one side exists + if !strings.HasPrefix(s, "{") && strings.HasSuffix(s, "}") { + s = "{" + s + } else if strings.HasPrefix(s, "{") && !strings.HasSuffix(s, "}") { + s = s + "}" + } + out, err := jsonrepair.JSONRepair(s) + if err != nil { + return s + } + return out +} diff --git a/components/tool/middlewares/jsonfix/json_fix_middleware_test.go b/components/tool/middlewares/jsonfix/json_fix_middleware_test.go new file mode 100644 index 0000000..3182ef2 --- /dev/null +++ b/components/tool/middlewares/jsonfix/json_fix_middleware_test.go @@ -0,0 +1,185 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package jsonfix + +import ( + "context" + "encoding/json" + "testing" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +func echoEndpoint(_ context.Context, in *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: in.Arguments}, nil +} + +func TestInvokableMiddleware_RepairsJSON(t *testing.T) { + mw := Invokable + chained := mw(echoEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + Arguments: "noise <|FunctionCallBegin|>{\"a\":1}<|FunctionCallEnd|> more", + CallID: "id1", + } + + out, err := chained(context.Background(), input) + if err != nil { + t.Fatalf("err: %v", err) + } + var m map[string]int + if e := json.Unmarshal([]byte(out.Result), &m); e != nil || m["a"] != 1 { + t.Fatalf("repair failed: %v %v", out.Result, e) + } +} + +func TestStreamableMiddleware_RepairsJSON(t *testing.T) { + mw := Streamable + var captured string + next := func(_ context.Context, in *compose.ToolInput) (*compose.StreamToolOutput, error) { + captured = in.Arguments + return &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{"ok"})}, nil + } + chained := mw(next) + _, err := chained(context.Background(), &compose.ToolInput{Name: "t", Arguments: "{\"text\": \"He said \"hello\" to me\"}"}) + if err != nil { + t.Fatalf("err: %v", err) + } + var m map[string]string + if e := json.Unmarshal([]byte(captured), &m); e != nil || m["text"] != "He said \"hello\" to me" { + t.Fatalf("repair failed: %v %v", captured, e) + } +} + +func TestInvokableMiddleware_NoChangeForValidJSON(t *testing.T) { + mw := Invokable + chained := mw(echoEndpoint) + + original := "{\"a\":1}" + out, err := chained(context.Background(), &compose.ToolInput{Name: "t", Arguments: original}) + if err != nil { + t.Fatalf("err: %v", err) + } + if out.Result != original { + t.Fatalf("should not change valid json: got %s", out.Result) + } +} + +func TestInvokableMiddleware_MissingBracesAndUnicode(t *testing.T) { + mw := Invokable + chained := mw(echoEndpoint) + + inputs := []string{ + "{\"key\":\"value\"", // missing tail + "\"key\":\"value\"}", // missing head + "{\"emoji\": \"😀😎\"}", + "{\"text\": \"line1\nline2\tTabbed\"}", + } + for _, in := range inputs { + out, err := chained(context.Background(), &compose.ToolInput{Name: "t", Arguments: in}) + if err != nil { + t.Fatalf("err: %v", err) + } + var any map[string]string + if e := json.Unmarshal([]byte(out.Result), &any); e != nil { + t.Fatalf("unmarshal failed: %v for %v", e, out.Result) + } + } +} + +type spyInvokable struct{} + +func (s *spyInvokable) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "spy", Desc: ""}, nil +} +func (s *spyInvokable) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return argumentsInJSON, nil +} + +func TestToolsNodeMiddleware_RepairsInvokable(t *testing.T) { + tn, err := compose.NewToolNode(context.Background(), &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&spyInvokable{}}, + ToolCallMiddlewares: []compose.ToolMiddleware{Middleware()}, + }) + if err != nil { + t.Fatalf("new tools node err: %v", err) + } + msg := schema.AssistantMessage("", nil) + msg.ToolCalls = []schema.ToolCall{{ + ID: "1", + Function: schema.FunctionCall{Name: "spy", Arguments: "garbage {\"x\": \"a \"quote\" b\"} tail"}, + }} + outs, err := tn.Invoke(context.Background(), msg) + if err != nil { + t.Fatalf("invoke err: %v", err) + } + if len(outs) != 1 { + t.Fatalf("unexpected outs: %d", len(outs)) + } + var m map[string]string + if e := json.Unmarshal([]byte(outs[0].Content), &m); e != nil || m["x"] != "a \"quote\" b" { + t.Fatalf("repair failed: %v %v", outs[0].Content, e) + } +} + +func TestInvokableMiddleware_UnescapedBackslashesAndSingleQuotesAndTrailingComma(t *testing.T) { + mw := Invokable + chained := mw(echoEndpoint) + + cases := []string{ + // unescaped backslashes in Windows path + "{\"path\": \"C:\\Users\\name\\file.txt\"}", + // single-quoted JSON + "{'a': 'b'}", + // trailing comma + "{\"a\": 1, \"b\": 2, }", + } + + for _, in := range cases { + out, err := chained(context.Background(), &compose.ToolInput{Name: "t", Arguments: in}) + if err != nil { + t.Fatalf("err: %v", err) + } + var any map[string]any + if e := json.Unmarshal([]byte(out.Result), &any); e != nil { + t.Fatalf("unmarshal failed: %v for %v", e, out.Result) + } + } +} + +func TestInvokableMiddleware_RawQuotesInsideStringValue(t *testing.T) { + mw := Invokable + chained := mw(echoEndpoint) + + // invalid JSON: raw quotes inside value + in := "{\"a\":\"b\"1\"\"}" + out, err := chained(context.Background(), &compose.ToolInput{Name: "t", Arguments: in}) + if err != nil { + t.Fatalf("err: %v", err) + } + var m map[string]string + if e := json.Unmarshal([]byte(out.Result), &m); e != nil { + t.Fatalf("unmarshal failed: %v for %v", e, out.Result) + } + v := m["a"] + if v == "" || v[0] != 'b' { + t.Fatalf("unexpected repaired value: %q", v) + } +}