go-llm/v2/registry.go

package llm

import (
	"os"
	"sync"

	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/deepseek"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/groq"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/moonshot"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/ollama"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/openai"
	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/xai"
)

// ProviderInfo describes a registered provider for discovery purposes (CLI
// pickers, wiring layers, admin tools). It is the single source of truth for
// "what providers exist and how do I instantiate one."
type ProviderInfo struct {
	// Name is the short lowercase identifier used in provider/model strings
	// (e.g., "openai", "deepseek", "moonshot").
	Name string

	// DisplayName is a human-readable label for UIs.
	DisplayName string

	// EnvKey is the conventional environment variable that holds the API key
	// for this provider. Empty string means "no key needed" (e.g., Ollama).
	EnvKey string

	// DefaultURL is the default base URL used when no override is supplied.
	DefaultURL string

	// Models is a list of well-known model names, populated for CLI pickers
	// and similar. It is not exhaustive and not validated against the API.
	Models []string

	// New returns a ready-to-use Client for this provider, given an API key
	// (ignored for key-less providers like Ollama) and optional ClientOptions.
	New func(apiKey string, opts ...ClientOption) *Client
}

// providerRegistry is the in-process list of known providers. Order is
// intentional: the three original providers first, then OpenAI-compatible
// additions in the order they were added. This slice seeds NewRegistry().
var providerRegistry = []ProviderInfo{
	{
		Name:        "openai",
		DisplayName: "OpenAI",
		EnvKey:      "OPENAI_API_KEY",
		DefaultURL:  openai.DefaultBaseURL,
		Models: []string{
			"gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano",
			"gpt-4o", "gpt-4o-mini",
			"gpt-4-turbo", "gpt-3.5-turbo",
			"o1", "o1-mini", "o1-preview", "o3-mini",
		},
		New: OpenAI,
	},
	{
		Name:        "anthropic",
		DisplayName: "Anthropic",
		EnvKey:      "ANTHROPIC_API_KEY",
		DefaultURL:  "https://api.anthropic.com",
		Models: []string{
			"claude-opus-4-7",
			"claude-sonnet-4-6",
			"claude-haiku-4-5-20251001",
			"claude-opus-4-20250514",
			"claude-sonnet-4-20250514",
			"claude-3-7-sonnet-20250219",
			"claude-3-5-sonnet-20241022",
			"claude-3-5-haiku-20241022",
		},
		New: Anthropic,
	},
	{
		Name:        "google",
		DisplayName: "Google",
		EnvKey:      "GOOGLE_API_KEY",
		DefaultURL:  "https://generativelanguage.googleapis.com",
		Models: []string{
			"gemini-2.0-flash", "gemini-2.0-flash-lite",
			"gemini-1.5-pro", "gemini-1.5-flash", "gemini-1.5-flash-8b",
		},
		New: Google,
	},
	{
		Name:        "deepseek",
		DisplayName: "DeepSeek",
		EnvKey:      "DEEPSEEK_API_KEY",
		DefaultURL:  deepseek.DefaultBaseURL,
		Models:      []string{"deepseek-chat", "deepseek-reasoner"},
		New:         DeepSeek,
	},
	{
		Name:        "moonshot",
		DisplayName: "Moonshot (Kimi)",
		EnvKey:      "MOONSHOT_API_KEY",
		DefaultURL:  moonshot.DefaultBaseURL,
		Models: []string{
			"kimi-k2-0711-preview",
			"moonshot-v1-8k", "moonshot-v1-32k", "moonshot-v1-128k",
			"moonshot-v1-8k-vision-preview",
		},
		New: Moonshot,
	},
	{
		Name:        "xai",
		DisplayName: "xAI (Grok)",
		EnvKey:      "XAI_API_KEY",
		DefaultURL:  xai.DefaultBaseURL,
		Models: []string{
			"grok-2", "grok-2-mini", "grok-2-vision", "grok-beta",
		},
		New: XAI,
	},
	{
		Name:        "groq",
		DisplayName: "Groq",
		EnvKey:      "GROQ_API_KEY",
		DefaultURL:  groq.DefaultBaseURL,
		Models: []string{
			"llama-3.3-70b-versatile",
			"llama-3.1-8b-instant",
			"mixtral-8x7b-32768",
			"gemma2-9b-it",
			"llama-3.2-90b-vision-preview",
		},
		New: Groq,
	},
	{
		Name:        "ollama",
		DisplayName: "Ollama (local)",
		EnvKey:      "", // no key needed
		DefaultURL:  ollama.DefaultLocalBaseURL,
		Models: []string{
			"llama3.2", "llama3.1", "qwen2.5", "mistral", "gemma2", "phi4",
		},
		New: func(_ string, opts ...ClientOption) *Client { return Ollama(opts...) },
	},
	{
		Name:        "ollama-cloud",
		DisplayName: "Ollama Cloud",
		EnvKey:      "OLLAMA_API_KEY",
		DefaultURL:  ollama.DefaultCloudBaseURL,
		Models: []string{
			"ministral-3:14b",
			"kimi-k2.5", "kimi-k2.6",
			"qwen3.5:122b",
			"gemma4:31b",
			"deepseek-v4-flash", "deepseek-v4-pro",
			"glm-5.1",
			"gemini-3-flash-preview",
		},
		New: OllamaCloud,
	},
	{
		Name:        "foreman",
		DisplayName: "Foreman",
		EnvKey:      "", // no single env key; discovered via LLM_* DSNs
		DefaultURL:  "", // always requires a URL
		Models:      []string{},
		New:         Foreman,
	},
}

// ---------------------------------------------------------------------------
// Resolver — dynamic model-spec resolution
// ---------------------------------------------------------------------------

// Resolver resolves an alias or short name to a full spec string. Consumers
// register resolvers for dynamic lookups (e.g., database-backed tier aliases).
type Resolver interface {
	// Resolve returns the resolved spec and an optional default reasoning
	// level. ok is false when the resolver does not handle this name.
	Resolve(name string) (spec string, defaultReasoning ReasoningLevel, ok bool)
}

// ResolverFunc adapts a plain function to the Resolver interface.
type ResolverFunc func(name string) (string, ReasoningLevel, bool)

// Resolve implements Resolver by calling the underlying function.
func (f ResolverFunc) Resolve(name string) (string, ReasoningLevel, bool) {
	return f(name)
}

// ---------------------------------------------------------------------------
// Registry — extensible provider/alias/resolver store
// ---------------------------------------------------------------------------

// Registry holds providers, static aliases, and dynamic resolvers. Use
// NewRegistry to create one pre-populated with the built-in providers, or
// use the package-level DefaultRegistry.
type Registry struct {
	mu        sync.RWMutex
	providers map[string]ProviderInfo
	order     []string // insertion order for Providers()
	aliases   map[string]string
	resolvers []Resolver
	envLookup func(string) string // defaults to os.Getenv
}

// NewRegistry creates a Registry pre-populated with all built-in providers
// (the same set returned by the providerRegistry package variable).
//
// Why: provides a fresh, isolated registry for testing or multi-tenant
// scenarios while reusing the canonical provider list.
// What: copies every entry from providerRegistry into a new Registry.
// Test: call NewRegistry(), verify Providers() length matches providerRegistry
// and ProviderByName("openai") is non-nil.
func NewRegistry() *Registry {
	r := &Registry{
		providers: make(map[string]ProviderInfo, len(providerRegistry)),
		order:     make([]string, 0, len(providerRegistry)),
		aliases:   make(map[string]string),
		envLookup: os.Getenv,
	}
	for _, info := range providerRegistry {
		r.providers[info.Name] = info
		r.order = append(r.order, info.Name)
	}
	return r
}

// DefaultRegistry is the package-level registry used by the convenience
// functions Parse, Providers, ProviderByName, RegisterProvider, RegisterAlias,
// and RegisterResolver. Initialized in init() with all built-in providers.
var DefaultRegistry *Registry

func init() {
	DefaultRegistry = NewRegistry()
}

// RegisterProvider adds or replaces a provider in the registry. When
// replacing, the provider keeps its original position in the ordered list.
//
// Why: allows consumers to override built-in factories (e.g., wrapping with
// middleware) or add entirely new providers at runtime.
// What: upserts info by Name into the provider map and order slice.
// Test: register a custom "openai" factory, verify ProviderByName returns it.
func (r *Registry) RegisterProvider(info ProviderInfo) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if _, exists := r.providers[info.Name]; !exists {
		r.order = append(r.order, info.Name)
	}
	r.providers[info.Name] = info
}

// RegisterAlias maps a short name to a full spec string. The spec is resolved
// recursively by Parse, so an alias can point to another alias or to a
// "provider/model" string.
//
// Why: lets consumers define convenient shortcuts like "fast" → "openai/gpt-4o-mini".
// What: stores name→spec in the alias map.
// Test: register "fast" → "openai/gpt-4o-mini", parse "fast", verify model.
func (r *Registry) RegisterAlias(name, spec string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.aliases[name] = spec
}

// RegisterResolver appends a dynamic resolver. Resolvers are checked in
// registration order after static aliases. A resolver may return a spec
// string that is itself an alias or "provider/model" — it will be recursed.
//
// Why: supports dynamic alias sources (databases, remote config) without
// requiring static registration of every possible name.
// What: appends res to the resolver list.
// Test: register a ResolverFunc, parse a name it handles, verify resolution.
func (r *Registry) RegisterResolver(res Resolver) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.resolvers = append(r.resolvers, res)
}

// ProviderByName returns the registered ProviderInfo with the given name, or
// nil if no such provider is registered. Name matching is exact.
//
// Why: callers need to look up provider metadata by name for factory calls,
// discovery, and DSN scheme resolution.
// What: returns a copy of the ProviderInfo or nil.
// Test: verify ProviderByName("openai") is non-nil, ProviderByName("nope") is nil.
func (r *Registry) ProviderByName(name string) *ProviderInfo {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if info, ok := r.providers[name]; ok {
		return &info
	}
	return nil
}

// Providers returns a copy of all registered providers in insertion order.
//
// Why: CLI pickers and admin tools need the full list for display.
// What: returns a freshly allocated slice of ProviderInfo copies.
// Test: verify length matches expected count after registration.
func (r *Registry) Providers() []ProviderInfo {
	r.mu.RLock()
	defer r.mu.RUnlock()

	out := make([]ProviderInfo, 0, len(r.order))
	for _, name := range r.order {
		if info, ok := r.providers[name]; ok {
			out = append(out, info)
		}
	}
	return out
}

// ---------------------------------------------------------------------------
// Package-level convenience functions — delegate to DefaultRegistry
// ---------------------------------------------------------------------------

// Providers returns a copy of the registered provider list so callers cannot
// mutate library state.
func Providers() []ProviderInfo {
	return DefaultRegistry.Providers()
}

// ProviderByName returns the registered ProviderInfo with the given name, or
// nil if no such provider is registered. Name matching is exact.
func ProviderByName(name string) *ProviderInfo {
	return DefaultRegistry.ProviderByName(name)
}

// RegisterProvider adds or replaces a provider in the DefaultRegistry.
func RegisterProvider(info ProviderInfo) {
	DefaultRegistry.RegisterProvider(info)
}

// RegisterAlias maps a short name to a full spec in the DefaultRegistry.
func RegisterAlias(name, spec string) {
	DefaultRegistry.RegisterAlias(name, spec)
}

// RegisterResolver appends a dynamic resolver to the DefaultRegistry.
func RegisterResolver(res Resolver) {
	DefaultRegistry.RegisterResolver(res)
}