arp/arp_agent/llm_test.go

package main

import (
	"encoding/json"
	"fmt"
	"testing"

	"github.com/sashabaranov/go-openai"
)

// TestLLM_ConvertMCPToolsToOpenAI tests the MCP to OpenAI tool conversion
func TestLLM_ConvertMCPToolsToOpenAI(t *testing.T) {
	tests := []struct {
		name     string
		mcpTools []Tool
		wantLen  int
	}{
		{
			name:     "EmptyTools",
			mcpTools: []Tool{},
			wantLen:  0,
		},
		{
			name: "SingleTool",
			mcpTools: []Tool{
				{
					Name:        "introspect",
					Description: "Discover the GraphQL schema",
					InputSchema: InputSchema{
						Type: "object",
						Properties: map[string]Property{
							"typeName": {Type: "string", Description: "The type to introspect"},
						},
						Required:             []string{},
						AdditionalProperties: false,
					},
				},
			},
			wantLen: 1,
		},
		{
			name: "MultipleTools",
			mcpTools: []Tool{
				{
					Name:        "query",
					Description: "Execute a GraphQL query",
					InputSchema: InputSchema{
						Type: "object",
						Properties: map[string]Property{
							"query": {Type: "string", Description: "The GraphQL query"},
						},
						Required:             []string{"query"},
						AdditionalProperties: false,
					},
				},
				{
					Name:        "mutate",
					Description: "Execute a GraphQL mutation",
					InputSchema: InputSchema{
						Type: "object",
						Properties: map[string]Property{
							"mutation": {Type: "string", Description: "The GraphQL mutation"},
						},
						Required:             []string{"mutation"},
						AdditionalProperties: false,
					},
				},
			},
			wantLen: 2,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tools := ConvertMCPToolsToOpenAI(tt.mcpTools)
			if len(tools) != tt.wantLen {
				t.Errorf("Expected %d tools, got %d", tt.wantLen, len(tools))
			}

			// Verify tool conversion details
			for i, tool := range tools {
				if tool.Type != openai.ToolTypeFunction {
					t.Errorf("Tool %d: Expected type %s, got %s", i, openai.ToolTypeFunction, tool.Type)
				}
				if tool.Function.Name != tt.mcpTools[i].Name {
					t.Errorf("Tool %d: Expected name %s, got %s", i, tt.mcpTools[i].Name, tool.Function.Name)
				}
				if tool.Function.Description != tt.mcpTools[i].Description {
					t.Errorf("Tool %d: Expected description %s, got %s", i, tt.mcpTools[i].Description, tool.Function.Description)
				}
			}
		})
	}
}

// TestLLM_ConvertMCPToolsToOpenAI_ObjectProperties tests that object-type properties
// get additionalProperties: true to allow arbitrary key-value pairs
func TestLLM_ConvertMCPToolsToOpenAI_ObjectProperties(t *testing.T) {
	mcpTools := []Tool{
		{
			Name:        "query",
			Description: "Execute a GraphQL query",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {
						Type:        "string",
						Description: "The GraphQL query string",
					},
					"variables": {
						Type:        "object",
						Description: "Optional query variables as key-value pairs",
					},
				},
				Required:             []string{"query"},
				AdditionalProperties: false,
			},
		},
	}

	tools := ConvertMCPToolsToOpenAI(mcpTools)
	if len(tools) != 1 {
		t.Fatalf("Expected 1 tool, got %d", len(tools))
	}

	// Check that parameters don't have additionalProperties at top level
	params := tools[0].Function.Parameters.(map[string]interface{})
	if _, hasAdditionalProps := params["additionalProperties"]; hasAdditionalProps {
		t.Error("Top-level parameters should NOT have additionalProperties field")
	}

	// Check that the variables property has additionalProperties: true
	props := params["properties"].(map[string]interface{})
	variablesProp, ok := props["variables"].(map[string]interface{})
	if !ok {
		t.Fatal("variables property not found")
	}

	if additionalProps, ok := variablesProp["additionalProperties"]; !ok {
		t.Error("Object property 'variables' should have additionalProperties field")
	} else if additionalProps != true {
		t.Errorf("Object property 'variables' additionalProperties should be true, got %v", additionalProps)
	}

	// Check that string property does NOT have additionalProperties
	queryProp, ok := props["query"].(map[string]interface{})
	if !ok {
		t.Fatal("query property not found")
	}

	if _, hasAdditionalProps := queryProp["additionalProperties"]; hasAdditionalProps {
		t.Error("String property 'query' should NOT have additionalProperties field")
	}
}

// TestLLM_ParseToolCall tests parsing tool calls from LLM responses
func TestLLM_ParseToolCall(t *testing.T) {
	tests := []struct {
		name     string
		toolCall openai.ToolCall
		wantName string
		wantArgs map[string]interface{}
		wantErr  bool
	}{
		{
			name: "ValidToolCall",
			toolCall: openai.ToolCall{
				ID: "call-123",
				Function: openai.FunctionCall{
					Name:      "query",
					Arguments: `{"query": "{ users { email } }"}`,
				},
			},
			wantName: "query",
			wantArgs: map[string]interface{}{
				"query": "{ users { email } }",
			},
			wantErr: false,
		},
		{
			name: "EmptyArguments",
			toolCall: openai.ToolCall{
				ID: "call-456",
				Function: openai.FunctionCall{
					Name:      "introspect",
					Arguments: `{}`,
				},
			},
			wantName: "introspect",
			wantArgs: map[string]interface{}{},
			wantErr:  false,
		},
		{
			name: "InvalidJSON",
			toolCall: openai.ToolCall{
				ID: "call-789",
				Function: openai.FunctionCall{
					Name:      "mutate",
					Arguments: `invalid json`,
				},
			},
			wantName: "mutate",
			wantArgs: nil,
			wantErr:  true,
		},
		{
			name: "NestedArguments",
			toolCall: openai.ToolCall{
				ID: "call-abc",
				Function: openai.FunctionCall{
					Name:      "createTask",
					Arguments: `{"title": "Test Task", "priority": "high", "assigneeId": "user-123"}`,
				},
			},
			wantName: "createTask",
			wantArgs: map[string]interface{}{
				"title":      "Test Task",
				"priority":   "high",
				"assigneeId": "user-123",
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			name, args, err := ParseToolCall(tt.toolCall)
			if (err != nil) != tt.wantErr {
				t.Errorf("ParseToolCall() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if name != tt.wantName {
				t.Errorf("ParseToolCall() name = %v, want %v", name, tt.wantName)
			}
			if !tt.wantErr && args != nil {
				// Compare args
				argsJSON, _ := json.Marshal(args)
				wantJSON, _ := json.Marshal(tt.wantArgs)
				if string(argsJSON) != string(wantJSON) {
					t.Errorf("ParseToolCall() args = %v, want %v", args, tt.wantArgs)
				}
			}
		})
	}
}

// TestLLM_ToolConversionSnapshot tests tool conversion with snapshot
func TestLLM_ToolConversionSnapshot(t *testing.T) {
	mcpTools := []Tool{
		{
			Name:        "introspect",
			Description: "Discover the GraphQL schema structure",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"typeName": {
						Type:        "string",
						Description: "Optional type name to introspect",
					},
				},
				Required:             []string{},
				AdditionalProperties: false,
			},
		},
		{
			Name:        "query",
			Description: "Execute a GraphQL query",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {
						Type:        "string",
						Description: "The GraphQL query string",
					},
				},
				Required:             []string{"query"},
				AdditionalProperties: false,
			},
		},
		{
			Name:        "mutate",
			Description: "Execute a GraphQL mutation",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"mutation": {
						Type:        "string",
						Description: "The GraphQL mutation string",
					},
				},
				Required:             []string{"mutation"},
				AdditionalProperties: false,
			},
		},
	}

	openaiTools := ConvertMCPToolsToOpenAI(mcpTools)
	testSnapshotResult(t, "converted_tools", openaiTools)
}

// TestIsRetryableError tests the isRetryableError function
func TestIsRetryableError(t *testing.T) {
	tests := []struct {
		name string
		err  error
		want bool
	}{
		{
			name: "nil error",
			err:  nil,
			want: false,
		},
		{
			name: "context canceled",
			err:  fmt.Errorf("context canceled"),
			want: true,
		},
		{
			name: "context deadline exceeded",
			err:  fmt.Errorf("context deadline exceeded"),
			want: true,
		},
		{
			name: "connection refused",
			err:  fmt.Errorf("dial tcp 127.0.0.1:8080: connect: connection refused"),
			want: true,
		},
		{
			name: "connection reset",
			err:  fmt.Errorf("read tcp 127.0.0.1:8080: read: connection reset by peer"),
			want: true,
		},
		{
			name: "connection closed",
			err:  fmt.Errorf("connection closed"),
			want: true,
		},
		{
			name: "timeout",
			err:  fmt.Errorf("dial tcp 127.0.0.1:8080: i/o timeout"),
			want: true,
		},
		{
			name: "service unavailable",
			err:  fmt.Errorf("service unavailable (HTTP 503)"),
			want: true,
		},
		{
			name: "rate limit",
			err:  fmt.Errorf("rate limit exceeded (HTTP 429)"),
			want: true,
		},
		{
			name: "server error 500",
			err:  fmt.Errorf("internal server error (HTTP 500)"),
			want: true,
		},
		{
			name: "server error 502",
			err:  fmt.Errorf("bad gateway (HTTP 502)"),
			want: true,
		},
		{
			name: "server error 503",
			err:  fmt.Errorf("service unavailable (HTTP 503)"),
			want: true,
		},
		{
			name: "server error 504",
			err:  fmt.Errorf("gateway timeout (HTTP 504)"),
			want: true,
		},
		{
			name: "invalid API key",
			err:  fmt.Errorf("invalid API key (HTTP 401)"),
			want: false,
		},
		{
			name: "bad request",
			err:  fmt.Errorf("bad request (HTTP 400)"),
			want: false,
		},
		{
			name: "generic error",
			err:  fmt.Errorf("something went wrong"),
			want: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := isRetryableError(tt.err)
			if got != tt.want {
				t.Errorf("isRetryableError(%v) = %v, want %v", tt.err, got, tt.want)
			}
		})
	}
}