ZovOS_AI/internal/agent/agent.go

package agent

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/config"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/tools"
)

// StreamCallback is called for each chunk of streamed output.
type StreamCallback func(chunk string, toolCall *ai.ToolCall, toolResult *string, done bool)

// Agent orchestrates the AI reasoning loop.
type Agent struct {
	provider ai.Provider
	tools    *tools.Registry
	cfg      config.AgentConfig
	log      *slog.Logger
}

// New creates a new Agent.
func New(provider ai.Provider, toolsRegistry *tools.Registry, cfg config.AgentConfig, log *slog.Logger) *Agent {
	return &Agent{
		provider: provider,
		tools:    toolsRegistry,
		cfg:      cfg,
		log:      log,
	}
}

// SetProvider updates the active AI provider.
func (a *Agent) SetProvider(p ai.Provider) {
	a.provider = p
}

// Run executes the agent loop: send messages → get response → execute tools → repeat.
func (a *Agent) Run(ctx context.Context, userMessage string, history []ai.Message, callback StreamCallback) ([]ai.Message, error) {
	timeout := time.Duration(a.cfg.TimeoutSeconds) * time.Second
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	messages := make([]ai.Message, 0, len(history)+2)

	// Add system prompt
	if a.cfg.SystemPrompt != "" {
		messages = append(messages, ai.Message{
			Role:    "system",
			Content: a.cfg.SystemPrompt,
		})
	}

	// Add conversation history
	messages = append(messages, history...)

	// Add new user message
	messages = append(messages, ai.Message{
		Role:    "user",
		Content: userMessage,
	})

	availableTools := a.tools.List()

	for iteration := 0; iteration < a.cfg.MaxIterations; iteration++ {
		a.log.Info("agent iteration", "iteration", iteration+1, "messages", len(messages))

		// Try streaming first
		streamCh, err := a.provider.StreamChat(ctx, messages, availableTools)
		if err != nil {
			return messages, fmt.Errorf("AI provider error: %w", err)
		}

		// Collect full response from stream
		var fullContent string
		var allToolCalls []ai.ToolCall

		for chunk := range streamCh {
			if chunk.Error != nil {
				return messages, fmt.Errorf("stream error: %w", chunk.Error)
			}
			if chunk.Content != "" {
				fullContent += chunk.Content
				if callback != nil {
					callback(chunk.Content, nil, nil, false)
				}
			}
			if len(chunk.ToolCalls) > 0 {
				allToolCalls = append(allToolCalls, chunk.ToolCalls...)
			}
		}

		// Add assistant response to messages
		assistantMsg := ai.Message{
			Role:      "assistant",
			Content:   fullContent,
			ToolCalls: allToolCalls,
		}
		messages = append(messages, assistantMsg)

		// If no tool calls, we're done
		if len(allToolCalls) == 0 {
			if callback != nil {
				callback("", nil, nil, true)
			}
			return messages, nil
		}

		// Execute tool calls
		for _, tc := range allToolCalls {
			a.log.Info("executing tool call", "tool", tc.Function.Name, "id", tc.ID)

			if callback != nil {
				callback("", &tc, nil, false)
			}

			result, err := a.tools.Execute(ctx, tc.Function.Name, tc.Function.Arguments)
			if err != nil {
				result = fmt.Sprintf("Tool execution error: %v", err)
			}

			if callback != nil {
				callback("", &tc, &result, false)
			}

			messages = append(messages, ai.Message{
				Role:       "tool",
				Content:    result,
				ToolCallID: tc.ID,
				Name:       tc.Function.Name,
			})
		}
	}

	if callback != nil {
		callback("", nil, nil, true)
	}

	return messages, fmt.Errorf("max iterations (%d) reached", a.cfg.MaxIterations)
}