gochat/backend/internal/llm/provider_manager.go

package llm

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
)

var ErrProviderNotConfigured = errors.New("Copilot provider is not configured")

const (
	EmbeddingModeReuseChat = "reuse_chat_credentials"
	EmbeddingModeSeparate  = "separate"
)

// RuntimeProviderConfig is the database-backed Copilot provider configuration.
// It intentionally contains no enabled flag: Copilot is globally available and
// reports ErrProviderNotConfigured until an administrator saves credentials.
type RuntimeProviderConfig struct {
	ChatProvider string
	ChatBaseURL  string
	ChatAPIKey   string
	ChatModel    string

	EmbeddingMode       string
	EmbeddingProvider   string
	EmbeddingBaseURL    string
	EmbeddingAPIKey     string
	EmbeddingModel      string
	EmbeddingDimensions int

	Temperature    float64
	MaxTokens      int
	TimeoutSeconds int
	MaxRetries     int
}

// AccountModelResolver returns an account-specific model for a feature. An
// empty model means the platform ChatModel should be used.
type AccountModelResolver func(ctx context.Context, accountID uint, feature string) (string, error)

type accountFeatureContextKey struct{}

type accountFeatureContext struct {
	AccountID uint
	Feature   string
}

type generationOverrideContextKey struct{}

type generationOverrideContext struct {
	Temperature *float64
}

// WithAccountFeature annotates an LLM request so ProviderManager can resolve
// the account-specific model selected on the Copilot settings page.
func WithAccountFeature(ctx context.Context, accountID uint, feature string) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, accountFeatureContextKey{}, accountFeatureContext{
		AccountID: accountID,
		Feature:   strings.TrimSpace(feature),
	})
}

// WithTemperatureOverride preserves an explicitly configured Assistant
// temperature while keeping the platform value as the default elsewhere.
func WithTemperatureOverride(ctx context.Context, temperature float64) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, generationOverrideContextKey{}, generationOverrideContext{
		Temperature: &temperature,
	})
}

type providerSnapshot struct {
	chat      Provider
	embedding Provider
	config    RuntimeProviderConfig
}

// ProviderManager is a hot-swappable Provider implementation. Services keep a
// stable reference to the manager while settings updates atomically replace the
// providers and defaults used by new requests.
type ProviderManager struct {
	mu            sync.RWMutex
	snapshot      *providerSnapshot
	modelResolver AccountModelResolver
}

func NewProviderManager() *ProviderManager { return &ProviderManager{} }

func (m *ProviderManager) SetAccountModelResolver(resolver AccountModelResolver) {
	m.mu.Lock()
	m.modelResolver = resolver
	m.mu.Unlock()
}

// Configure validates and constructs all providers before atomically swapping
// the active snapshot. A failed configuration never damages the working one.
func (m *ProviderManager) Configure(cfg RuntimeProviderConfig) error {
	snapshot, err := buildProviderSnapshot(cfg)
	if err != nil {
		return err
	}
	m.mu.Lock()
	m.snapshot = snapshot
	m.mu.Unlock()
	return nil
}

func buildProviderSnapshot(cfg RuntimeProviderConfig) (*providerSnapshot, error) {
	cfg = normalizeRuntimeProviderConfig(cfg)
	if cfg.ChatProvider == "" || cfg.ChatAPIKey == "" || cfg.ChatModel == "" {
		return nil, ErrProviderNotConfigured
	}

	chat, err := buildChatProvider(cfg)
	if err != nil {
		return nil, err
	}
	embedding, err := buildEmbeddingProvider(cfg, chat)
	if err != nil {
		return nil, err
	}
	return &providerSnapshot{chat: chat, embedding: embedding, config: cfg}, nil
}

func normalizeRuntimeProviderConfig(cfg RuntimeProviderConfig) RuntimeProviderConfig {
	cfg.ChatProvider = strings.ToLower(strings.TrimSpace(cfg.ChatProvider))
	cfg.ChatBaseURL = strings.TrimRight(strings.TrimSpace(cfg.ChatBaseURL), "/")
	cfg.ChatAPIKey = strings.TrimSpace(cfg.ChatAPIKey)
	cfg.ChatModel = strings.TrimSpace(cfg.ChatModel)
	cfg.EmbeddingMode = strings.ToLower(strings.TrimSpace(cfg.EmbeddingMode))
	cfg.EmbeddingProvider = strings.ToLower(strings.TrimSpace(cfg.EmbeddingProvider))
	cfg.EmbeddingBaseURL = strings.TrimRight(strings.TrimSpace(cfg.EmbeddingBaseURL), "/")
	cfg.EmbeddingAPIKey = strings.TrimSpace(cfg.EmbeddingAPIKey)
	cfg.EmbeddingModel = strings.TrimSpace(cfg.EmbeddingModel)

	if cfg.EmbeddingMode == "" {
		cfg.EmbeddingMode = EmbeddingModeReuseChat
	}
	if cfg.EmbeddingProvider == "" {
		cfg.EmbeddingProvider = "openai"
	}
	if cfg.EmbeddingModel == "" {
		cfg.EmbeddingModel = "text-embedding-3-small"
	}
	if cfg.EmbeddingDimensions == 0 {
		cfg.EmbeddingDimensions = 1536
	}
	if cfg.MaxTokens == 0 {
		cfg.MaxTokens = 1024
	}
	if cfg.Temperature < 0 {
		cfg.Temperature = 0.7
	}
	if cfg.TimeoutSeconds == 0 {
		cfg.TimeoutSeconds = 60
	}
	if cfg.MaxRetries < 0 {
		cfg.MaxRetries = 3
	}
	return cfg
}

func buildChatProvider(cfg RuntimeProviderConfig) (Provider, error) {
	switch cfg.ChatProvider {
	case "openai", "openai_compatible":
		if cfg.ChatProvider == "openai_compatible" && cfg.ChatBaseURL == "" {
			return nil, fmt.Errorf("Copilot chat base URL is required for openai_compatible")
		}
		return NewOpenAIProvider(OpenAIProviderConfig{
			APIKey:        cfg.ChatAPIKey,
			BaseURL:       cfg.ChatBaseURL,
			Model:         cfg.ChatModel,
			EmbedModel:    cfg.EmbeddingModel,
			MaxRetries:    cfg.MaxRetries,
			MaxRetriesSet: true,
			Timeout:       cfg.TimeoutSeconds,
		}), nil
	case "anthropic":
		return NewAnthropicProvider(AnthropicProviderConfig{
			APIKey:        cfg.ChatAPIKey,
			BaseURL:       cfg.ChatBaseURL,
			Model:         cfg.ChatModel,
			MaxRetries:    cfg.MaxRetries,
			MaxRetriesSet: true,
			Timeout:       cfg.TimeoutSeconds,
		}), nil
	default:
		return nil, fmt.Errorf("unsupported Copilot chat provider: %s", cfg.ChatProvider)
	}
}

func buildEmbeddingProvider(cfg RuntimeProviderConfig, chat Provider) (Provider, error) {
	switch cfg.EmbeddingMode {
	case EmbeddingModeReuseChat:
		if cfg.ChatProvider == "anthropic" {
			return nil, fmt.Errorf("Anthropic chat requires a separate OpenAI-compatible embedding provider")
		}
		return chat, nil
	case EmbeddingModeSeparate:
		if cfg.EmbeddingAPIKey == "" {
			return nil, fmt.Errorf("Copilot embedding API key is required")
		}
		if cfg.EmbeddingProvider != "openai" && cfg.EmbeddingProvider != "openai_compatible" {
			return nil, fmt.Errorf("unsupported Copilot embedding provider: %s", cfg.EmbeddingProvider)
		}
		if cfg.EmbeddingProvider == "openai_compatible" && cfg.EmbeddingBaseURL == "" {
			return nil, fmt.Errorf("Copilot embedding base URL is required for openai_compatible")
		}
		return NewOpenAIProvider(OpenAIProviderConfig{
			APIKey:        cfg.EmbeddingAPIKey,
			BaseURL:       cfg.EmbeddingBaseURL,
			Model:         cfg.ChatModel,
			EmbedModel:    cfg.EmbeddingModel,
			MaxRetries:    cfg.MaxRetries,
			MaxRetriesSet: true,
			Timeout:       cfg.TimeoutSeconds,
		}), nil
	default:
		return nil, fmt.Errorf("unsupported Copilot embedding mode: %s", cfg.EmbeddingMode)
	}
}

func (m *ProviderManager) Clear() {
	m.mu.Lock()
	m.snapshot = nil
	m.mu.Unlock()
}

func (m *ProviderManager) Snapshot() (RuntimeProviderConfig, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.snapshot == nil {
		return RuntimeProviderConfig{}, false
	}
	return m.snapshot.config, true
}

func (m *ProviderManager) current() (*providerSnapshot, AccountModelResolver, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.snapshot == nil {
		return nil, m.modelResolver, ErrProviderNotConfigured
	}
	return m.snapshot, m.modelResolver, nil
}

func resolveFeatureModel(ctx context.Context, resolver AccountModelResolver, fallback string) string {
	if resolver == nil || ctx == nil {
		return fallback
	}
	featureCtx, ok := ctx.Value(accountFeatureContextKey{}).(accountFeatureContext)
	if !ok || featureCtx.AccountID == 0 || featureCtx.Feature == "" {
		return fallback
	}
	model, err := resolver(ctx, featureCtx.AccountID, featureCtx.Feature)
	if err != nil || strings.TrimSpace(model) == "" {
		return fallback
	}
	return strings.TrimSpace(model)
}

func resolveRuntimeChatModel(ctx context.Context, cfg RuntimeProviderConfig, resolver AccountModelResolver) string {
	if cfg.ChatProvider == "openai" {
		return resolveFeatureModel(ctx, resolver, cfg.ChatModel)
	}
	return cfg.ChatModel
}

// ResolveChatModel returns the model ChatCompletion will use after applying an
// account feature override.
func (m *ProviderManager) ResolveChatModel(ctx context.Context) (string, error) {
	snapshot, resolver, err := m.current()
	if err != nil {
		return "", err
	}
	return resolveRuntimeChatModel(ctx, snapshot.config, resolver), nil
}

func applyRuntimeChatConfig(ctx context.Context, req ChatRequest, cfg RuntimeProviderConfig, resolver AccountModelResolver) ChatRequest {
	// Account overrides store only a model name. Keep provider/model pairs intact
	// for compatible endpoints, where an OpenAI model may not exist.
	req.Model = resolveRuntimeChatModel(ctx, cfg, resolver)
	req.Temperature = cfg.Temperature
	if override, ok := ctx.Value(generationOverrideContextKey{}).(generationOverrideContext); ok && override.Temperature != nil {
		req.Temperature = *override.Temperature
	}
	req.MaxTokens = cfg.MaxTokens
	return req
}

func (m *ProviderManager) ChatCompletion(ctx context.Context, req ChatRequest) (*ChatResponse, error) {
	snapshot, resolver, err := m.current()
	if err != nil {
		return nil, err
	}
	req = applyRuntimeChatConfig(ctx, req, snapshot.config, resolver)
	return snapshot.chat.ChatCompletion(ctx, req)
}

func (m *ProviderManager) CreateEmbedding(ctx context.Context, req EmbeddingRequest) (*EmbeddingResponse, error) {
	snapshot, _, err := m.current()
	if err != nil {
		return nil, err
	}
	req.Model = snapshot.config.EmbeddingModel
	req.Dimensions = snapshot.config.EmbeddingDimensions
	return snapshot.embedding.CreateEmbedding(ctx, req)
}

func (m *ProviderManager) ChatCompletionStream(ctx context.Context, req ChatRequest, onChunk func(StreamChunk) error) error {
	snapshot, resolver, err := m.current()
	if err != nil {
		return err
	}
	req = applyRuntimeChatConfig(ctx, req, snapshot.config, resolver)
	return snapshot.chat.ChatCompletionStream(ctx, req, onChunk)
}

var _ Provider = (*ProviderManager)(nil)