arp/arp_agent/llm.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"time"

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

// LLM is an OpenAI LLM wrapper with tool-calling support
type LLM struct {
	client      *openai.Client
	model       string
	temperature float32
	maxTokens   int
	// Retry configuration
	maxRetries int
	retryDelay time.Duration
}

// NewLLM creates a new LLM instance
func NewLLM(apiKey, model string, temperature float32, baseURL string, maxTokens int) *LLM {
	return NewLLMWithRetry(apiKey, model, temperature, baseURL, maxTokens, 3, 1*time.Second)
}

// NewLLMWithRetry creates a new LLM instance with custom retry configuration
func NewLLMWithRetry(apiKey, model string, temperature float32, baseURL string, maxTokens int, maxRetries int, retryDelay time.Duration) *LLM {
	config := openai.DefaultConfig(apiKey)
	if baseURL != "" {
		config.BaseURL = baseURL
	}

	return &LLM{
		client:      openai.NewClientWithConfig(config),
		model:       model,
		temperature: temperature,
		maxTokens:   maxTokens,
		maxRetries:  maxRetries,
		retryDelay:  retryDelay,
	}
}

// ChatCompletionRequest is a request for chat completion
type ChatCompletionRequest struct {
	Messages []openai.ChatCompletionMessage
	Tools    []openai.Tool
}

// ChatCompletionResponse is a response from chat completion
type ChatCompletionResponse struct {
	Message openai.ChatCompletionMessage
}

// Chat sends a chat completion request with retry logic
func (l *LLM) Chat(ctx context.Context, messages []openai.ChatCompletionMessage, tools []openai.Tool) (*openai.ChatCompletionMessage, error) {
	var lastErr error
	delay := l.retryDelay

	for attempt := 0; attempt <= l.maxRetries; attempt++ {
		select {
		case <-ctx.Done():
			return nil, fmt.Errorf("context canceled: %w", ctx.Err())
		default:
		}

		req := openai.ChatCompletionRequest{
			Model:       l.model,
			Messages:    messages,
			Temperature: l.temperature,
			MaxTokens:   l.maxTokens,
		}

		if len(tools) > 0 {
			req.Tools = tools
		}

		resp, err := l.client.CreateChatCompletion(ctx, req)
		if err == nil {
			if len(resp.Choices) == 0 {
				return nil, fmt.Errorf("no response choices returned")
			}

			// Log warning if finish reason indicates an issue
			choice := resp.Choices[0]
			if choice.FinishReason == "length" {
				// Model hit token limit - may have incomplete response
				// This is common with reasoning models that need more tokens
				return nil, fmt.Errorf("response truncated: model hit token limit (finish_reason: length). Consider increasing OPENAI_MAX_TOKENS (current: %d). Usage: prompt=%d, completion=%d, total=%d",
					l.maxTokens, resp.Usage.PromptTokens, resp.Usage.CompletionTokens, resp.Usage.TotalTokens)
			}

			return &choice.Message, nil
		}

		lastErr = err

		// Check if this error is retryable
		if !isRetryableError(err) {
			return nil, fmt.Errorf("failed to create chat completion: %w", err)
		}

		// Don't wait after the last attempt
		if attempt < l.maxRetries {
			log.Printf("LLM request failed (attempt %d/%d): %v. Retrying in %v...",
				attempt+1, l.maxRetries, err, delay)
			select {
			case <-ctx.Done():
				return nil, fmt.Errorf("context canceled during retry wait: %w", ctx.Err())
			case <-time.After(delay):
			}
			// Exponential backoff
			delay *= 2
		}
	}

	return nil, fmt.Errorf("failed to create chat completion after %d retries: %w", l.maxRetries+1, lastErr)
}

// isRetryableError checks if an error is transient and worth retrying
func isRetryableError(err error) bool {
	if err == nil {
		return false
	}

	errStr := err.Error()

	// Context cancellation - retry if not explicitly canceled
	if strings.Contains(errStr, "context canceled") || strings.Contains(errStr, "context deadline exceeded") {
		// These can be transient if the context was canceled due to connection issues
		// But we should check if it's a genuine cancellation vs. a timeout
		return true
	}

	// Network-related errors
	retryablePatterns := []string{
		"connection refused",
		"connection reset",
		"connection closed",
		"network is unreachable",
		"no route to host",
		"timeout",
		"i/o timeout",
		"temporary failure",
		"server misbehaving",
		"service unavailable",
		"too many requests",
		"rate limit",
		"429",
		"500",
		"502",
		"503",
		"504",
	}

	lowerErr := strings.ToLower(errStr)
	for _, pattern := range retryablePatterns {
		if strings.Contains(lowerErr, strings.ToLower(pattern)) {
			return true
		}
	}

	return false
}

// ConvertMCPToolsToOpenAI converts MCP tools to OpenAI tool format
func ConvertMCPToolsToOpenAI(mcpTools []Tool) []openai.Tool {
	tools := make([]openai.Tool, len(mcpTools))
	for i, t := range mcpTools {
		// Convert InputSchema to JSON schema format using map[string]interface{}
		props := make(map[string]interface{})
		for name, prop := range t.InputSchema.Properties {
			propMap := map[string]interface{}{
				"type":        prop.Type,
				"description": prop.Description,
			}
			// For object types without explicit nested properties,
			// allow additionalProperties so the LLM can pass any key-value pairs
			// This is important for tools like 'query' and 'mutate' that accept
			// arbitrary variables objects
			if prop.Type == "object" {
				propMap["additionalProperties"] = true
			}
			props[name] = propMap
		}

		// Build parameters map, omitting empty required array
		params := map[string]interface{}{
			"type":       t.InputSchema.Type,
			"properties": props,
		}
		// Only include required if it has elements - empty slice marshals as null
		if len(t.InputSchema.Required) > 0 {
			params["required"] = t.InputSchema.Required
		}

		tools[i] = openai.Tool{
			Type: openai.ToolTypeFunction,
			Function: &openai.FunctionDefinition{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  params,
			},
		}
	}
	return tools
}

// ParseToolCall parses a tool call from the LLM response
func ParseToolCall(toolCall openai.ToolCall) (string, map[string]interface{}, error) {
	name := toolCall.Function.Name

	var args map[string]interface{}
	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
		return name, nil, fmt.Errorf("failed to parse tool arguments: %w", err)
	}

	return name, args, nil
}

// TestConnection tests the connection to OpenAI API
func (l *LLM) TestConnection(ctx context.Context) error {
	// Simple test request - use enough tokens for reasoning models
	// Reasoning models need more tokens for their thinking process
	req := openai.ChatCompletionRequest{
		Model: l.model,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleUser,
				Content: "Hello",
			},
		},
		MaxTokens: 100,
	}

	_, err := l.client.CreateChatCompletion(ctx, req)
	if err != nil {
		return fmt.Errorf("failed to connect to OpenAI API: %w", err)
	}

	return nil
}