69 lines
1.9 KiB
Go
69 lines
1.9 KiB
Go
package main
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
// TestProcessLLMResponse tests the processing of LLM responses to detect tool usage
|
|
func TestProcessLLMResponse(t *testing.T) {
|
|
// Test data
|
|
testCases := []struct {
|
|
name string
|
|
input string
|
|
expectedOutput string
|
|
expectedToolName string
|
|
expectToolCall bool
|
|
}{
|
|
{
|
|
name: "RegularTextResponse",
|
|
input: "The capital of France is Paris.",
|
|
expectedOutput: "The capital of France is Paris.",
|
|
expectedToolName: "",
|
|
expectToolCall: false,
|
|
},
|
|
{
|
|
name: "ToolCallResponse",
|
|
input: "[TOOL_CALL:get_current_timestamp]",
|
|
expectedOutput: "Used tool: get_current_timestamp",
|
|
expectedToolName: "get_current_timestamp",
|
|
expectToolCall: true,
|
|
},
|
|
{
|
|
name: "AnotherToolCallResponse",
|
|
input: "[TOOL_CALL:send_email]",
|
|
expectedOutput: "Used tool: send_email",
|
|
expectedToolName: "send_email",
|
|
expectToolCall: true,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
// This mimics the logic in processLLMResponse function
|
|
respText := tc.input
|
|
var toolName string
|
|
|
|
if strings.HasPrefix(respText, "[TOOL_CALL:") && strings.HasSuffix(respText, "]") {
|
|
// Extract tool name from the marker
|
|
toolName = strings.TrimPrefix(strings.TrimSuffix(respText, "]"), "[TOOL_CALL:")
|
|
// Remove the marker from the response text
|
|
respText = "Used tool: " + toolName
|
|
}
|
|
|
|
if respText != tc.expectedOutput {
|
|
t.Errorf("Expected output '%s', got '%s'", tc.expectedOutput, respText)
|
|
}
|
|
|
|
if tc.expectToolCall {
|
|
if toolName != tc.expectedToolName {
|
|
t.Errorf("Expected tool name '%s', got '%s'", tc.expectedToolName, toolName)
|
|
}
|
|
} else {
|
|
if toolName != "" {
|
|
t.Errorf("Expected no tool call, but got tool name '%s'", toolName)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
} |