Cyrene/backend/ai-core/internal/llm/selector.go

package llm

import (
	"context"
	"fmt"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene/ai-core/internal/config"
)

// ModelPurpose identifies the kind of LLM task.
type ModelPurpose string

const (
	PurposeChat             ModelPurpose = "chat"
	PurposeDeepThinking     ModelPurpose = "deep_thinking"
	PurposeIntentAnalysis   ModelPurpose = "intent_analysis"
	PurposeToolCalling      ModelPurpose = "tool_calling"
	PurposeMemoryExtraction ModelPurpose = "memory_extraction"
	PurposeVision           ModelPurpose = "vision"
	PurposeOCR              ModelPurpose = "ocr"
)

// ErrModelNotRequired is returned when an optional model is unavailable.
var ErrModelNotRequired = fmt.Errorf("model not required, caller should degrade gracefully")

// ModelSelector routes requests to the best available LLMProvider based on purpose.
type ModelSelector struct {
	loader    *config.Loader
	envCfg    OpenAIConfig
	mu        sync.RWMutex
	cache     map[string]LLMProvider
	cachedEnv LLMProvider // cached env fallback, created once
}

// NewModelSelector creates a ModelSelector. If loader is nil or has no config,
// all calls fall back to envCfg.
func NewModelSelector(loader *config.Loader, envFallback OpenAIConfig) *ModelSelector {
	return &ModelSelector{
		loader: loader,
		envCfg: envFallback,
		cache:  make(map[string]LLMProvider),
	}
}

// Select returns an LLMProvider for the given purpose. Falls back through the
// routing fallback chain; returns the env provider if nothing matches.
func (s *ModelSelector) Select(ctx context.Context, purpose ModelPurpose) (LLMProvider, error) {
	if s.loader == nil || !s.loader.HasConfig() {
		return s.envProvider(), nil
	}

	cfg := s.loader.GetConfig()
	if cfg == nil {
		return s.envProvider(), nil
	}

	route, ok := cfg.Routing[string(purpose)]
	if !ok || len(route.FallbackChain) == 0 {
		return s.envProvider(), nil
	}

	for _, modelID := range route.FallbackChain {
		provider, err := s.getOrCreateProvider(modelID, cfg)
		if err != nil {
			continue
		}
		return provider, nil
	}

	if route.Required {
		return nil, fmt.Errorf("all models unavailable for purpose %s", purpose)
	}
	return s.envProvider(), nil
}

// DefaultAdapter returns an *Adapter backed by the chat-purpose provider.
// This is the backward-compatible entry point: all existing consumers
// (Orchestrator, Synthesizer, BackgroundThinker, etc.) use this.
func (s *ModelSelector) DefaultAdapter() *Adapter {
	provider, _ := s.Select(context.Background(), PurposeChat)
	return NewAdapter(provider)
}

func (s *ModelSelector) envProvider() LLMProvider {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.cachedEnv == nil {
		s.cachedEnv = NewOpenAIProvider(s.envCfg)
	}
	return s.cachedEnv
}

func (s *ModelSelector) getOrCreateProvider(modelID string, cfg *config.ModelsConfigData) (LLMProvider, error) {
	s.mu.RLock()
	if p, ok := s.cache[modelID]; ok {
		s.mu.RUnlock()
		return p, nil
	}
	s.mu.RUnlock()

	modelCfg, ok := cfg.Models[modelID]
	if !ok {
		return nil, fmt.Errorf("model %s not found", modelID)
	}
	if !modelCfg.Enabled {
		return nil, fmt.Errorf("model %s is disabled", modelID)
	}

	provCfg, ok := cfg.Providers[modelCfg.Provider]
	if !ok {
		return nil, fmt.Errorf("provider %s not found for model %s", modelCfg.Provider, modelID)
	}

	timeout := time.Duration(provCfg.TimeoutSec) * time.Second
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	maxRetries := provCfg.MaxRetries
	if maxRetries <= 0 {
		maxRetries = 3
	}

	provider := NewOpenAIProvider(OpenAIConfig{
		BaseURL:       provCfg.BaseURL,
		APIKey:        provCfg.APIKey,
		Model:         modelCfg.Name,
		FallbackModel: modelCfg.Name,
		MaxRetries:    maxRetries,
		Timeout:       timeout,
	})

	s.mu.Lock()
	s.cache[modelID] = provider
	s.mu.Unlock()

	return provider, nil
}