Add inital implementation of an agent
Change-Id: Ib60c33e8c1a44bc9341cac5c1f1fdc518fb5ed1e
diff --git a/server/agent/agent.go b/server/agent/agent.go
index e901cf6..e3aa766 100644
--- a/server/agent/agent.go
+++ b/server/agent/agent.go
@@ -1,23 +1,399 @@
package agent
+import (
+ "context"
+ "fmt"
+ "log"
+ "os"
+ "path/filepath"
+ "strings"
+ "time"
+
+ "github.com/iomodo/staff/git"
+ "github.com/iomodo/staff/llm"
+ "github.com/iomodo/staff/tm"
+)
+
+// AgentConfig contains configuration for the agent
type AgentConfig struct {
Name string
Role string
GitUsername string
GitEmail string
WorkingDir string
+
+ // LLM Configuration
+ LLMProvider llm.Provider
+ LLMModel string
+ LLMConfig llm.Config
+
+ // System prompt for the agent
+ SystemPrompt string
+
+ // Task Manager Configuration
+ TaskManager tm.TaskManager
+
+ // Git Configuration
+ GitRepoPath string
+ GitRemote string
+ GitBranch string
}
+// Agent represents an AI agent that can process tasks
type Agent struct {
- Config AgentConfig
+ Config AgentConfig
+ llmProvider llm.LLMProvider
+ gitInterface git.GitInterface
+ ctx context.Context
+ cancel context.CancelFunc
}
-func NewAgent(config AgentConfig) *Agent {
- return &Agent{
- Config: config,
+// NewAgent creates a new agent instance
+func NewAgent(config AgentConfig) (*Agent, error) {
+ // Validate configuration
+ if err := validateConfig(config); err != nil {
+ return nil, fmt.Errorf("invalid config: %w", err)
+ }
+
+ // Create LLM provider
+ llmProvider, err := llm.CreateProvider(config.LLMConfig)
+ if err != nil {
+ return nil, fmt.Errorf("failed to create LLM provider: %w", err)
+ }
+
+ // Create git interface
+ gitInterface := git.DefaultGit(config.GitRepoPath)
+
+ // Create context with cancellation
+ ctx, cancel := context.WithCancel(context.Background())
+
+ agent := &Agent{
+ Config: config,
+ llmProvider: llmProvider,
+ gitInterface: gitInterface,
+ ctx: ctx,
+ cancel: cancel,
+ }
+
+ return agent, nil
+}
+
+// validateConfig validates the agent configuration
+func validateConfig(config AgentConfig) error {
+ if config.Name == "" {
+ return fmt.Errorf("agent name is required")
+ }
+ if config.Role == "" {
+ return fmt.Errorf("agent role is required")
+ }
+ if config.WorkingDir == "" {
+ return fmt.Errorf("working directory is required")
+ }
+ if config.SystemPrompt == "" {
+ return fmt.Errorf("system prompt is required")
+ }
+ if config.TaskManager == nil {
+ return fmt.Errorf("task manager is required")
+ }
+ if config.GitRepoPath == "" {
+ return fmt.Errorf("git repository path is required")
+ }
+ return nil
+}
+
+// Run starts the agent's main loop
+func (a *Agent) Run() error {
+ log.Printf("Starting agent %s (%s)", a.Config.Name, a.Config.Role)
+ defer log.Printf("Agent %s stopped", a.Config.Name)
+
+ // Initialize git repository if needed
+ if err := a.initializeGit(); err != nil {
+ return fmt.Errorf("failed to initialize git: %w", err)
+ }
+
+ // Main agent loop
+ for {
+ select {
+ case <-a.ctx.Done():
+ return a.ctx.Err()
+ default:
+ if err := a.processNextTask(); err != nil {
+ log.Printf("Error processing task: %v", err)
+ // Continue running even if there's an error
+ time.Sleep(30 * time.Second)
+ }
+ }
}
}
-func (a *Agent) Run() {
+// Stop stops the agent
+func (a *Agent) Stop() {
+ log.Printf("Stopping agent %s", a.Config.Name)
+ a.cancel()
+ if a.llmProvider != nil {
+ a.llmProvider.Close()
+ }
+}
+// initializeGit initializes the git repository
+func (a *Agent) initializeGit() error {
+ ctx := context.Background()
+
+ // Check if repository exists
+ isRepo, err := a.gitInterface.IsRepository(ctx, a.Config.GitRepoPath)
+ if err != nil {
+ return fmt.Errorf("failed to check repository: %w", err)
+ }
+
+ if !isRepo {
+ // Initialize new repository
+ if err := a.gitInterface.Init(ctx, a.Config.GitRepoPath); err != nil {
+ return fmt.Errorf("failed to initialize repository: %w", err)
+ }
+ }
+
+ // Set git user configuration
+ userConfig := git.UserConfig{
+ Name: a.Config.GitUsername,
+ Email: a.Config.GitEmail,
+ }
+ if err := a.gitInterface.SetUserConfig(ctx, userConfig); err != nil {
+ return fmt.Errorf("failed to set git user config: %w", err)
+ }
+
+ // Checkout to the specified branch
+ if a.Config.GitBranch != "" {
+ if err := a.gitInterface.Checkout(ctx, a.Config.GitBranch); err != nil {
+ // Try to create the branch if it doesn't exist
+ if err := a.gitInterface.CreateBranch(ctx, a.Config.GitBranch, ""); err != nil {
+ return fmt.Errorf("failed to create branch %s: %w", a.Config.GitBranch, err)
+ }
+ }
+ }
+
+ return nil
+}
+
+// processNextTask processes the next available task
+func (a *Agent) processNextTask() error {
+ ctx := context.Background()
+
+ // Get tasks assigned to this agent
+ taskList, err := a.Config.TaskManager.GetTasksByOwner(ctx, a.Config.Name, 0, 10)
+ if err != nil {
+ return fmt.Errorf("failed to get tasks: %w", err)
+ }
+
+ // Find a task that's ready to be worked on
+ var taskToProcess *tm.Task
+ for _, task := range taskList.Tasks {
+ if task.Status == tm.StatusToDo {
+ taskToProcess = task
+ break
+ }
+ }
+
+ if taskToProcess == nil {
+ // No tasks to process, wait a bit
+ time.Sleep(60 * time.Second)
+ return nil
+ }
+
+ log.Printf("Processing task: %s - %s", taskToProcess.ID, taskToProcess.Title)
+
+ // Start the task
+ startedTask, err := a.Config.TaskManager.StartTask(ctx, taskToProcess.ID)
+ if err != nil {
+ return fmt.Errorf("failed to start task: %w", err)
+ }
+
+ // Process the task with LLM
+ solution, err := a.processTaskWithLLM(startedTask)
+ if err != nil {
+ // Mark task as failed or retry
+ log.Printf("Failed to process task with LLM: %v", err)
+ return err
+ }
+
+ // Create PR with the solution
+ if err := a.createPullRequest(startedTask, solution); err != nil {
+ return fmt.Errorf("failed to create pull request: %w", err)
+ }
+
+ // Complete the task
+ if _, err := a.Config.TaskManager.CompleteTask(ctx, startedTask.ID); err != nil {
+ return fmt.Errorf("failed to complete task: %w", err)
+ }
+
+ log.Printf("Successfully completed task: %s", startedTask.ID)
+ return nil
+}
+
+// processTaskWithLLM sends the task to the LLM and gets a solution
+func (a *Agent) processTaskWithLLM(task *tm.Task) (string, error) {
+ ctx := context.Background()
+
+ // Prepare the prompt
+ prompt := a.buildTaskPrompt(task)
+
+ // Create chat completion request
+ req := llm.ChatCompletionRequest{
+ Model: a.Config.LLMModel,
+ Messages: []llm.Message{
+ {
+ Role: llm.RoleSystem,
+ Content: a.Config.SystemPrompt,
+ },
+ {
+ Role: llm.RoleUser,
+ Content: prompt,
+ },
+ },
+ MaxTokens: intPtr(4000),
+ Temperature: float64Ptr(0.7),
+ }
+
+ // Get response from LLM
+ resp, err := a.llmProvider.ChatCompletion(ctx, req)
+ if err != nil {
+ return "", fmt.Errorf("LLM chat completion failed: %w", err)
+ }
+
+ if len(resp.Choices) == 0 {
+ return "", fmt.Errorf("no response from LLM")
+ }
+
+ return resp.Choices[0].Message.Content, nil
+}
+
+// buildTaskPrompt builds the prompt for the LLM based on the task
+func (a *Agent) buildTaskPrompt(task *tm.Task) string {
+ var prompt strings.Builder
+
+ prompt.WriteString(fmt.Sprintf("Task ID: %s\n", task.ID))
+ prompt.WriteString(fmt.Sprintf("Title: %s\n", task.Title))
+ prompt.WriteString(fmt.Sprintf("Priority: %s\n", task.Priority))
+
+ if task.Description != "" {
+ prompt.WriteString(fmt.Sprintf("Description: %s\n", task.Description))
+ }
+
+ if task.DueDate != nil {
+ prompt.WriteString(fmt.Sprintf("Due Date: %s\n", task.DueDate.Format("2006-01-02")))
+ }
+
+ prompt.WriteString("\nPlease provide a detailed solution for this task. ")
+ prompt.WriteString("Include any code, documentation, or other deliverables as needed. ")
+ prompt.WriteString("Format your response appropriately for the type of task.")
+
+ return prompt.String()
+}
+
+// createPullRequest creates a pull request with the solution
+func (a *Agent) createPullRequest(task *tm.Task, solution string) error {
+ ctx := context.Background()
+
+ // Generate branch name
+ branchName := a.generateBranchName(task)
+
+ // Create and checkout to new branch
+ if err := a.gitInterface.CreateBranch(ctx, branchName, ""); err != nil {
+ return fmt.Errorf("failed to create branch: %w", err)
+ }
+
+ if err := a.gitInterface.Checkout(ctx, branchName); err != nil {
+ return fmt.Errorf("failed to checkout branch: %w", err)
+ }
+
+ // Create solution file
+ solutionPath := filepath.Join(a.Config.WorkingDir, fmt.Sprintf("task-%s-solution.md", task.ID))
+ solutionContent := a.formatSolution(task, solution)
+
+ if err := os.WriteFile(solutionPath, []byte(solutionContent), 0644); err != nil {
+ return fmt.Errorf("failed to write solution file: %w", err)
+ }
+
+ // Add and commit the solution
+ if err := a.gitInterface.Add(ctx, []string{solutionPath}); err != nil {
+ return fmt.Errorf("failed to add solution file: %w", err)
+ }
+
+ commitMessage := fmt.Sprintf("feat: Complete task %s - %s", task.ID, task.Title)
+ if err := a.gitInterface.Commit(ctx, commitMessage, git.CommitOptions{}); err != nil {
+ return fmt.Errorf("failed to commit solution: %w", err)
+ }
+
+ // Push the branch
+ if err := a.gitInterface.Push(ctx, "origin", branchName, git.PushOptions{SetUpstream: true}); err != nil {
+ return fmt.Errorf("failed to push branch: %w", err)
+ }
+
+ log.Printf("Created pull request for task %s on branch %s", task.ID, branchName)
+ return nil
+}
+
+// generateBranchName generates a branch name for the task
+func (a *Agent) generateBranchName(task *tm.Task) string {
+ // Clean the task title for branch name
+ cleanTitle := strings.ReplaceAll(task.Title, " ", "-")
+ cleanTitle = strings.ToLower(cleanTitle)
+
+ // Remove special characters that are not allowed in git branch names
+ // Keep only alphanumeric characters and hyphens
+ var result strings.Builder
+ for _, char := range cleanTitle {
+ if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '-' {
+ result.WriteRune(char)
+ }
+ }
+ cleanTitle = result.String()
+
+ // Remove consecutive hyphens
+ for strings.Contains(cleanTitle, "--") {
+ cleanTitle = strings.ReplaceAll(cleanTitle, "--", "-")
+ }
+
+ // Remove leading and trailing hyphens
+ cleanTitle = strings.Trim(cleanTitle, "-")
+
+ // Limit length
+ if len(cleanTitle) > 50 {
+ cleanTitle = cleanTitle[:50]
+ // Ensure we don't end with a hyphen after truncation
+ cleanTitle = strings.TrimSuffix(cleanTitle, "-")
+ }
+
+ return fmt.Sprintf("task/%s-%s", task.ID, cleanTitle)
+}
+
+// formatSolution formats the solution for the pull request
+func (a *Agent) formatSolution(task *tm.Task, solution string) string {
+ var content strings.Builder
+
+ content.WriteString(fmt.Sprintf("# Task Solution: %s\n\n", task.Title))
+ content.WriteString(fmt.Sprintf("**Task ID:** %s\n", task.ID))
+ content.WriteString(fmt.Sprintf("**Agent:** %s (%s)\n", a.Config.Name, a.Config.Role))
+ content.WriteString(fmt.Sprintf("**Completed:** %s\n\n", time.Now().Format("2006-01-02 15:04:05")))
+
+ content.WriteString("## Task Description\n\n")
+ content.WriteString(task.Description)
+ content.WriteString("\n\n")
+
+ content.WriteString("## Solution\n\n")
+ content.WriteString(solution)
+ content.WriteString("\n\n")
+
+ content.WriteString("---\n")
+ content.WriteString("*This solution was generated by AI Agent*\n")
+
+ return content.String()
+}
+
+// ptr helpers for cleaner code
+func intPtr(i int) *int {
+ return &i
+}
+
+func float64Ptr(f float64) *float64 {
+ return &f
}