gochat/backend/internal/service/tool_execution_service.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// ToolExecutionService handles LLM function calling: converts CaptainCustomTool
// definitions to LLM ToolDefinitions, executes HTTP tool calls, and runs the
// tool_call loop (LLM → tool_call → execute → result → LLM → final answer).
//
// Reference: AI_FEATURE_ROADMAP.md §3.2 — Function Calling complete implementation
type ToolExecutionService struct {
	toolRepo    *repository.CaptainCustomToolRepo
	llmProvider llm.Provider
	httpClient  *http.Client
}

// NewToolExecutionService creates a new ToolExecutionService.
func NewToolExecutionService(toolRepo *repository.CaptainCustomToolRepo, llmProvider llm.Provider) *ToolExecutionService {
	return &ToolExecutionService{
		toolRepo:    toolRepo,
		llmProvider: llmProvider,
		httpClient:  &http.Client{Timeout: 30 * time.Second},
	}
}

// GetToolsForAssistant returns enabled custom tools for an account as LLM ToolDefinitions.
func (s *ToolExecutionService) GetToolsForAccount(ctx context.Context, accountID uint) ([]llm.ToolDefinition, error) {
	// ListByAccount returns all tools; we filter for enabled ones
	tools, _, err := s.toolRepo.ListByAccount(ctx, accountID, 0, 100)
	if err != nil {
		return nil, fmt.Errorf("fetch custom tools: %w", err)
	}

	defs := make([]llm.ToolDefinition, 0, len(tools))
	for _, tool := range tools {
		if !tool.Enabled {
			continue
		}
		def := customToolToDefinition(tool)
		defs = append(defs, def)
	}
	return defs, nil
}

// customToolToDefinition converts a CaptainCustomTool to an LLM ToolDefinition.
func customToolToDefinition(tool model.CaptainCustomTool) llm.ToolDefinition {
	var params map[string]interface{}
	if len(tool.ParamSchema) > 0 && string(tool.ParamSchema) != "null" {
		_ = json.Unmarshal(tool.ParamSchema, &params)
	}
	if params == nil {
		params = map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}
	}

	return llm.ToolDefinition{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        tool.Slug,
			Description: tool.Description,
			Parameters:  params,
		},
	}
}

// ExecuteToolCall executes a single tool call by making the configured HTTP request.
// Returns the result as a string (typically JSON).
func (s *ToolExecutionService) ExecuteToolCall(ctx context.Context, accountID uint, call llm.ToolCall) (string, error) {
	// Find the tool by slug (function name)
	tools, _, err := s.toolRepo.ListByAccount(ctx, accountID, 0, 100)
	if err != nil {
		return "", fmt.Errorf("fetch tools: %w", err)
	}

	var tool *model.CaptainCustomTool
	for i := range tools {
		if tools[i].Enabled && tools[i].Slug == call.Function.Name {
			tool = &tools[i]
			break
		}
	}
	if tool == nil {
		return "", fmt.Errorf("tool %s not found or not enabled", call.Function.Name)
	}

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

	// Build HTTP request
	method := strings.ToUpper(tool.HTTPMethod)
	if method == "" {
		method = "GET"
	}

	var bodyReader io.Reader
	if method == "POST" || method == "PUT" || method == "PATCH" {
		bodyBytes, _ := json.Marshal(args)
		bodyReader = bytes.NewReader(bodyBytes)
	}

	req, err := http.NewRequestWithContext(ctx, method, tool.EndpointURL, bodyReader)
	if err != nil {
		return "", fmt.Errorf("create tool request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	// Apply auth
	applyToolAuth(req, tool)

	// Execute
	resp, err := s.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("execute tool request: %w", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(io.LimitReader(resp.Body, 10*1024)) // max 10KB response

	if resp.StatusCode >= 400 {
		return "", fmt.Errorf("tool endpoint returned status %d: %s", resp.StatusCode, string(body))
	}

	applogger.L().Infof("ToolExecutionService: tool %s returned status %d, body=%s",
		call.Function.Name, resp.StatusCode, string(body))

	return string(body), nil
}

// applyToolAuth applies authentication to the HTTP request based on the tool's auth config.
func applyToolAuth(req *http.Request, tool *model.CaptainCustomTool) {
	switch tool.AuthType {
	case model.ToolAuthTypeBearer:
		var authCfg struct {
			Token string `json:"token"`
		}
		_ = json.Unmarshal(tool.AuthConfig, &authCfg)
		if authCfg.Token != "" {
			req.Header.Set("Authorization", "Bearer "+authCfg.Token)
		}
	case model.ToolAuthTypeBasic:
		var authCfg struct {
			Username string `json:"username"`
			Password string `json:"password"`
		}
		_ = json.Unmarshal(tool.AuthConfig, &authCfg)
		if authCfg.Username != "" {
			req.SetBasicAuth(authCfg.Username, authCfg.Password)
		}
	case model.ToolAuthTypeApiKey:
		var authCfg struct {
			Header string `json:"header"`
			Key    string `json:"key"`
		}
		_ = json.Unmarshal(tool.AuthConfig, &authCfg)
		if authCfg.Header != "" && authCfg.Key != "" {
			req.Header.Set(authCfg.Header, authCfg.Key)
		}
	}
}

// RunToolCallLoop executes the full LLM tool_call loop:
// 1. Send messages + tools to LLM
// 2. If LLM returns tool_calls, execute each and add results as tool messages
// 3. Re-send to LLM with tool results
// 4. Repeat until LLM returns a normal content response (no tool_calls)
// 5. Return final content
//
// maxIterations prevents infinite loops (default 5).
func (s *ToolExecutionService) RunToolCallLoop(
	ctx context.Context,
	accountID uint,
	messages []llm.ChatMessage,
	modelName string,
	temperature float64,
	maxTokens int,
	maxIterations int,
) (string, error) {
	if s.llmProvider == nil {
		return "", fmt.Errorf("LLM provider not configured")
	}
	if maxIterations <= 0 {
		maxIterations = 5
	}

	// Get tools for this account
	tools, err := s.GetToolsForAccount(ctx, accountID)
	if err != nil {
		applogger.L().Warnf("ToolExecutionService: failed to get tools: %v, continuing without tools", err)
		tools = nil
	}

	for iteration := 0; iteration < maxIterations; iteration++ {
		req := llm.ChatRequest{
			Model:       modelName,
			Messages:    messages,
			Temperature: temperature,
			MaxTokens:   maxTokens,
		}
		if len(tools) > 0 {
			req.Tools = tools
		}

		resp, err := s.llmProvider.ChatCompletion(ctx, req)
		if err != nil {
			return "", fmt.Errorf("LLM call iteration %d: %w", iteration, err)
		}
		if resp == nil || len(resp.Choices) == 0 {
			return "", fmt.Errorf("empty LLM response at iteration %d", iteration)
		}

		choice := resp.Choices[0]

		// If finish_reason is "tool_calls" or message has tool_calls, execute them
		if len(choice.Message.ToolCalls) > 0 {
			// Add assistant message with tool_calls to conversation
			messages = append(messages, llm.ChatMessage{
				Role:      "assistant",
				ToolCalls: choice.Message.ToolCalls,
				Content:   choice.Message.Content,
			})

			// Execute each tool call and add results
			for _, call := range choice.Message.ToolCalls {
				result, execErr := s.ExecuteToolCall(ctx, accountID, call)
				if execErr != nil {
					applogger.L().Errorf("ToolExecutionService: tool %s failed: %v", call.Function.Name, execErr)
					result = fmt.Sprintf(`{"error": "%s"}`, escapeJSONString(execErr.Error()))
				}

				messages = append(messages, llm.ChatMessage{
					Role:       "tool",
					Content:    result,
					ToolCallID: call.ID,
					Name:       call.Function.Name,
				})
			}
			continue
		}

		// No tool_calls — this is the final answer
		return choice.Message.Content, nil
	}

	return "", fmt.Errorf("tool call loop exceeded %d iterations", maxIterations)
}

func escapeJSONString(s string) string {
	var b strings.Builder
	for _, r := range s {
		switch r {
		case '"':
			b.WriteString(`\"`)
		case '\\':
			b.WriteString(`\\`)
		case '\n':
			b.WriteString(`\n`)
		case '\t':
			b.WriteString(`\t`)
		case '\r':
			b.WriteString(`\r`)
		default:
			b.WriteRune(r)
		}
	}
	return b.String()
}