Fix: Refactor LLM response parser to separate tool calls and strip them from text content in openai_compat providers
This commit is contained in:
ManManavadaria
parent
12f0c4a6cf
commit
78319093c9
7 changed files with 787 additions and 267 deletions
|
|
@ -7,6 +7,8 @@ import (
|
|||
"fmt"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/toolcall"
|
||||
)
|
||||
|
||||
// ClaudeCliProvider implements LLMProvider using the claude CLI as a subprocess.
|
||||
|
|
@ -173,28 +175,12 @@ func (p *ClaudeCliProvider) parseClaudeCliResponse(output string) (*LLMResponse,
|
|||
|
||||
// extractToolCalls delegates to the shared extractToolCallsFromText function.
|
||||
func (p *ClaudeCliProvider) extractToolCalls(text string) []ToolCall {
|
||||
return extractToolCallsFromText(text)
|
||||
return toolcall.ExtractToolCallsFromText(text)
|
||||
}
|
||||
|
||||
// stripToolCallsJSON delegates to the shared stripToolCallsFromText function.
|
||||
func (p *ClaudeCliProvider) stripToolCallsJSON(text string) string {
|
||||
return stripToolCallsFromText(text)
|
||||
}
|
||||
|
||||
// findMatchingBrace finds the index after the closing brace matching the opening brace at pos.
|
||||
func findMatchingBrace(text string, pos int) int {
|
||||
depth := 0
|
||||
for i := pos; i < len(text); i++ {
|
||||
if text[i] == '{' {
|
||||
depth++
|
||||
} else if text[i] == '}' {
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return i + 1
|
||||
}
|
||||
}
|
||||
}
|
||||
return pos
|
||||
return toolcall.StripToolCallsFromText(text)
|
||||
}
|
||||
|
||||
// claudeCliJSONResponse represents the JSON output from the claude CLI.
|
||||
|
|
|
|||
|
|
@ -824,158 +824,3 @@ func TestParseClaudeCliResponse_WhitespaceResult(t *testing.T) {
|
|||
t.Errorf("Content = %q, want %q (should be trimmed)", resp.Content, "hello")
|
||||
}
|
||||
}
|
||||
|
||||
// --- extractToolCalls tests ---
|
||||
|
||||
func TestExtractToolCalls_NoToolCalls(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
got := p.extractToolCalls("Just a regular response.")
|
||||
if len(got) != 0 {
|
||||
t.Errorf("extractToolCalls() = %d, want 0", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCalls_WithToolCalls(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := `Here's the result:
|
||||
{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"test","arguments":"{}"}}]}`
|
||||
|
||||
got := p.extractToolCalls(text)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("extractToolCalls() = %d, want 1", len(got))
|
||||
}
|
||||
if got[0].ID != "call_1" {
|
||||
t.Errorf("ID = %q, want %q", got[0].ID, "call_1")
|
||||
}
|
||||
if got[0].Name != "test" {
|
||||
t.Errorf("Name = %q, want %q", got[0].Name, "test")
|
||||
}
|
||||
if got[0].Type != "function" {
|
||||
t.Errorf("Type = %q, want %q", got[0].Type, "function")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCalls_InvalidJSON(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
got := p.extractToolCalls(`{"tool_calls":invalid}`)
|
||||
if len(got) != 0 {
|
||||
t.Errorf("extractToolCalls() with invalid JSON = %d, want 0", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCalls_MultipleToolCalls(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"/tmp/test\"}"}},{"id":"call_2","type":"function","function":{"name":"write_file","arguments":"{\"path\":\"/tmp/out\",\"content\":\"hello\"}"}}]}`
|
||||
|
||||
got := p.extractToolCalls(text)
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("extractToolCalls() = %d, want 2", len(got))
|
||||
}
|
||||
if got[0].Name != "read_file" {
|
||||
t.Errorf("[0].Name = %q, want %q", got[0].Name, "read_file")
|
||||
}
|
||||
if got[1].Name != "write_file" {
|
||||
t.Errorf("[1].Name = %q, want %q", got[1].Name, "write_file")
|
||||
}
|
||||
// Verify arguments were parsed
|
||||
if got[0].Arguments["path"] != "/tmp/test" {
|
||||
t.Errorf("[0].Arguments[path] = %v, want /tmp/test", got[0].Arguments["path"])
|
||||
}
|
||||
if got[1].Arguments["content"] != "hello" {
|
||||
t.Errorf("[1].Arguments[content] = %v, want hello", got[1].Arguments["content"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCalls_UnmatchedBrace(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
got := p.extractToolCalls(`{"tool_calls":[{"id":"call_1"`)
|
||||
if len(got) != 0 {
|
||||
t.Errorf("extractToolCalls() with unmatched brace = %d, want 0", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCalls_ToolCallArgumentsParsing(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"fn","arguments":"{\"num\":42,\"flag\":true,\"name\":\"test\"}"}}]}`
|
||||
|
||||
got := p.extractToolCalls(text)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("len = %d, want 1", len(got))
|
||||
}
|
||||
// Verify different argument types
|
||||
if got[0].Arguments["num"] != float64(42) {
|
||||
t.Errorf("Arguments[num] = %v (%T), want 42", got[0].Arguments["num"], got[0].Arguments["num"])
|
||||
}
|
||||
if got[0].Arguments["flag"] != true {
|
||||
t.Errorf("Arguments[flag] = %v, want true", got[0].Arguments["flag"])
|
||||
}
|
||||
if got[0].Arguments["name"] != "test" {
|
||||
t.Errorf("Arguments[name] = %v, want test", got[0].Arguments["name"])
|
||||
}
|
||||
// Verify raw arguments string is preserved in FunctionCall
|
||||
if got[0].Function.Arguments == "" {
|
||||
t.Error("Function.Arguments should contain raw JSON string")
|
||||
}
|
||||
}
|
||||
|
||||
// --- stripToolCallsJSON tests ---
|
||||
|
||||
func TestStripToolCallsJSON(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := `Let me check the weather.
|
||||
{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"test","arguments":"{}"}}]}
|
||||
Done.`
|
||||
|
||||
got := p.stripToolCallsJSON(text)
|
||||
if strings.Contains(got, "tool_calls") {
|
||||
t.Errorf("should remove tool_calls JSON, got %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "Let me check the weather.") {
|
||||
t.Errorf("should keep text before, got %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "Done.") {
|
||||
t.Errorf("should keep text after, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripToolCallsJSON_NoToolCalls(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := "Just regular text."
|
||||
got := p.stripToolCallsJSON(text)
|
||||
if got != text {
|
||||
t.Errorf("stripToolCallsJSON() = %q, want %q", got, text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripToolCallsJSON_OnlyToolCalls(t *testing.T) {
|
||||
p := NewClaudeCliProvider("/workspace")
|
||||
text := `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"fn","arguments":"{}"}}]}`
|
||||
got := p.stripToolCallsJSON(text)
|
||||
if got != "" {
|
||||
t.Errorf("stripToolCallsJSON() = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
// --- findMatchingBrace tests ---
|
||||
|
||||
func TestFindMatchingBrace(t *testing.T) {
|
||||
tests := []struct {
|
||||
text string
|
||||
pos int
|
||||
want int
|
||||
}{
|
||||
{`{"a":1}`, 0, 7},
|
||||
{`{"a":{"b":2}}`, 0, 13},
|
||||
{`text {"a":1} more`, 5, 12},
|
||||
{`{unclosed`, 0, 0}, // no match returns pos
|
||||
{`{}`, 0, 2}, // empty object
|
||||
{`{{{}}}`, 0, 6}, // deeply nested
|
||||
{`{"a":"b{c}d"}`, 0, 13}, // braces in strings (simplified matcher)
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got := findMatchingBrace(tt.text, tt.pos)
|
||||
if got != tt.want {
|
||||
t.Errorf("findMatchingBrace(%q, %d) = %d, want %d", tt.text, tt.pos, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ import (
|
|||
"fmt"
|
||||
"os/exec"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/toolcall"
|
||||
)
|
||||
|
||||
// CodexCliProvider implements LLMProvider by wrapping the codex CLI as a subprocess.
|
||||
|
|
@ -234,12 +236,12 @@ func (p *CodexCliProvider) parseJSONLEvents(output string) (*LLMResponse, error)
|
|||
content := strings.Join(contentParts, "\n")
|
||||
|
||||
// Extract tool calls from response text (same pattern as ClaudeCliProvider)
|
||||
toolCalls := extractToolCallsFromText(content)
|
||||
toolCalls := toolcall.ExtractToolCallsFromText(content)
|
||||
|
||||
finishReason := "stop"
|
||||
if len(toolCalls) > 0 {
|
||||
finishReason = "tool_calls"
|
||||
content = stripToolCallsFromText(content)
|
||||
content = toolcall.StripToolCallsFromText(content)
|
||||
}
|
||||
|
||||
return &LLMResponse{
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||
"github.com/sipeed/picoclaw/pkg/providers/toolcall"
|
||||
)
|
||||
|
||||
type ToolCall = protocoltypes.ToolCall
|
||||
|
|
@ -152,36 +153,66 @@ func parseResponse(body []byte) (*LLMResponse, error) {
|
|||
}
|
||||
|
||||
choice := apiResponse.Choices[0]
|
||||
toolCalls := make([]ToolCall, 0, len(choice.Message.ToolCalls))
|
||||
for _, tc := range choice.Message.ToolCalls {
|
||||
arguments := make(map[string]interface{})
|
||||
name := ""
|
||||
|
||||
if tc.Function != nil {
|
||||
name = tc.Function.Name
|
||||
if tc.Function.Arguments != "" {
|
||||
if err := json.Unmarshal([]byte(tc.Function.Arguments), &arguments); err != nil {
|
||||
log.Printf("openai_compat: failed to decode tool call arguments for %q: %v", name, err)
|
||||
arguments["raw"] = tc.Function.Arguments
|
||||
}
|
||||
}
|
||||
// Parse structured tool calls from the standard tool_calls array
|
||||
toolCalls := toolcall.ParseStructuredToolCalls(choice.Message.ToolCalls)
|
||||
|
||||
// Extract any additional tool calls from the content text (tool calls in text content)
|
||||
content := choice.Message.Content
|
||||
textToolCalls := toolcall.ExtractToolCallsFromText(content)
|
||||
|
||||
toolCalls = mergeToolCalls(toolCalls, textToolCalls)
|
||||
|
||||
if len(textToolCalls) > 0 {
|
||||
content = toolcall.StripToolCallsFromText(content)
|
||||
}
|
||||
|
||||
toolCalls = append(toolCalls, ToolCall{
|
||||
ID: tc.ID,
|
||||
Name: name,
|
||||
Arguments: arguments,
|
||||
})
|
||||
}
|
||||
// Determine finish reason based on tool calls
|
||||
finishReason := determineFinishReason(choice.FinishReason, len(toolCalls))
|
||||
|
||||
return &LLMResponse{
|
||||
Content: choice.Message.Content,
|
||||
Content: strings.TrimSpace(content),
|
||||
ToolCalls: toolCalls,
|
||||
FinishReason: choice.FinishReason,
|
||||
FinishReason: finishReason,
|
||||
Usage: apiResponse.Usage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// mergeToolCalls merges embedded tool calls with structured tool calls.
|
||||
func mergeToolCalls(structured []ToolCall, embedded []ToolCall) []ToolCall {
|
||||
if len(embedded) == 0 {
|
||||
return structured
|
||||
}
|
||||
|
||||
existingIDs := make(map[string]bool, len(structured))
|
||||
for _, tc := range structured {
|
||||
existingIDs[tc.ID] = true
|
||||
}
|
||||
|
||||
// Append embedded tool calls that don't already exist
|
||||
for _, etc := range embedded {
|
||||
if !existingIDs[etc.ID] {
|
||||
structured = append(structured, etc)
|
||||
}
|
||||
}
|
||||
|
||||
return structured
|
||||
}
|
||||
|
||||
// determineFinishReason determines the appropriate finish reason based on the API response
|
||||
// and whether tool calls are present.
|
||||
func determineFinishReason(apiFinishReason string, toolCallCount int) string {
|
||||
if toolCallCount > 0 && apiFinishReason == "stop" {
|
||||
return "tool_calls"
|
||||
}
|
||||
|
||||
if apiFinishReason != "" {
|
||||
return apiFinishReason
|
||||
}
|
||||
|
||||
return "stop"
|
||||
}
|
||||
|
||||
func normalizeModel(model, apiBase string) string {
|
||||
idx := strings.Index(model, "/")
|
||||
if idx == -1 {
|
||||
|
|
|
|||
|
|
@ -1,72 +0,0 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// extractToolCallsFromText parses tool call JSON from response text.
|
||||
// Both ClaudeCliProvider and CodexCliProvider use this to extract
|
||||
// tool calls that the model outputs in its response text.
|
||||
func extractToolCallsFromText(text string) []ToolCall {
|
||||
start := strings.Index(text, `{"tool_calls"`)
|
||||
if start == -1 {
|
||||
return nil
|
||||
}
|
||||
|
||||
end := findMatchingBrace(text, start)
|
||||
if end == start {
|
||||
return nil
|
||||
}
|
||||
|
||||
jsonStr := text[start:end]
|
||||
|
||||
var wrapper struct {
|
||||
ToolCalls []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal([]byte(jsonStr), &wrapper); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var result []ToolCall
|
||||
for _, tc := range wrapper.ToolCalls {
|
||||
var args map[string]interface{}
|
||||
json.Unmarshal([]byte(tc.Function.Arguments), &args)
|
||||
|
||||
result = append(result, ToolCall{
|
||||
ID: tc.ID,
|
||||
Type: tc.Type,
|
||||
Name: tc.Function.Name,
|
||||
Arguments: args,
|
||||
Function: &FunctionCall{
|
||||
Name: tc.Function.Name,
|
||||
Arguments: tc.Function.Arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// stripToolCallsFromText removes tool call JSON from response text.
|
||||
func stripToolCallsFromText(text string) string {
|
||||
start := strings.Index(text, `{"tool_calls"`)
|
||||
if start == -1 {
|
||||
return text
|
||||
}
|
||||
|
||||
end := findMatchingBrace(text, start)
|
||||
if end == start {
|
||||
return text
|
||||
}
|
||||
|
||||
return strings.TrimSpace(text[:start] + text[end:])
|
||||
}
|
||||
145
pkg/providers/toolcall/tool_call_extract.go
Normal file
145
pkg/providers/toolcall/tool_call_extract.go
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
package toolcall
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||
)
|
||||
|
||||
// parseStructuredToolCalls converts API response tool calls to the internal ToolCall format.
|
||||
func ParseStructuredToolCalls(apiToolCalls []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}) []protocoltypes.ToolCall {
|
||||
if len(apiToolCalls) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
toolCalls := make([]protocoltypes.ToolCall, 0, len(apiToolCalls))
|
||||
for _, tc := range apiToolCalls {
|
||||
if tc.Function == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
arguments := ParseToolCallArguments(tc.Function.Arguments)
|
||||
|
||||
toolCalls = append(toolCalls, protocoltypes.ToolCall{
|
||||
ID: tc.ID,
|
||||
Type: tc.Type,
|
||||
Name: tc.Function.Name,
|
||||
Arguments: arguments,
|
||||
Function: &protocoltypes.FunctionCall{
|
||||
Name: tc.Function.Name,
|
||||
Arguments: tc.Function.Arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
return toolCalls
|
||||
}
|
||||
|
||||
// parseToolCallArguments parses JSON arguments string into a map.
|
||||
func ParseToolCallArguments(argsStr string) map[string]interface{} {
|
||||
if argsStr == "" {
|
||||
return make(map[string]interface{})
|
||||
}
|
||||
|
||||
var arguments map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &arguments); err != nil {
|
||||
log.Printf("openai_compat: failed to decode tool call arguments: %v", err)
|
||||
return map[string]interface{}{"raw": argsStr}
|
||||
}
|
||||
|
||||
return arguments
|
||||
}
|
||||
|
||||
// extractToolCallsFromText parses tool call JSON from response text.
|
||||
// This handles cases where models embed tool calls in the content field.
|
||||
func ExtractToolCallsFromText(text string) []protocoltypes.ToolCall {
|
||||
start := strings.Index(text, `{"tool_calls"`)
|
||||
if start == -1 {
|
||||
return nil
|
||||
}
|
||||
|
||||
end := FindMatchingBrace(text, start)
|
||||
if end == start {
|
||||
return nil
|
||||
}
|
||||
|
||||
jsonStr := text[start:end]
|
||||
|
||||
var wrapper struct {
|
||||
ToolCalls []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal([]byte(jsonStr), &wrapper); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var result []protocoltypes.ToolCall
|
||||
for _, tc := range wrapper.ToolCalls {
|
||||
args := ParseToolCallArguments(tc.Function.Arguments)
|
||||
|
||||
result = append(result, protocoltypes.ToolCall{
|
||||
ID: tc.ID,
|
||||
Type: tc.Type,
|
||||
Name: tc.Function.Name,
|
||||
Arguments: args,
|
||||
Function: &protocoltypes.FunctionCall{
|
||||
Name: tc.Function.Name,
|
||||
Arguments: tc.Function.Arguments,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// FindMatchingBrace finds the index after the closing brace matching the opening brace at pos.
|
||||
func FindMatchingBrace(text string, pos int) int {
|
||||
depth := 0
|
||||
for i := pos; i < len(text); i++ {
|
||||
if text[i] == '{' {
|
||||
depth++
|
||||
} else if text[i] == '}' {
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return i + 1
|
||||
}
|
||||
}
|
||||
}
|
||||
return pos
|
||||
}
|
||||
|
||||
// stripJSONObject removes a JSON object starting with the specified pattern from text.
|
||||
func StripJSONObject(text, pattern string) string {
|
||||
start := strings.Index(text, pattern)
|
||||
if start == -1 {
|
||||
return text
|
||||
}
|
||||
|
||||
end := FindMatchingBrace(text, start)
|
||||
if end == start {
|
||||
return text
|
||||
}
|
||||
|
||||
return strings.TrimSpace(text[:start] + text[end:])
|
||||
}
|
||||
|
||||
// StripToolCallsFromText removes tool call JSON from response text.
|
||||
func StripToolCallsFromText(text string) string {
|
||||
return StripJSONObject(text, `{"tool_calls"`)
|
||||
}
|
||||
583
pkg/providers/toolcall/tool_call_extract_test.go
Normal file
583
pkg/providers/toolcall/tool_call_extract_test.go
Normal file
|
|
@ -0,0 +1,583 @@
|
|||
package toolcall
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||
)
|
||||
|
||||
// --- FindMatchingBrace tests ---
|
||||
|
||||
func TestFindMatchingBrace(t *testing.T) {
|
||||
tests := []struct {
|
||||
text string
|
||||
pos int
|
||||
want int
|
||||
}{
|
||||
{`{"a":1}`, 0, 7},
|
||||
{`{"a":{"b":2}}`, 0, 13},
|
||||
{`text {"a":1} more`, 5, 12},
|
||||
{`{unclosed`, 0, 0}, // no match returns pos
|
||||
{`{}`, 0, 2}, // empty object
|
||||
{`{{{}}}`, 0, 6}, // deeply nested
|
||||
{`{"a":"b{c}d"}`, 0, 13}, // braces in strings (simplified matcher)
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got := FindMatchingBrace(tt.text, tt.pos)
|
||||
if got != tt.want {
|
||||
t.Errorf("FindMatchingBrace(%q, %d) = %d, want %d", tt.text, tt.pos, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- StripJSONObject tests ---
|
||||
|
||||
func TestStripJSONObject(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
pattern string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "removes matching object",
|
||||
text: `before {"tool_calls":[]} after`,
|
||||
pattern: `{"tool_calls"`,
|
||||
want: `before after`,
|
||||
},
|
||||
{
|
||||
name: "pattern not found returns original",
|
||||
text: "no match here",
|
||||
pattern: `{"tool_calls"`,
|
||||
want: "no match here",
|
||||
},
|
||||
{
|
||||
name: "unmatched brace returns original",
|
||||
text: `{"tool_calls"`,
|
||||
pattern: `{"tool_calls"`,
|
||||
want: `{"tool_calls"`,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := StripJSONObject(tt.text, tt.pattern)
|
||||
if got != tt.want {
|
||||
t.Errorf("StripJSONObject() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// --- StripToolCallsFromText tests ---
|
||||
|
||||
func TestStripToolCallsFromText(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "removes tool_calls JSON",
|
||||
text: `Let me check.` + "\n" + `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"fn","arguments":"{}"}}]}`,
|
||||
want: "Let me check.",
|
||||
},
|
||||
{
|
||||
name: "no tool_calls returns original",
|
||||
text: "Just regular text.",
|
||||
want: "Just regular text.",
|
||||
},
|
||||
{
|
||||
name: "only tool_calls returns empty",
|
||||
text: `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"fn","arguments":"{}"}}]}`,
|
||||
want: "",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := StripToolCallsFromText(tt.text)
|
||||
if got != tt.want {
|
||||
t.Errorf("StripToolCallsFromText() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// --- ParseToolCallArguments tests ---
|
||||
|
||||
func TestParseToolCallArguments(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
argsStr string
|
||||
want map[string]interface{}
|
||||
}{
|
||||
{
|
||||
name: "empty string returns empty map",
|
||||
argsStr: "",
|
||||
want: make(map[string]interface{}),
|
||||
},
|
||||
{
|
||||
name: "valid JSON object",
|
||||
argsStr: `{"location":"Tokyo","unit":"celsius"}`,
|
||||
want: map[string]interface{}{
|
||||
"location": "Tokyo",
|
||||
"unit": "celsius",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "valid JSON with nested objects",
|
||||
argsStr: `{"query":{"text":"hello","lang":"en"}}`,
|
||||
want: map[string]interface{}{
|
||||
"query": map[string]interface{}{
|
||||
"text": "hello",
|
||||
"lang": "en",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "valid JSON with arrays",
|
||||
argsStr: `{"items":["a","b","c"],"count":3}`,
|
||||
want: map[string]interface{}{
|
||||
"items": []interface{}{"a", "b", "c"},
|
||||
"count": float64(3),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "valid JSON with numbers",
|
||||
argsStr: `{"temperature":72.5,"humidity":60}`,
|
||||
want: map[string]interface{}{
|
||||
"temperature": 72.5,
|
||||
"humidity": float64(60),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "valid JSON with booleans",
|
||||
argsStr: `{"enabled":true,"active":false}`,
|
||||
want: map[string]interface{}{
|
||||
"enabled": true,
|
||||
"active": false,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "invalid JSON returns raw string in map",
|
||||
argsStr: `{invalid json}`,
|
||||
want: map[string]interface{}{
|
||||
"raw": `{invalid json}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "malformed JSON returns raw string",
|
||||
argsStr: `{"key":value}`,
|
||||
want: map[string]interface{}{
|
||||
"raw": `{"key":value}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "empty JSON object",
|
||||
argsStr: `{}`,
|
||||
want: make(map[string]interface{}),
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ParseToolCallArguments(tt.argsStr)
|
||||
if len(got) != len(tt.want) {
|
||||
t.Errorf("ParseToolCallArguments() length = %d, want %d", len(got), len(tt.want))
|
||||
return
|
||||
}
|
||||
for key, wantVal := range tt.want {
|
||||
gotVal, ok := got[key]
|
||||
if !ok {
|
||||
t.Errorf("ParseToolCallArguments() missing key %q", key)
|
||||
continue
|
||||
}
|
||||
if key == "raw" {
|
||||
// For raw error cases, just check it exists
|
||||
if _, ok := gotVal.(string); !ok {
|
||||
t.Errorf("ParseToolCallArguments() raw value is not string")
|
||||
}
|
||||
} else {
|
||||
// For other cases, do deep comparison
|
||||
if gotVal != wantVal {
|
||||
// Handle nested maps
|
||||
if gotMap, ok := gotVal.(map[string]interface{}); ok {
|
||||
if wantMap, ok := wantVal.(map[string]interface{}); ok {
|
||||
if len(gotMap) != len(wantMap) {
|
||||
t.Errorf("ParseToolCallArguments()[%q] nested map length = %d, want %d", key, len(gotMap), len(wantMap))
|
||||
}
|
||||
continue
|
||||
}
|
||||
}
|
||||
t.Errorf("ParseToolCallArguments()[%q] = %v, want %v", key, gotVal, wantVal)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// --- ExtractToolCallsFromText tests ---
|
||||
|
||||
func TestExtractToolCallsFromText(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
want []protocoltypes.ToolCall
|
||||
wantLen int
|
||||
checkIDs bool
|
||||
}{
|
||||
{
|
||||
name: "no tool_calls returns nil",
|
||||
text: "Just regular text without tool calls",
|
||||
want: nil,
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "text with single tool call",
|
||||
text: `Let me check.` + "\n" + `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\"}"}}]}`,
|
||||
wantLen: 1,
|
||||
checkIDs: true,
|
||||
},
|
||||
{
|
||||
name: "text with multiple tool calls",
|
||||
text: `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"/tmp/test.txt\"}"}},{"id":"call_2","type":"function","function":{"name":"write_file","arguments":"{\"path\":\"/tmp/out.txt\",\"content\":\"hello\"}"}}]}`,
|
||||
wantLen: 2,
|
||||
checkIDs: true,
|
||||
},
|
||||
{
|
||||
name: "tool_calls with empty arguments",
|
||||
text: `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"no_args","arguments":"{}"}}]}`,
|
||||
wantLen: 1,
|
||||
checkIDs: true,
|
||||
},
|
||||
{
|
||||
name: "tool_calls with nested JSON arguments",
|
||||
text: `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"complex","arguments":"{\"query\":{\"text\":\"hello\",\"lang\":\"en\"},\"options\":{\"case\":\"lower\"}}"}}]}`,
|
||||
wantLen: 1,
|
||||
checkIDs: true,
|
||||
},
|
||||
{
|
||||
name: "malformed JSON returns nil",
|
||||
text: `{"tool_calls":[invalid json]}`,
|
||||
want: nil,
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "unclosed brace returns nil",
|
||||
text: `{"tool_calls":[{"id":"call_1"`,
|
||||
want: nil,
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "empty tool_calls array",
|
||||
text: `{"tool_calls":[]}`,
|
||||
want: nil,
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "text before and after tool_calls",
|
||||
text: `Before text.` + "\n" + `{"tool_calls":[{"id":"call_1","type":"function","function":{"name":"test","arguments":"{}"}}]}` + "\n" + `After text.`,
|
||||
wantLen: 1,
|
||||
checkIDs: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ExtractToolCallsFromText(tt.text)
|
||||
if tt.want == nil && got != nil {
|
||||
t.Errorf("ExtractToolCallsFromText() = %v, want nil", got)
|
||||
return
|
||||
}
|
||||
if len(got) != tt.wantLen {
|
||||
t.Errorf("ExtractToolCallsFromText() length = %d, want %d", len(got), tt.wantLen)
|
||||
return
|
||||
}
|
||||
if tt.checkIDs && len(got) > 0 {
|
||||
// Verify structure of first tool call
|
||||
tc := got[0]
|
||||
if tc.ID == "" {
|
||||
t.Error("ExtractToolCallsFromText() tool call ID is empty")
|
||||
}
|
||||
if tc.Name == "" {
|
||||
t.Error("ExtractToolCallsFromText() tool call Name is empty")
|
||||
}
|
||||
if tc.Function == nil {
|
||||
t.Error("ExtractToolCallsFromText() tool call Function is nil")
|
||||
} else {
|
||||
if tc.Function.Name == "" {
|
||||
t.Error("ExtractToolCallsFromText() Function.Name is empty")
|
||||
}
|
||||
if tc.Function.Arguments == "" {
|
||||
t.Error("ExtractToolCallsFromText() Function.Arguments is empty")
|
||||
}
|
||||
}
|
||||
if tc.Arguments == nil {
|
||||
t.Error("ExtractToolCallsFromText() Arguments map is nil")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractToolCallsFromText_Detailed(t *testing.T) {
|
||||
text := `Let me check the weather.` + "\n" + `{"tool_calls":[{"id":"call_123","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\",\"unit\":\"celsius\"}"}}]}`
|
||||
|
||||
toolCalls := ExtractToolCallsFromText(text)
|
||||
if len(toolCalls) != 1 {
|
||||
t.Fatalf("ExtractToolCallsFromText() length = %d, want 1", len(toolCalls))
|
||||
}
|
||||
|
||||
tc := toolCalls[0]
|
||||
if tc.ID != "call_123" {
|
||||
t.Errorf("ToolCall.ID = %q, want %q", tc.ID, "call_123")
|
||||
}
|
||||
if tc.Type != "function" {
|
||||
t.Errorf("ToolCall.Type = %q, want %q", tc.Type, "function")
|
||||
}
|
||||
if tc.Name != "get_weather" {
|
||||
t.Errorf("ToolCall.Name = %q, want %q", tc.Name, "get_weather")
|
||||
}
|
||||
if tc.Function == nil {
|
||||
t.Fatal("ToolCall.Function is nil")
|
||||
}
|
||||
if tc.Function.Name != "get_weather" {
|
||||
t.Errorf("Function.Name = %q, want %q", tc.Function.Name, "get_weather")
|
||||
}
|
||||
if tc.Function.Arguments != `{"location":"Tokyo","unit":"celsius"}` {
|
||||
t.Errorf("Function.Arguments = %q, want %q", tc.Function.Arguments, `{"location":"Tokyo","unit":"celsius"}`)
|
||||
}
|
||||
if tc.Arguments["location"] != "Tokyo" {
|
||||
t.Errorf("Arguments[location] = %v, want Tokyo", tc.Arguments["location"])
|
||||
}
|
||||
if tc.Arguments["unit"] != "celsius" {
|
||||
t.Errorf("Arguments[unit] = %v, want celsius", tc.Arguments["unit"])
|
||||
}
|
||||
}
|
||||
|
||||
// --- ParseStructuredToolCalls tests ---
|
||||
|
||||
func TestParseStructuredToolCalls(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
apiToolCalls []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}
|
||||
wantLen int
|
||||
}{
|
||||
{
|
||||
name: "empty array returns nil",
|
||||
apiToolCalls: []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{},
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "single tool call",
|
||||
apiToolCalls: []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"NYC"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "multiple tool calls",
|
||||
apiToolCalls: []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "read_file",
|
||||
Arguments: `{"path":"/tmp/a.txt"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "call_2",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "write_file",
|
||||
Arguments: `{"path":"/tmp/b.txt","content":"hello"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
wantLen: 2,
|
||||
},
|
||||
{
|
||||
name: "nil function is skipped",
|
||||
apiToolCalls: []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: nil,
|
||||
},
|
||||
{
|
||||
ID: "call_2",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "valid_call",
|
||||
Arguments: `{}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "empty arguments string",
|
||||
apiToolCalls: []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "no_args",
|
||||
Arguments: "",
|
||||
},
|
||||
},
|
||||
},
|
||||
wantLen: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ParseStructuredToolCalls(tt.apiToolCalls)
|
||||
if tt.wantLen == 0 {
|
||||
if len(got) > 0 {
|
||||
t.Errorf("ParseStructuredToolCalls() = %v, want nil or empty", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if len(got) != tt.wantLen {
|
||||
t.Errorf("ParseStructuredToolCalls() length = %d, want %d", len(got), tt.wantLen)
|
||||
return
|
||||
}
|
||||
// Verify structure
|
||||
for i, tc := range got {
|
||||
if tc.ID == "" {
|
||||
t.Errorf("ParseStructuredToolCalls()[%d].ID is empty", i)
|
||||
}
|
||||
if tc.Name == "" {
|
||||
t.Errorf("ParseStructuredToolCalls()[%d].Name is empty", i)
|
||||
}
|
||||
if tc.Function == nil {
|
||||
t.Errorf("ParseStructuredToolCalls()[%d].Function is nil", i)
|
||||
}
|
||||
if tc.Arguments == nil {
|
||||
t.Errorf("ParseStructuredToolCalls()[%d].Arguments is nil", i)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseStructuredToolCalls_Detailed(t *testing.T) {
|
||||
apiToolCalls := []struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function *struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}{
|
||||
{
|
||||
ID: "call_abc",
|
||||
Type: "function",
|
||||
Function: &struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"city":"SF","unit":"fahrenheit"}`,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
toolCalls := ParseStructuredToolCalls(apiToolCalls)
|
||||
if len(toolCalls) != 1 {
|
||||
t.Fatalf("ParseStructuredToolCalls() length = %d, want 1", len(toolCalls))
|
||||
}
|
||||
|
||||
tc := toolCalls[0]
|
||||
if tc.ID != "call_abc" {
|
||||
t.Errorf("ToolCall.ID = %q, want %q", tc.ID, "call_abc")
|
||||
}
|
||||
if tc.Type != "function" {
|
||||
t.Errorf("ToolCall.Type = %q, want %q", tc.Type, "function")
|
||||
}
|
||||
if tc.Name != "get_weather" {
|
||||
t.Errorf("ToolCall.Name = %q, want %q", tc.Name, "get_weather")
|
||||
}
|
||||
if tc.Function == nil {
|
||||
t.Fatal("ToolCall.Function is nil")
|
||||
}
|
||||
if tc.Function.Name != "get_weather" {
|
||||
t.Errorf("Function.Name = %q, want %q", tc.Function.Name, "get_weather")
|
||||
}
|
||||
if tc.Function.Arguments != `{"city":"SF","unit":"fahrenheit"}` {
|
||||
t.Errorf("Function.Arguments = %q, want %q", tc.Function.Arguments, `{"city":"SF","unit":"fahrenheit"}`)
|
||||
}
|
||||
if tc.Arguments["city"] != "SF" {
|
||||
t.Errorf("Arguments[city] = %v, want SF", tc.Arguments["city"])
|
||||
}
|
||||
if tc.Arguments["unit"] != "fahrenheit" {
|
||||
t.Errorf("Arguments[unit] = %v, want fahrenheit", tc.Arguments["unit"])
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue