unibots/shell/orchestration/router.go

package orchestration

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	coretypes "github.com/enmanuel/agents/pkg/llm"
	"github.com/enmanuel/agents/pkg/orchestration"
)

// routeInitial asks the LLM which bot should handle the question first.
func (o *Orchestrator) routeInitial(ctx context.Context, question string, participants []string) (orchestration.RoutingDecision, error) {
	systemPrompt := strings.ReplaceAll(o.routingPrompt, "{{PARTICIPANTS}}", o.buildParticipantsList(participants, ""))

	resp, err := o.llm(ctx, coretypes.CompletionRequest{
		Model:        o.cfg.LLM.Primary.Model,
		MaxTokens:    o.cfg.LLM.Primary.MaxTokens,
		Temperature:  o.cfg.LLM.Primary.Temperature,
		SystemPrompt: systemPrompt,
		Messages: []coretypes.Message{
			{Role: coretypes.RoleUser, Content: question},
		},
	})
	if err != nil {
		return orchestration.RoutingDecision{}, fmt.Errorf("LLM routing call: %w", err)
	}

	var rd orchestration.RoutingDecision
	if err := json.Unmarshal([]byte(strings.TrimSpace(resp.Content)), &rd); err != nil {
		o.logger.Warn("failed to parse routing response, raw", "content", resp.Content, "err", err)
		return orchestration.RoutingDecision{}, fmt.Errorf("parse routing decision: %w", err)
	}

	// Validate the chosen bot is actually a participant
	if !contains(participants, rd.TargetBotID) {
		o.logger.Warn("LLM chose unknown bot, falling back to first", "chosen", rd.TargetBotID)
		rd.TargetBotID = participants[0]
		rd.Confidence = 0.5
		rd.Reason = "fallback: LLM chose unknown bot"
	}

	return rd, nil
}

// routeRefinement asks the LLM which bot should improve the response,
// excluding the last respondent.
func (o *Orchestrator) routeRefinement(
	ctx context.Context,
	question string,
	responses []orchestration.BotResponse,
	participants []string,
	excludeBot string,
) (orchestration.RoutingDecision, error) {
	lastResponse := ""
	if len(responses) > 0 {
		lastResponse = responses[len(responses)-1].Text
	}

	systemPrompt := strings.ReplaceAll(o.refinementPrompt, "{{PARTICIPANTS}}", o.buildParticipantsList(participants, excludeBot))
	systemPrompt = strings.ReplaceAll(systemPrompt, "{{LAST_RESPONSE}}", lastResponse)

	userContent := fmt.Sprintf("Original question: %s\n\nCurrent response that needs improvement:\n%s", question, lastResponse)

	resp, err := o.llm(ctx, coretypes.CompletionRequest{
		Model:        o.cfg.LLM.Primary.Model,
		MaxTokens:    o.cfg.LLM.Primary.MaxTokens,
		Temperature:  o.cfg.LLM.Primary.Temperature,
		SystemPrompt: systemPrompt,
		Messages: []coretypes.Message{
			{Role: coretypes.RoleUser, Content: userContent},
		},
	})
	if err != nil {
		return orchestration.RoutingDecision{}, fmt.Errorf("LLM refinement call: %w", err)
	}

	var rd orchestration.RoutingDecision
	if err := json.Unmarshal([]byte(strings.TrimSpace(resp.Content)), &rd); err != nil {
		o.logger.Warn("failed to parse refinement response", "content", resp.Content, "err", err)
		return orchestration.RoutingDecision{}, fmt.Errorf("parse refinement decision: %w", err)
	}

	// Validate: must be a participant and not the excluded bot
	if rd.TargetBotID == excludeBot || !contains(participants, rd.TargetBotID) {
		// Pick first available that isn't excluded
		for _, p := range participants {
			if p != excludeBot {
				rd.TargetBotID = p
				rd.Reason = "fallback: LLM chose excluded or unknown bot"
				break
			}
		}
	}

	return rd, nil
}

func contains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}