server/agent/agent.go

package agent

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

	"github.com/iomodo/staff/config"
	"github.com/iomodo/staff/llm"
	"github.com/iomodo/staff/llm/provider"
	"github.com/iomodo/staff/tm"
)

type Agent struct {
	// Identity
	Name string
	Role string

	// LLM Configuration
	Model        string
	SystemPrompt string
	MaxTokens    *int
	Temperature  *float64

	// Runtime
	Provider    llm.LLMProvider
	CurrentTask *string // Task ID currently being processed

	IsRunning bool
	StopChan  chan struct{}

	logger *slog.Logger

	taskManager tm.TaskManager
	thinker     *Thinker
}

func NewAgent(agentConfig config.AgentConfig, llmConfig llm.Config, taskManager tm.TaskManager, agentRoles []string, logger *slog.Logger) (*Agent, error) {
	// Load system prompt
	systemPrompt, err := loadSystemPrompt(agentConfig.SystemPromptFile)
	if err != nil {
		return nil, fmt.Errorf("failed to load system prompt: %w", err)
	}

	prov := provider.CreateProvider(llmConfig)
	thinker := NewThinker(prov, agentConfig.Model, systemPrompt, *agentConfig.MaxTokens, *agentConfig.Temperature, agentRoles, logger)

	agent := &Agent{
		Name:         agentConfig.Name,
		Role:         agentConfig.Role,
		Model:        agentConfig.Model,
		SystemPrompt: systemPrompt,
		Provider:     prov,
		MaxTokens:    agentConfig.MaxTokens,
		Temperature:  agentConfig.Temperature,
		taskManager:  taskManager,
		logger:       logger,
		thinker:      thinker,
	}

	return agent, nil
}

// Start starts an agent to process tasks in a loop
func (a *Agent) Start(loopInterval time.Duration) error {
	if a.IsRunning {
		return fmt.Errorf("agent %s is already running", a.Name)
	}

	a.IsRunning = true
	a.StopChan = make(chan struct{})

	go a.runLoop(loopInterval)

	a.logger.Info("Started agent",
		slog.String("name", a.Name),
		slog.String("role", a.Role),
		slog.String("model", a.Model))
	return nil
}

func (a *Agent) Stop() {
	close(a.StopChan)
	a.IsRunning = false
}

func (a *Agent) runLoop(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-a.StopChan:
			a.logger.Info("Agent stopping", slog.String("name", a.Name))
			return
		case <-ticker.C:
			if err := a.processTasks(); err != nil {
				a.logger.Error("Error processing tasks for agent",
					slog.String("agent", a.Name),
					slog.String("error", err.Error()))
			}
		}
	}
}

// processAgentTasks processes all assigned tasks for an agent
func (a *Agent) processTasks() error {
	if a.CurrentTask != nil {
		return nil
	}

	// Get tasks assigned to this agent
	tasks, err := a.taskManager.GetTasksByAssignee(a.Name)
	if err != nil {
		return fmt.Errorf("failed to get tasks for agent %s: %w", a.Name, err)
	}

	a.logger.Info("Processing tasks for agent",
		slog.Int("task_count", len(tasks)),
		slog.String("agent", a.Name))

	for _, task := range tasks {
		if task.Status == tm.StatusToDo {
			err2 := a.processTask(task)
			if err2 == nil {
				return nil
			}
			a.logger.Error("Error processing task",
				slog.String("task_id", task.ID),
				slog.String("error", err2.Error()))

		}
	}

	return nil
}

// processTask processes a single task with an agent
func (a *Agent) processTask(task *tm.Task) error {
	ctx := context.Background()
	startTime := time.Now()

	a.logger.Info("Agent processing task",
		slog.String("agent", a.Name),
		slog.String("task_id", task.ID),
		slog.String("title", task.Title))

	// Mark task as in progress
	task.Status = tm.StatusInProgress
	a.CurrentTask = &task.ID
	defer func() {
		if r := recover(); r != nil {
			a.logger.Error("Task processing panicked, clearing agent state",
				slog.String("task_id", task.ID),
				slog.String("agent", a.Name))
			a.CurrentTask = nil
			panic(r)
		}
	}()

	// Check if this task should generate subtasks (with LLM decision)
	if a.thinker.ShouldGenerateSubtasks(task) {
		err := a.processSubtask(ctx, task)
		if err == nil {
			a.logger.Info("Task converted to subtasks by agent using LLM analysis",
				slog.String("task_id", task.ID),
				slog.String("agent", a.Name))
			return nil
		}
		a.logger.Error("Subtask processing failed, cleared agent state",
			slog.String("task_id", task.ID),
			slog.String("agent", a.Name),
			slog.String("error", err.Error()))
	}

	err := a.processSolution(ctx, task)
	if err != nil {
		return fmt.Errorf("failed to process solution for task: %w", err)
	}
	duration := time.Since(startTime)
	a.logger.Info("Task completed by agent",
		slog.String("task_id", task.ID),
		slog.String("agent", a.Name),
		slog.Duration("duration", duration))
	return nil
}

func (a *Agent) processSubtask(ctx context.Context, task *tm.Task) error {
	a.logger.Info("LLM determined task should generate subtasks", slog.String("task_id", task.ID))
	analysis, err := a.thinker.GenerateSubtasksForTask(ctx, task)
	if err != nil {
		return fmt.Errorf("failed to generate subtasks for task: %w", err)
	}

	solutionURL, err2 := a.taskManager.ProposeSubTasks(ctx, task, analysis, a.Name)
	if err2 != nil {
		return fmt.Errorf("failed to propose subtasks for task: %w", err2)
	}
	task.SolutionURL = solutionURL

	a.logger.Info("Generated subtask Solution for task",
		slog.String("task_id", task.ID),
		slog.String("solution_url", solutionURL))
	a.logger.Info("Proposed subtasks and new agents for task",
		slog.String("task_id", task.ID),
		slog.Int("subtask_count", len(analysis.Subtasks)),
		slog.Int("new_agent_count", len(analysis.AgentCreations)))

	// Log proposed new agents if any
	if len(analysis.AgentCreations) > 0 {
		for _, agent := range analysis.AgentCreations {
			a.logger.Info("Proposed new agent",
				slog.String("role", agent.Role),
				slog.Any("skills", agent.Skills))
		}
	}

	return nil
}

func (a *Agent) processSolution(ctx context.Context, task *tm.Task) error {
	solution, err := a.thinker.GenerateSolution(ctx, task)
	if err != nil {
		return fmt.Errorf("failed to generate solution: %w", err)
	}

	solutionURL, err := a.taskManager.ProposeSolution(ctx, task, solution, a.Name)
	if err != nil {
		return fmt.Errorf("failed to propose solution: %w", err)
	}
	task.SolutionURL = solutionURL

	a.logger.Info("Generated Solution for task",
		slog.String("task_id", task.ID),
		slog.String("agent", a.Name),
		slog.String("solution_url", solutionURL))
	return nil
}

// loadSystemPrompt loads the system prompt from file
func loadSystemPrompt(filePath string) (string, error) {
	content, err := os.ReadFile(filePath)
	if err != nil {
		return "", fmt.Errorf("failed to read system prompt file %s: %w", filePath, err)
	}
	return string(content), nil
}