You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
114 lines
3.7 KiB
Go
114 lines
3.7 KiB
Go
/*
|
|
* 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, "<think>")
|
|
|
|
// 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
|
|
}
|