GoClaw/gateway/internal/orchestrator/orchestrator.go

// Package orchestrator implements the GoClaw main AI orchestration loop.
// It loads config from DB, calls the LLM with tool definitions,
// executes tool calls, and returns the final response.
package orchestrator

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

	"git.softuniq.eu/UniqAI/GoClaw/gateway/internal/db"
	"git.softuniq.eu/UniqAI/GoClaw/gateway/internal/llm"
	"git.softuniq.eu/UniqAI/GoClaw/gateway/internal/tools"
)

// ─── Types ────────────────────────────────────────────────────────────────────

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ToolCallStep struct {
	Tool       string `json:"tool"`
	Args       any    `json:"args"`
	Result     any    `json:"result,omitempty"`
	Error      string `json:"error,omitempty"`
	Success    bool   `json:"success"`
	DurationMs int64  `json:"durationMs"`
}

type ChatResult struct {
	Success   bool           `json:"success"`
	Response  string         `json:"response"`
	ToolCalls []ToolCallStep `json:"toolCalls"`
	Model     string         `json:"model"`
	Usage     *llm.Usage     `json:"usage,omitempty"`
	Error     string         `json:"error,omitempty"`
}

// OrchestratorConfig is the runtime config loaded from DB or defaults.
type OrchestratorConfig struct {
	ID           int
	Name         string
	Model        string
	SystemPrompt string
	AllowedTools []string
	Temperature  float64
	MaxTokens    int
}

// ─── Default System Prompt ────────────────────────────────────────────────────

const defaultSystemPrompt = `You are GoClaw Orchestrator — the main AI agent managing the GoClaw distributed AI system.

You have full access to:
1. **Specialized Agents**: Browser Agent (web browsing), Tool Builder (create tools), Agent Compiler (create agents)
2. **System Tools**: shell_exec (run commands), file_read/write (manage files), http_request (web requests), docker_exec (Docker management)

Your responsibilities:
- Answer user questions directly when possible
- Delegate complex web tasks to Browser Agent
- Execute shell commands to manage the system
- Read and write files to modify the codebase
- Monitor Docker containers and services

Decision making:
- For simple questions: answer directly without tools
- For system tasks: use shell_exec, file_read/write
- For Docker: use docker_exec
- Always use list_agents first if you're unsure which agent to delegate to

Response style:
- Be concise and actionable
- Show what tools you used and their results
- Respond in the same language as the user

You are running on a Linux server with Docker and full internet access.`

// ─── Orchestrator ─────────────────────────────────────────────────────────────

type Orchestrator struct {
	llmClient   *llm.Client
	executor    *tools.Executor
	database    *db.DB
	projectRoot string
}

func New(llmClient *llm.Client, database *db.DB, projectRoot string) *Orchestrator {
	o := &Orchestrator{
		llmClient:   llmClient,
		database:    database,
		projectRoot: projectRoot,
	}
	// Inject agent list function to avoid circular dependency
	o.executor = tools.NewExecutor(projectRoot, o.listAgentsFn)
	// Inject DB so delegate_to_agent can resolve live agent container addresses
	o.executor.SetDatabase(database)
	return o
}

// GetConfig loads orchestrator config from DB, falls back to defaults.
func (o *Orchestrator) GetConfig() *OrchestratorConfig {
	if o.database != nil {
		cfg, err := o.database.GetOrchestratorConfig()
		if err == nil && cfg != nil {
			systemPrompt := cfg.SystemPrompt
			if systemPrompt == "" {
				systemPrompt = defaultSystemPrompt
			}
			return &OrchestratorConfig{
				ID:           cfg.ID,
				Name:         cfg.Name,
				Model:        cfg.Model,
				SystemPrompt: systemPrompt,
				AllowedTools: cfg.AllowedTools,
				Temperature:  cfg.Temperature,
				MaxTokens:    cfg.MaxTokens,
			}
		}
		log.Printf("[Orchestrator] Failed to load config from DB: %v — using defaults", err)
	}
	return &OrchestratorConfig{
		Name:         "GoClaw Orchestrator",
		Model:        "qwen2.5:7b",
		SystemPrompt: defaultSystemPrompt,
		Temperature:  0.5,
		MaxTokens:    8192,
	}
}

// Chat runs the full orchestration loop: LLM → tool calls → LLM → response.
func (o *Orchestrator) Chat(ctx context.Context, messages []Message, overrideModel string, maxIter int) ChatResult {
	if maxIter <= 0 {
		maxIter = 10
	}

	cfg := o.GetConfig()
	model := cfg.Model
	if overrideModel != "" {
		model = overrideModel
	}

	log.Printf("[Orchestrator] Chat started: model=%s, messages=%d", model, len(messages))

	// Build conversation
	conv := []llm.Message{
		{Role: "system", Content: cfg.SystemPrompt},
	}
	for _, m := range messages {
		conv = append(conv, llm.Message{Role: m.Role, Content: m.Content})
	}

	// Build tools list
	toolDefs := tools.OrchestratorTools()
	llmTools := make([]llm.Tool, len(toolDefs))
	for i, t := range toolDefs {
		llmTools[i] = llm.Tool{
			Type: t.Type,
			Function: llm.ToolFunction{
				Name:        t.Function.Name,
				Description: t.Function.Description,
				Parameters:  t.Function.Parameters,
			},
		}
	}

	temp := cfg.Temperature
	maxTok := cfg.MaxTokens

	var toolCallSteps []ToolCallStep
	var finalResponse string
	var lastUsage *llm.Usage
	var lastModel string

	for iter := 0; iter < maxIter; iter++ {
		req := llm.ChatRequest{
			Model:       model,
			Messages:    conv,
			Temperature: &temp,
			MaxTokens:   &maxTok,
			Tools:       llmTools,
			ToolChoice:  "auto",
		}

		resp, err := o.llmClient.Chat(ctx, req)
		if err != nil {
			// Fallback: try without tools
			log.Printf("[Orchestrator] LLM error with tools: %v — retrying without tools", err)
			req.Tools = nil
			req.ToolChoice = ""
			resp2, err2 := o.llmClient.Chat(ctx, req)
			if err2 != nil {
				return ChatResult{
					Success: false,
					Error:   fmt.Sprintf("LLM error (model: %s): %v", model, err2),
				}
			}
			if len(resp2.Choices) > 0 {
				finalResponse = resp2.Choices[0].Message.Content
				lastUsage = resp2.Usage
				lastModel = resp2.Model
			}
			break
		}

		if len(resp.Choices) == 0 {
			break
		}

		choice := resp.Choices[0]
		lastUsage = resp.Usage
		lastModel = resp.Model
		if lastModel == "" {
			lastModel = model
		}

		// Check if LLM wants to call tools
		if choice.FinishReason == "tool_calls" && len(choice.Message.ToolCalls) > 0 {
			// Add assistant message with tool calls to conversation
			conv = append(conv, choice.Message)

			// Execute each tool call
			for _, tc := range choice.Message.ToolCalls {
				toolName := tc.Function.Name
				argsJSON := tc.Function.Arguments

				log.Printf("[Orchestrator] Executing tool: %s args=%s", toolName, argsJSON)
				start := time.Now()

				result := o.executor.Execute(ctx, toolName, argsJSON)

				step := ToolCallStep{
					Tool:       toolName,
					Success:    result.Success,
					DurationMs: time.Since(start).Milliseconds(),
				}

				// Parse args for display
				var argsMap any
				_ = json.Unmarshal([]byte(argsJSON), &argsMap)
				step.Args = argsMap

				var toolResultContent string
				if result.Success {
					step.Result = result.Result
					resultBytes, _ := json.Marshal(result.Result)
					toolResultContent = string(resultBytes)
				} else {
					step.Error = result.Error
					toolResultContent = fmt.Sprintf(`{"error": %q}`, result.Error)
				}

				toolCallSteps = append(toolCallSteps, step)

				// Add tool result to conversation
				conv = append(conv, llm.Message{
					Role:       "tool",
					Content:    toolResultContent,
					ToolCallID: tc.ID,
					Name:       toolName,
				})
			}
			// Continue loop — LLM will process tool results
			continue
		}

		// LLM finished — extract final response
		finalResponse = choice.Message.Content
		break
	}

	return ChatResult{
		Success:   true,
		Response:  finalResponse,
		ToolCalls: toolCallSteps,
		Model:     lastModel,
		Usage:     lastUsage,
	}
}

// listAgentsFn is injected into the tool executor to list agents from DB.
func (o *Orchestrator) listAgentsFn() ([]map[string]any, error) {
	if o.database == nil {
		return []map[string]any{}, nil
	}
	rows, err := o.database.ListAgents()
	if err != nil {
		return nil, err
	}
	result := make([]map[string]any, len(rows))
	for i, r := range rows {
		result[i] = map[string]any{
			"id":             r.ID,
			"name":           r.Name,
			"role":           r.Role,
			"model":          r.Model,
			"description":    r.Description,
			"isActive":       r.IsActive,
			"isSystem":       r.IsSystem,
			"isOrchestrator": r.IsOrchestrator,
		}
	}
	return result, nil
}