diff --git a/CLAUDE.md b/CLAUDE.md index 3a8f7718..356c66cc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,13 +1,21 @@ # CLAUDE.md -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. +This file provides guidance to Claude and other AI models and agents when working with the code in this repository. ## Project Overview -GOModel is a high-performance AI gateway written in Go that routes requests to multiple LLM providers (OpenAI, Anthropic, Google Gemini). It's designed as a drop-in replacement for LiteLLM with superior concurrency and strict type safety. +GOModel is a high-performance AI gateway written in Go that routes requests to multiple LLM providers (OpenAI, Anthropic, Google Gemini). It is designed as a drop-in replacement for LiteLLM, offering superior quality, speed, concurrency, and strict type safety. **Module name:** `gomodel` **Go version:** 1.24.0 +**GitHub link:** https://github.com/ENTERPILOT/GOModel +**Stage:** Development — not used in a real project, so backward compatibility is not a concern. + +### Introductory Information + +Please note that some of the information below may be slightly outdated. It is always best to verify the current project structure. These are guidelines only. + +Feel free to suggest improvements to this file as you work with it. ## Common Development Commands @@ -124,6 +132,22 @@ type Provider interface { - Automatically creates provider configs from environment variables - At least one provider API key is required +**Provider naming in config.yaml:** + +Provider names (map keys) are arbitrary identifiers. Only the `type` field determines which provider implementation is used. You can use any name and have multiple providers of the same type: + +```yaml +providers: + my-openai: # Any name works + type: "openai" + api_key: "${OPENAI_API_KEY}" + + backup-openai: # Multiple providers of same type supported + type: "openai" + api_key: "${BACKUP_OPENAI_KEY}" + base_url: "https://custom.endpoint.com" +``` + ### Adding a New Provider 1. Create package in `internal/providers/{name}/` diff --git a/METRICS_CONFIGURATION.md b/METRICS_CONFIGURATION.md index 9a3ef6c5..64bd2abc 100644 --- a/METRICS_CONFIGURATION.md +++ b/METRICS_CONFIGURATION.md @@ -105,7 +105,7 @@ metrics: endpoint: "/internal/prometheus" # Custom path providers: - openai-primary: + openai: type: "openai" api_key: "${OPENAI_API_KEY}" ``` diff --git a/cmd/gomodel/main.go b/cmd/gomodel/main.go index f8511cd7..ad546b07 100644 --- a/cmd/gomodel/main.go +++ b/cmd/gomodel/main.go @@ -3,28 +3,23 @@ package main import ( "context" - "errors" "flag" "fmt" "log/slog" - "net/http" "os" "os/signal" "syscall" "time" "gomodel/config" - "gomodel/internal/auditlog" + "gomodel/internal/app" "gomodel/internal/observability" "gomodel/internal/providers" - - // Import provider packages to trigger their init() registration - _ "gomodel/internal/providers/anthropic" - _ "gomodel/internal/providers/gemini" - _ "gomodel/internal/providers/groq" - _ "gomodel/internal/providers/openai" - _ "gomodel/internal/providers/xai" - "gomodel/internal/server" + "gomodel/internal/providers/anthropic" + "gomodel/internal/providers/gemini" + "gomodel/internal/providers/groq" + "gomodel/internal/providers/openai" + "gomodel/internal/providers/xai" "gomodel/internal/version" ) @@ -62,62 +57,30 @@ func main() { os.Exit(1) } - // Setup observability hooks for metrics collection (if enabled) - // This must be done BEFORE creating providers so they can use the hooks + // Create provider factory and register all providers explicitly + factory := providers.NewProviderFactory() + + // Set observability hooks before registering providers if cfg.Metrics.Enabled { - metricsHooks := observability.NewPrometheusHooks() - providers.SetGlobalHooks(metricsHooks) - slog.Info("prometheus metrics enabled", "endpoint", cfg.Metrics.Endpoint) - } else { - slog.Info("prometheus metrics disabled") + factory.SetHooks(observability.NewPrometheusHooks()) } - // Initialize provider infrastructure (cache, registry, router) - providerResult, err := providers.Init(context.Background(), cfg) + // Register all providers with the factory + factory.Register(openai.Registration) + factory.Register(anthropic.Registration) + factory.Register(gemini.Registration) + factory.Register(groq.Registration) + factory.Register(xai.Registration) + + // Create the application + application, err := app.New(context.Background(), app.Config{ + AppConfig: cfg, + Factory: factory, + }) if err != nil { - slog.Error("failed to initialize providers", "error", err) + slog.Error("failed to initialize application", "error", err) os.Exit(1) } - defer providerResult.Close() - - // Security check: warn if no master key is configured - if cfg.Server.MasterKey == "" { - slog.Warn("SECURITY WARNING: GOMODEL_MASTER_KEY not set - server running in UNSAFE MODE", - "security_risk", "unauthenticated access allowed", - "recommendation", "set GOMODEL_MASTER_KEY environment variable to secure this gateway") - } else { - slog.Info("authentication enabled", "mode", "master_key") - } - - // Initialize audit logging - auditResult, err := auditlog.New(context.Background(), cfg) - if err != nil { - slog.Error("failed to initialize audit logging", "error", err) - os.Exit(1) - } - defer auditResult.Close() - - if cfg.Logging.Enabled { - slog.Info("audit logging enabled", - "storage_type", cfg.Logging.StorageType, - "log_bodies", cfg.Logging.LogBodies, - "log_headers", cfg.Logging.LogHeaders, - "retention_days", cfg.Logging.RetentionDays, - ) - } else { - slog.Info("audit logging disabled") - } - - // Create and start server - serverCfg := &server.Config{ - MasterKey: cfg.Server.MasterKey, - MetricsEnabled: cfg.Metrics.Enabled, - MetricsEndpoint: cfg.Metrics.Endpoint, - BodySizeLimit: cfg.Server.BodySizeLimit, - AuditLogger: auditResult.Logger, - LogOnlyModelInteractions: cfg.Logging.OnlyModelInteractions, - } - srv := server.New(providerResult.Router, serverCfg) // Handle graceful shutdown go func() { @@ -125,25 +88,18 @@ func main() { signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) <-quit - slog.Info("shutting down server...") - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - if err := srv.Shutdown(ctx); err != nil { - slog.Error("server shutdown error", "error", err) + if err := application.Shutdown(ctx); err != nil { + slog.Error("application shutdown error", "error", err) } }() + // Start the server (blocking) addr := ":" + cfg.Server.Port - slog.Info("starting server", "address", addr) - - if err := srv.Start(addr); err != nil { - if errors.Is(err, http.ErrServerClosed) { - slog.Info("server stopped gracefully") - } else { - slog.Error("server failed to start", "error", err) - os.Exit(1) - } + if err := application.Start(addr); err != nil { + slog.Error("application failed", "error", err) + os.Exit(1) } } diff --git a/config/config.go b/config/config.go index 2fde8f10..bee69d15 100644 --- a/config/config.go +++ b/config/config.go @@ -8,6 +8,7 @@ import ( "strconv" "strings" + "github.com/go-viper/mapstructure/v2" "github.com/joho/godotenv" "github.com/spf13/viper" ) @@ -24,132 +25,152 @@ var bodySizeLimitRegex = regexp.MustCompile(`(?i)^(\d+)([KMG])?B?$`) // Config holds the application configuration type Config struct { - Server ServerConfig `mapstructure:"server"` - Cache CacheConfig `mapstructure:"cache"` - Storage StorageConfig `mapstructure:"storage"` - Logging LogConfig `mapstructure:"logging"` - Metrics MetricsConfig `mapstructure:"metrics"` - Providers map[string]ProviderConfig `mapstructure:"providers"` + Server ServerConfig + Cache CacheConfig + Storage StorageConfig + Logging LogConfig + Metrics MetricsConfig + Providers map[string]ProviderConfig } // LogConfig holds audit logging configuration type LogConfig struct { // Enabled controls whether audit logging is active // Default: false - Enabled bool `mapstructure:"enabled"` + Enabled bool // StorageType specifies the storage backend for audit logs: "sqlite" (default), "postgresql", or "mongodb" // This selects which of the storage backends (configured separately) to use for audit logs - StorageType string `mapstructure:"storage_type"` + StorageType string // LogBodies enables logging of full request/response bodies // WARNING: May contain sensitive data (PII, API keys in prompts) // Default: true - LogBodies bool `mapstructure:"log_bodies"` + LogBodies bool // LogHeaders enables logging of request/response headers // Sensitive headers (Authorization, Cookie, etc.) are auto-redacted // Default: true - LogHeaders bool `mapstructure:"log_headers"` + LogHeaders bool // BufferSize is the number of log entries to buffer before flushing // Default: 1000 - BufferSize int `mapstructure:"buffer_size"` + BufferSize int // FlushInterval is how often to flush buffered logs (in seconds) // Default: 5 - FlushInterval int `mapstructure:"flush_interval"` + FlushInterval int // RetentionDays is how long to keep logs (0 = forever) // Default: 30 - RetentionDays int `mapstructure:"retention_days"` + RetentionDays int // OnlyModelInteractions limits audit logging to AI model endpoints only // When true, only /v1/chat/completions, /v1/responses, /v1/models are logged // Endpoints like /health, /metrics, /admin are skipped // Default: true - OnlyModelInteractions bool `mapstructure:"only_model_interactions"` + OnlyModelInteractions bool } // StorageConfig holds database storage configuration (used by audit logging, future IAM, etc.) type StorageConfig struct { // SQLite configuration - SQLite SQLiteStorageConfig `mapstructure:"sqlite"` + SQLite SQLiteStorageConfig // PostgreSQL configuration - PostgreSQL PostgreSQLStorageConfig `mapstructure:"postgresql"` + PostgreSQL PostgreSQLStorageConfig // MongoDB configuration - MongoDB MongoDBStorageConfig `mapstructure:"mongodb"` + MongoDB MongoDBStorageConfig } // SQLiteStorageConfig holds SQLite-specific storage configuration type SQLiteStorageConfig struct { // Path is the database file path (default: .cache/gomodel.db) - Path string `mapstructure:"path"` + Path string } // PostgreSQLStorageConfig holds PostgreSQL-specific storage configuration type PostgreSQLStorageConfig struct { // URL is the connection string (e.g., postgres://user:pass@localhost/dbname) - URL string `mapstructure:"url"` + URL string // MaxConns is the maximum connection pool size (default: 10) - MaxConns int `mapstructure:"max_conns"` + MaxConns int } // MongoDBStorageConfig holds MongoDB-specific storage configuration type MongoDBStorageConfig struct { // URL is the connection string (e.g., mongodb://localhost:27017) - URL string `mapstructure:"url"` + URL string // Database is the database name (default: gomodel) - Database string `mapstructure:"database"` + Database string } // CacheConfig holds cache configuration for model storage type CacheConfig struct { // Type specifies the cache backend: "local" (default) or "redis" - Type string `mapstructure:"type"` + Type string // Redis configuration (only used when Type is "redis") - Redis RedisConfig `mapstructure:"redis"` + Redis RedisConfig } // RedisConfig holds Redis-specific configuration type RedisConfig struct { // URL is the Redis connection URL (e.g., "redis://localhost:6379") - URL string `mapstructure:"url"` + URL string // Key is the Redis key for storing the model cache (default: "gomodel:models") - Key string `mapstructure:"key"` + Key string // TTL is the time-to-live for cached data in seconds (default: 86400 = 24 hours) - TTL int `mapstructure:"ttl"` + TTL int } // ServerConfig holds HTTP server configuration type ServerConfig struct { - Port string `mapstructure:"port"` - MasterKey string `mapstructure:"master_key"` // Optional: Master key for authentication - BodySizeLimit string `mapstructure:"body_size_limit"` // Max request body size (e.g., "10M", "1024K") + Port string + MasterKey string // Optional: Master key for authentication + BodySizeLimit string // Max request body size (e.g., "10M", "1024K") } // MetricsConfig holds observability configuration for Prometheus metrics type MetricsConfig struct { // Enabled controls whether Prometheus metrics are collected and exposed // Default: false - Enabled bool `mapstructure:"enabled"` + Enabled bool // Endpoint is the HTTP path where metrics are exposed // Default: "/metrics" - Endpoint string `mapstructure:"endpoint"` + Endpoint string } // ProviderConfig holds generic provider configuration type ProviderConfig struct { - Type string `mapstructure:"type"` // e.g., "openai", "anthropic", "gemini" - APIKey string `mapstructure:"api_key"` // API key for authentication - BaseURL string `mapstructure:"base_url"` // Optional: override default base URL - Models []string `mapstructure:"models"` // Optional: restrict to specific models + Type string // e.g., "openai", "anthropic", "gemini" + APIKey string // API key for authentication + BaseURL string // Optional: override default base URL + Models []string // Optional: restrict to specific models +} + +// snakeCaseMatchName returns a DecoderConfigOption that matches +// snake_case map keys to PascalCase struct field names. +func snakeCaseMatchName() viper.DecoderConfigOption { + return func(c *mapstructure.DecoderConfig) { + c.MatchName = func(mapKey, fieldName string) bool { + // Reject malformed keys: leading/trailing underscores or consecutive underscores + if strings.HasPrefix(mapKey, "_") || strings.HasSuffix(mapKey, "_") { + return false + } + if strings.Contains(mapKey, "__") { + return false + } + + // Remove underscores and compare case-insensitively + normalizedKey := strings.ReplaceAll(mapKey, "_", "") + return strings.EqualFold(normalizedKey, fieldName) + } + } } // Load reads configuration from file and environment @@ -202,7 +223,7 @@ func Load() (*Config, error) { // Read config file (optional, won't fail if not found) if err := viper.ReadInConfig(); err == nil { // Config file found, unmarshal it - if err := viper.Unmarshal(&cfg); err != nil { + if err := viper.Unmarshal(&cfg, snakeCaseMatchName()); err != nil { return nil, err } // Expand environment variables in config values @@ -250,31 +271,31 @@ func Load() (*Config, error) { // TODO: Similarly for ENV variables. All ENV variables like *_API_KEY should be taken and iterated over // Add providers from environment variables if available if apiKey := viper.GetString("OPENAI_API_KEY"); apiKey != "" { - cfg.Providers["openai-primary"] = ProviderConfig{ + cfg.Providers["openai"] = ProviderConfig{ Type: "openai", APIKey: apiKey, } } if apiKey := viper.GetString("ANTHROPIC_API_KEY"); apiKey != "" { - cfg.Providers["anthropic-primary"] = ProviderConfig{ + cfg.Providers["anthropic"] = ProviderConfig{ Type: "anthropic", APIKey: apiKey, } } if apiKey := viper.GetString("GEMINI_API_KEY"); apiKey != "" { - cfg.Providers["gemini-primary"] = ProviderConfig{ + cfg.Providers["gemini"] = ProviderConfig{ Type: "gemini", APIKey: apiKey, } } if apiKey := viper.GetString("XAI_API_KEY"); apiKey != "" { - cfg.Providers["xai-primary"] = ProviderConfig{ + cfg.Providers["xai"] = ProviderConfig{ Type: "xai", APIKey: apiKey, } } if apiKey := viper.GetString("GROQ_API_KEY"); apiKey != "" { - cfg.Providers["groq-primary"] = ProviderConfig{ + cfg.Providers["groq"] = ProviderConfig{ Type: "groq", APIKey: apiKey, } diff --git a/config/config.yaml b/config/config.yaml index e7e34f07..02d991d8 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -33,23 +33,23 @@ metrics: endpoint: "${METRICS_ENDPOINT:-/metrics}" providers: - openai-primary: + openai: type: "openai" api_key: "${OPENAI_API_KEY}" - anthropic-primary: + anthropic: type: "anthropic" api_key: "${ANTHROPIC_API_KEY}" - gemini-primary: + gemini: type: "gemini" api_key: "${GEMINI_API_KEY}" - xai-primary: + xai: type: "xai" api_key: "${XAI_API_KEY}" - groq-primary: + groq: type: "groq" api_key: "${GROQ_API_KEY}" diff --git a/config/config_defaults_test.go b/config/config_defaults_test.go index af294af9..bdc8d681 100644 --- a/config/config_defaults_test.go +++ b/config/config_defaults_test.go @@ -39,7 +39,7 @@ func TestLoad_WithDefaults(t *testing.T) { server: port: "${TEST_PORT_DEFAULTS:-9999}" providers: - openai-primary: + openai: type: "openai" api_key: "${TEST_KEY_DEFAULTS:-default-key}" ` @@ -70,7 +70,7 @@ providers: t.Errorf("Expected port 9999 (default), got %s", cfg.Server.Port) } - provider := cfg.Providers["openai-primary"] + provider := cfg.Providers["openai"] if provider.APIKey != "default-key" { t.Errorf("Expected API key 'default-key', got %s", provider.APIKey) } @@ -117,7 +117,7 @@ providers: server: port: "${TEST_PORT_DEFAULTS:-9999}" providers: - openai-primary: + openai: type: "openai" api_key: "${TEST_KEY_DEFAULTS:-default-key}" ` @@ -137,7 +137,7 @@ providers: t.Errorf("Expected port 1111 (env override), got %s", cfg.Server.Port) } - provider := cfg.Providers["openai-primary"] + provider := cfg.Providers["openai"] if provider.APIKey != "real-key" { t.Errorf("Expected API key 'real-key', got %s", provider.APIKey) } diff --git a/config/config_example_test.go b/config/config_example_test.go index 75348b95..7016a301 100644 --- a/config/config_example_test.go +++ b/config/config_example_test.go @@ -34,7 +34,7 @@ func TestLoad_FromEnvironment(t *testing.T) { } // Check OpenAI provider - if openaiCfg, exists := cfg.Providers["openai-primary"]; exists { + if openaiCfg, exists := cfg.Providers["openai"]; exists { if openaiCfg.Type != "openai" { t.Errorf("expected openai type, got '%s'", openaiCfg.Type) } @@ -42,11 +42,11 @@ func TestLoad_FromEnvironment(t *testing.T) { t.Errorf("expected openai key 'test-openai-key', got '%s'", openaiCfg.APIKey) } } else { - t.Error("expected 'openai-primary' provider to exist") + t.Error("expected 'openai' provider to exist") } // Check Anthropic provider - if anthropicCfg, exists := cfg.Providers["anthropic-primary"]; exists { + if anthropicCfg, exists := cfg.Providers["anthropic"]; exists { if anthropicCfg.Type != "anthropic" { t.Errorf("expected anthropic type, got '%s'", anthropicCfg.Type) } @@ -54,7 +54,7 @@ func TestLoad_FromEnvironment(t *testing.T) { t.Errorf("expected anthropic key 'test-anthropic-key', got '%s'", anthropicCfg.APIKey) } } else { - t.Error("expected 'anthropic-primary' provider to exist") + t.Error("expected 'anthropic' provider to exist") } } diff --git a/config/config_helpers_test.go b/config/config_helpers_test.go index 0dd65541..04e575a3 100644 --- a/config/config_helpers_test.go +++ b/config/config_helpers_test.go @@ -552,7 +552,7 @@ func TestRemoveEmptyProviders(t *testing.T) { input: Config{ Server: ServerConfig{Port: "8080"}, Providers: map[string]ProviderConfig{ - "openai-primary": { + "openai": { Type: "openai", APIKey: "sk-openai-valid", }, @@ -573,7 +573,7 @@ func TestRemoveEmptyProviders(t *testing.T) { expected: Config{ Server: ServerConfig{Port: "8080"}, Providers: map[string]ProviderConfig{ - "openai-primary": { + "openai": { Type: "openai", APIKey: "sk-openai-valid", }, @@ -785,7 +785,7 @@ func TestIntegration_ExpandAndFilter(t *testing.T) { Port: "${PORT:-8080}", }, Providers: map[string]ProviderConfig{ - "openai-primary": { + "openai": { Type: "openai", APIKey: "${OPENAI_API_KEY}", }, @@ -808,7 +808,7 @@ func TestIntegration_ExpandAndFilter(t *testing.T) { Port: "8080", }, Providers: map[string]ProviderConfig{ - "openai-primary": { + "openai": { Type: "openai", APIKey: "sk-openai-123", }, diff --git a/config/config_test.go b/config/config_test.go index 32c71626..351154c8 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -4,6 +4,7 @@ import ( "os" "testing" + "github.com/go-viper/mapstructure/v2" "github.com/spf13/viper" ) @@ -67,9 +68,9 @@ func TestLoad_OpenAIAPIKeyFromEnv(t *testing.T) { } // Check that OpenAI provider was created from environment variable - provider, exists := cfg.Providers["openai-primary"] + provider, exists := cfg.Providers["openai"] if !exists { - t.Fatal("expected 'openai-primary' provider to exist") + t.Fatal("expected 'openai' provider to exist") } if provider.Type != "openai" { @@ -131,17 +132,17 @@ func TestLoad_MultipleEnvVars(t *testing.T) { } // Check OpenAI provider - openaiProvider, exists := cfg.Providers["openai-primary"] + openaiProvider, exists := cfg.Providers["openai"] if !exists { - t.Error("expected 'openai-primary' provider to exist") + t.Error("expected 'openai' provider to exist") } else if openaiProvider.APIKey != testAPIKey { t.Errorf("expected OpenAI API key %s, got %s", testAPIKey, openaiProvider.APIKey) } // Check Anthropic provider - anthropicProvider, exists := cfg.Providers["anthropic-primary"] + anthropicProvider, exists := cfg.Providers["anthropic"] if !exists { - t.Error("expected 'anthropic-primary' provider to exist") + t.Error("expected 'anthropic' provider to exist") } else if anthropicProvider.APIKey != testAnthropicKey { t.Errorf("expected Anthropic API key %s, got %s", testAnthropicKey, anthropicProvider.APIKey) } @@ -184,7 +185,7 @@ OPENAI_API_KEY=sk-from-dotenv-file // Add provider from environment variable if apiKey := viper.GetString("OPENAI_API_KEY"); apiKey != "" { - cfg.Providers["openai-primary"] = ProviderConfig{ + cfg.Providers["openai"] = ProviderConfig{ Type: "openai", APIKey: apiKey, } @@ -195,9 +196,9 @@ OPENAI_API_KEY=sk-from-dotenv-file t.Errorf("expected port 7070 from .env file, got %s", cfg.Server.Port) } - openaiProvider, exists := cfg.Providers["openai-primary"] + openaiProvider, exists := cfg.Providers["openai"] if !exists { - t.Fatal("expected 'openai-primary' provider to exist") + t.Fatal("expected 'openai' provider to exist") } if openaiProvider.APIKey != "sk-from-dotenv-file" { @@ -246,7 +247,7 @@ OPENAI_API_KEY=sk-from-dotenv-file // Add provider from environment variable if apiKey := viper.GetString("OPENAI_API_KEY"); apiKey != "" { - cfg.Providers["openai-primary"] = ProviderConfig{ + cfg.Providers["openai"] = ProviderConfig{ Type: "openai", APIKey: apiKey, } @@ -257,9 +258,9 @@ OPENAI_API_KEY=sk-from-dotenv-file t.Errorf("expected port 9999 from environment variable (not .env file), got %s", cfg.Server.Port) } - openaiProvider, exists := cfg.Providers["openai-primary"] + openaiProvider, exists := cfg.Providers["openai"] if !exists { - t.Fatal("expected 'openai-primary' provider to exist") + t.Fatal("expected 'openai' provider to exist") } if openaiProvider.APIKey != "sk-from-real-env" { @@ -325,6 +326,162 @@ func TestLoggingOnlyModelInteractionsFromEnv(t *testing.T) { } } +func TestSnakeCaseMatchName(t *testing.T) { + tests := []struct { + name string + mapKey string + fieldName string + expected bool + }{ + // Snake case to PascalCase matches + {"body_size_limit matches BodySizeLimit", "body_size_limit", "BodySizeLimit", true}, + {"api_key matches APIKey", "api_key", "APIKey", true}, + {"base_url matches BaseURL", "base_url", "BaseURL", true}, + {"storage_type matches StorageType", "storage_type", "StorageType", true}, + {"log_bodies matches LogBodies", "log_bodies", "LogBodies", true}, + {"only_model_interactions matches OnlyModelInteractions", "only_model_interactions", "OnlyModelInteractions", true}, + {"max_conns matches MaxConns", "max_conns", "MaxConns", true}, + {"flush_interval matches FlushInterval", "flush_interval", "FlushInterval", true}, + {"retention_days matches RetentionDays", "retention_days", "RetentionDays", true}, + + // Simple case-insensitive matches (no underscores) + {"port matches Port", "port", "Port", true}, + {"enabled matches Enabled", "enabled", "Enabled", true}, + {"type matches Type", "type", "Type", true}, + {"url matches URL", "url", "URL", true}, + {"ttl matches TTL", "ttl", "TTL", true}, + {"redis matches Redis", "redis", "Redis", true}, + {"sqlite matches SQLite", "sqlite", "SQLite", true}, + {"postgresql matches PostgreSQL", "postgresql", "PostgreSQL", true}, + {"mongodb matches MongoDB", "mongodb", "MongoDB", true}, + + // Case variations + {"PORT matches Port", "PORT", "Port", true}, + {"Port matches Port", "Port", "Port", true}, + {"BODY_SIZE_LIMIT matches BodySizeLimit", "BODY_SIZE_LIMIT", "BodySizeLimit", true}, + + // Non-matches + {"different names don't match", "foo", "Bar", false}, + {"partial match fails", "body_size", "BodySizeLimit", false}, + + // Malformed keys are rejected + {"consecutive underscores rejected", "body__size_limit", "BodySizeLimit", false}, + {"leading underscore rejected", "_port", "Port", false}, + {"trailing underscore rejected", "port_", "Port", false}, + {"leading and trailing underscore rejected", "_port_", "Port", false}, + } + + // Get the MatchName function from snakeCaseMatchName + var decoderConfig mapstructure.DecoderConfig + opt := snakeCaseMatchName() + opt(&decoderConfig) + matchName := decoderConfig.MatchName + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := matchName(tt.mapKey, tt.fieldName) + if result != tt.expected { + t.Errorf("matchName(%q, %q) = %v, expected %v", + tt.mapKey, tt.fieldName, result, tt.expected) + } + }) + } +} + +func TestSnakeCaseMatchNameWithViper(t *testing.T) { + // Reset viper state + viper.Reset() + + // Create a map simulating YAML config with snake_case keys + configData := map[string]any{ + "server": map[string]any{ + "port": "9090", + "master_key": "test-master-key", + "body_size_limit": "50M", + }, + "logging": map[string]any{ + "enabled": true, + "storage_type": "postgresql", + "log_bodies": false, + "log_headers": true, + "buffer_size": 500, + "flush_interval": 10, + "retention_days": 60, + "only_model_interactions": false, + }, + "cache": map[string]any{ + "type": "redis", + "redis": map[string]any{ + "url": "redis://localhost:6379", + "key": "test:models", + "ttl": 3600, + }, + }, + } + + // Set the config data in viper + for k, v := range configData { + viper.Set(k, v) + } + + var cfg Config + err := viper.Unmarshal(&cfg, snakeCaseMatchName()) + if err != nil { + t.Fatalf("Unmarshal failed: %v", err) + } + + // Verify server config + if cfg.Server.Port != "9090" { + t.Errorf("expected Server.Port=9090, got %s", cfg.Server.Port) + } + if cfg.Server.MasterKey != "test-master-key" { + t.Errorf("expected Server.MasterKey=test-master-key, got %s", cfg.Server.MasterKey) + } + if cfg.Server.BodySizeLimit != "50M" { + t.Errorf("expected Server.BodySizeLimit=50M, got %s", cfg.Server.BodySizeLimit) + } + + // Verify logging config + if !cfg.Logging.Enabled { + t.Error("expected Logging.Enabled=true") + } + if cfg.Logging.StorageType != "postgresql" { + t.Errorf("expected Logging.StorageType=postgresql, got %s", cfg.Logging.StorageType) + } + if cfg.Logging.LogBodies { + t.Error("expected Logging.LogBodies=false") + } + if !cfg.Logging.LogHeaders { + t.Error("expected Logging.LogHeaders=true") + } + if cfg.Logging.BufferSize != 500 { + t.Errorf("expected Logging.BufferSize=500, got %d", cfg.Logging.BufferSize) + } + if cfg.Logging.FlushInterval != 10 { + t.Errorf("expected Logging.FlushInterval=10, got %d", cfg.Logging.FlushInterval) + } + if cfg.Logging.RetentionDays != 60 { + t.Errorf("expected Logging.RetentionDays=60, got %d", cfg.Logging.RetentionDays) + } + if cfg.Logging.OnlyModelInteractions { + t.Error("expected Logging.OnlyModelInteractions=false") + } + + // Verify cache config + if cfg.Cache.Type != "redis" { + t.Errorf("expected Cache.Type=redis, got %s", cfg.Cache.Type) + } + if cfg.Cache.Redis.URL != "redis://localhost:6379" { + t.Errorf("expected Cache.Redis.URL=redis://localhost:6379, got %s", cfg.Cache.Redis.URL) + } + if cfg.Cache.Redis.Key != "test:models" { + t.Errorf("expected Cache.Redis.Key=test:models, got %s", cfg.Cache.Redis.Key) + } + if cfg.Cache.Redis.TTL != 3600 { + t.Errorf("expected Cache.Redis.TTL=3600, got %d", cfg.Cache.Redis.TTL) + } +} + func TestValidateBodySizeLimit(t *testing.T) { tests := []struct { name string diff --git a/go.mod b/go.mod index c0741053..a418c1a2 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.24.0 require ( github.com/andybalholm/brotli v1.2.0 + github.com/go-viper/mapstructure/v2 v2.4.0 github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.8.0 github.com/joho/godotenv v1.5.1 @@ -23,7 +24,6 @@ require ( github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect - github.com/go-viper/mapstructure/v2 v2.4.0 // indirect github.com/golang/snappy v1.0.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect diff --git a/internal/app/app.go b/internal/app/app.go new file mode 100644 index 00000000..c41712f4 --- /dev/null +++ b/internal/app/app.go @@ -0,0 +1,218 @@ +// Package app provides the main application struct for centralized dependency management +// and lifecycle control of the GOModel server. +package app + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/http" + "sync" + "time" + + "gomodel/config" + "gomodel/internal/auditlog" + "gomodel/internal/core" + "gomodel/internal/providers" + "gomodel/internal/server" +) + +// App represents the main application with all its dependencies. +// It provides centralized lifecycle management for all components. +type App struct { + config *config.Config + providers *providers.InitResult + audit *auditlog.Result + server *server.Server + + shutdownMu sync.Mutex + shutdown bool +} + +// Config holds the configuration options for creating an App. +type Config struct { + // AppConfig is the main application configuration. + AppConfig *config.Config + + // RefreshInterval is how often to refresh the model registry. + // Default: 1 hour + RefreshInterval time.Duration + + // Factory is the provider factory with registered providers. + // Hooks should be set on the factory before passing it here. + Factory *providers.ProviderFactory +} + +// New creates a new App with all dependencies initialized. +// The caller must call Shutdown to release resources. +func New(ctx context.Context, cfg Config) (*App, error) { + if cfg.AppConfig == nil { + return nil, fmt.Errorf("app config is required") + } + if cfg.Factory == nil { + return nil, fmt.Errorf("factory is required") + } + + app := &App{ + config: cfg.AppConfig, + } + + // Initialize provider infrastructure + // RefreshInterval default (1 hour) is applied in providers.InitWithConfig if zero + initCfg := providers.InitConfig{ + RefreshInterval: cfg.RefreshInterval, + Factory: cfg.Factory, + } + + providerResult, err := providers.InitWithConfig(ctx, cfg.AppConfig, initCfg) + if err != nil { + return nil, fmt.Errorf("failed to initialize providers: %w", err) + } + app.providers = providerResult + + // Initialize audit logging + auditResult, err := auditlog.New(ctx, cfg.AppConfig) + if err != nil { + closeErr := app.providers.Close() + if closeErr != nil { + return nil, fmt.Errorf("failed to initialize audit logging: %w (also: providers close error: %v)", err, closeErr) + } + return nil, fmt.Errorf("failed to initialize audit logging: %w", err) + } + app.audit = auditResult + + // Log configuration status + app.logStartupInfo() + + // Create server + serverCfg := &server.Config{ + MasterKey: cfg.AppConfig.Server.MasterKey, + MetricsEnabled: cfg.AppConfig.Metrics.Enabled, + MetricsEndpoint: cfg.AppConfig.Metrics.Endpoint, + BodySizeLimit: cfg.AppConfig.Server.BodySizeLimit, + AuditLogger: auditResult.Logger, + LogOnlyModelInteractions: cfg.AppConfig.Logging.OnlyModelInteractions, + } + app.server = server.New(app.providers.Router, serverCfg) + + return app, nil +} + +// Router returns the core.RoutableProvider for request routing. +func (a *App) Router() core.RoutableProvider { + if a.providers == nil { + return nil + } + return a.providers.Router +} + +// AuditLogger returns the audit logger interface. +func (a *App) AuditLogger() auditlog.LoggerInterface { + if a.audit == nil { + return nil + } + return a.audit.Logger +} + +// Start starts the HTTP server on the given address. +// This is a blocking call that returns when the server stops. +func (a *App) Start(addr string) error { + if a.server == nil { + return fmt.Errorf("server is not initialized") + } + slog.Info("starting server", "address", addr) + if err := a.server.Start(addr); err != nil { + if errors.Is(err, http.ErrServerClosed) { + slog.Info("server stopped gracefully") + return nil + } + return fmt.Errorf("server failed to start: %w", err) + } + return nil +} + +// Shutdown gracefully shuts down all components in the correct order. +// It ensures proper cleanup of resources: +// 1. HTTP server (stop accepting new requests) +// 2. Background refresh goroutine and cache +// 3. Audit logging +// +// Safe to call multiple times; subsequent calls are no-ops. +func (a *App) Shutdown(ctx context.Context) error { + a.shutdownMu.Lock() + if a.shutdown { + a.shutdownMu.Unlock() + return nil + } + a.shutdown = true + a.shutdownMu.Unlock() + + slog.Info("shutting down application...") + + var errs []error + + // 1. Shutdown HTTP server first (stop accepting new requests) + if a.server != nil { + if err := a.server.Shutdown(ctx); err != nil { + slog.Error("server shutdown error", "error", err) + errs = append(errs, fmt.Errorf("server shutdown: %w", err)) + } + } + + // 2. Close providers (stops background refresh and cache) + if a.providers != nil { + if err := a.providers.Close(); err != nil { + slog.Error("providers close error", "error", err) + errs = append(errs, fmt.Errorf("providers close: %w", err)) + } + } + + // 3. Close audit logging (flushes pending logs) + if a.audit != nil { + if err := a.audit.Close(); err != nil { + slog.Error("audit logger close error", "error", err) + errs = append(errs, fmt.Errorf("audit close: %w", err)) + } + } + + if len(errs) > 0 { + return fmt.Errorf("shutdown errors: %w", errors.Join(errs...)) + } + + slog.Info("application shutdown complete") + return nil +} + +// logStartupInfo logs the application configuration on startup. +func (a *App) logStartupInfo() { + cfg := a.config + + // Security warnings + if cfg.Server.MasterKey == "" { + slog.Warn("SECURITY WARNING: GOMODEL_MASTER_KEY not set - server running in UNSAFE MODE", + "security_risk", "unauthenticated access allowed", + "recommendation", "set GOMODEL_MASTER_KEY environment variable to secure this gateway") + } else { + slog.Info("authentication enabled", "mode", "master_key") + } + + // Metrics configuration + if cfg.Metrics.Enabled { + slog.Info("prometheus metrics enabled", "endpoint", cfg.Metrics.Endpoint) + } else { + slog.Info("prometheus metrics disabled") + } + + // Audit logging configuration + if cfg.Logging.Enabled { + slog.Info("audit logging enabled", + "storage_type", cfg.Logging.StorageType, + "log_bodies", cfg.Logging.LogBodies, + "log_headers", cfg.Logging.LogHeaders, + "retention_days", cfg.Logging.RetentionDays, + ) + } else { + slog.Info("audit logging disabled") + } +} diff --git a/internal/auditlog/factory.go b/internal/auditlog/factory.go index a88618c1..127444f6 100644 --- a/internal/auditlog/factory.go +++ b/internal/auditlog/factory.go @@ -2,6 +2,7 @@ package auditlog import ( "context" + "errors" "fmt" "time" @@ -34,7 +35,7 @@ func (r *Result) Close() error { } } if len(errs) > 0 { - return fmt.Errorf("close errors: %v", errs) + return fmt.Errorf("close errors: %w", errors.Join(errs...)) } return nil } diff --git a/internal/core/interfaces.go b/internal/core/interfaces.go index eb05fc82..cf8c7a1b 100644 --- a/internal/core/interfaces.go +++ b/internal/core/interfaces.go @@ -37,3 +37,23 @@ type RoutableProvider interface { // Returns empty string if the model is not found. GetProviderType(model string) string } + +// ModelLookup defines the interface for looking up models and their providers. +// This abstraction allows the Router to be decoupled from the concrete ModelRegistry implementation. +type ModelLookup interface { + // Supports returns true if the registry has a provider for the given model + Supports(model string) bool + + // GetProvider returns the provider for the given model, or nil if not found + GetProvider(model string) Provider + + // GetProviderType returns the provider type string for the given model. + // Returns empty string if the model is not found. + GetProviderType(model string) string + + // ListModels returns all models in the registry + ListModels() []Model + + // ModelCount returns the number of registered models + ModelCount() int +} diff --git a/internal/providers/anthropic/anthropic.go b/internal/providers/anthropic/anthropic.go index 3bf13b83..77da3eb4 100644 --- a/internal/providers/anthropic/anthropic.go +++ b/internal/providers/anthropic/anthropic.go @@ -20,38 +20,41 @@ import ( "gomodel/internal/providers" ) +// Registration provides factory registration for the Anthropic provider. +var Registration = providers.Registration{ + Type: "anthropic", + New: New, +} + const ( defaultBaseURL = "https://api.anthropic.com/v1" anthropicAPIVersion = "2023-06-01" ) -func init() { - // Self-register with the factory - providers.RegisterProvider("anthropic", New) -} - // Provider implements the core.Provider interface for Anthropic type Provider struct { client *llmclient.Client apiKey string } -// New creates a new Anthropic provider -func New(apiKey string) *Provider { +// New creates a new Anthropic provider. +func New(apiKey string, hooks llmclient.Hooks) core.Provider { p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("anthropic", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.New(cfg, p.setHeaders) return p } -// NewWithHTTPClient creates a new Anthropic provider with a custom HTTP client -func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { +// NewWithHTTPClient creates a new Anthropic provider with a custom HTTP client. +// If httpClient is nil, http.DefaultClient is used. +func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { + if httpClient == nil { + httpClient = http.DefaultClient + } p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("anthropic", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) return p } diff --git a/internal/providers/anthropic/anthropic_test.go b/internal/providers/anthropic/anthropic_test.go index 0305e0e2..3872478e 100644 --- a/internal/providers/anthropic/anthropic_test.go +++ b/internal/providers/anthropic/anthropic_test.go @@ -10,11 +10,13 @@ import ( "testing" "gomodel/internal/core" + "gomodel/internal/llmclient" ) func TestNew(t *testing.T) { apiKey := "test-api-key" - provider := New(apiKey) + // Use NewWithHTTPClient to get concrete type for internal testing + provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -24,6 +26,15 @@ func TestNew(t *testing.T) { } } +func TestNew_ReturnsProvider(t *testing.T) { + apiKey := "test-api-key" + provider := New(apiKey, llmclient.Hooks{}) + + if provider == nil { + t.Error("provider should not be nil") + } +} + func TestChatCompletion(t *testing.T) { tests := []struct { name string @@ -136,7 +147,7 @@ func TestChatCompletion(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -263,7 +274,7 @@ data: {"type":"message_stop"} })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -292,7 +303,7 @@ data: {"type":"message_stop"} } func TestListModels(t *testing.T) { - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) resp, err := provider.ListModels(context.Background()) @@ -355,7 +366,7 @@ func TestChatCompletionWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -621,7 +632,7 @@ func TestResponses(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -690,7 +701,7 @@ func TestResponsesWithArrayInput(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -754,7 +765,7 @@ func TestResponsesWithInstructions(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -871,7 +882,7 @@ data: {"type":"message_stop"} })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -905,7 +916,7 @@ func TestResponsesWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) diff --git a/internal/providers/factory.go b/internal/providers/factory.go index 02a21c11..6addc65d 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -3,66 +3,84 @@ package providers import ( "fmt" + "sync" "gomodel/config" "gomodel/internal/core" "gomodel/internal/llmclient" ) -// Builder creates a provider instance from configuration -type Builder func(cfg config.ProviderConfig) (core.Provider, error) +// NewFunc is the constructor signature for providers. +type NewFunc func(apiKey string, hooks llmclient.Hooks) core.Provider -// registry holds all registered provider builders -var registry = make(map[string]Builder) +// Registration contains metadata for registering a provider with the factory. +type Registration struct { + // Type is the provider identifier (e.g., "openai", "anthropic") + Type string -// globalHooks holds the observability hooks to inject into all providers -var globalHooks llmclient.Hooks + // New creates a new provider instance + New NewFunc +} -// Register allows provider packages to register themselves -// This should be called from init() functions in provider packages -func Register(providerType string, builder Builder) { - registry[providerType] = builder +// ProviderFactory manages provider registration and creation. +type ProviderFactory struct { + mu sync.RWMutex + builders map[string]NewFunc + hooks llmclient.Hooks } -// RegisterProvider registers a provider constructor with base URL support -func RegisterProvider[T core.Provider](providerType string, newProvider func(string) T) { - Register(providerType, func(cfg config.ProviderConfig) (core.Provider, error) { - p := newProvider(cfg.APIKey) - if cfg.BaseURL != "" { - if setter, ok := any(p).(interface{ SetBaseURL(string) }); ok { - setter.SetBaseURL(cfg.BaseURL) - } - } - return p, nil - }) +// NewProviderFactory creates a new provider factory instance. +func NewProviderFactory() *ProviderFactory { + return &ProviderFactory{ + builders: make(map[string]NewFunc), + } +} + +// SetHooks configures observability hooks for all providers created by this factory. +func (f *ProviderFactory) SetHooks(hooks llmclient.Hooks) { + f.mu.Lock() + defer f.mu.Unlock() + f.hooks = hooks } -// Create instantiates a provider based on configuration -func Create(cfg config.ProviderConfig) (core.Provider, error) { - builder, ok := registry[cfg.Type] +// Register adds a provider to the factory. +func (f *ProviderFactory) Register(reg Registration) { + f.mu.Lock() + defer f.mu.Unlock() + f.builders[reg.Type] = reg.New +} + +// Create instantiates a provider based on configuration. +func (f *ProviderFactory) Create(cfg config.ProviderConfig) (core.Provider, error) { + f.mu.RLock() + builder, ok := f.builders[cfg.Type] + hooks := f.hooks + f.mu.RUnlock() + if !ok { return nil, fmt.Errorf("unknown provider type: %s", cfg.Type) } - return builder(cfg) -} -// ListRegistered returns a list of all registered provider types -func ListRegistered() []string { - types := make([]string, 0, len(registry)) - for t := range registry { - types = append(types, t) + p := builder(cfg.APIKey, hooks) + + // Set custom base URL if configured + if cfg.BaseURL != "" { + if setter, ok := p.(interface{ SetBaseURL(string) }); ok { + setter.SetBaseURL(cfg.BaseURL) + } } - return types -} -// SetGlobalHooks configures observability hooks that will be injected into all providers. -// This must be called before Create() to take effect. -// This enables metrics, tracing, and logging without modifying provider implementations. -func SetGlobalHooks(hooks llmclient.Hooks) { - globalHooks = hooks + return p, nil } -// GetGlobalHooks returns the currently configured global hooks -func GetGlobalHooks() llmclient.Hooks { - return globalHooks +// RegisteredTypes returns a list of all registered provider types. +func (f *ProviderFactory) RegisteredTypes() []string { + f.mu.RLock() + defer f.mu.RUnlock() + + types := make([]string, 0, len(f.builders)) + for t := range f.builders { + types = append(types, t) + } + return types } diff --git a/internal/providers/factory_test.go b/internal/providers/factory_test.go index 28bf0b63..ab84d57a 100644 --- a/internal/providers/factory_test.go +++ b/internal/providers/factory_test.go @@ -7,6 +7,7 @@ import ( "gomodel/config" "gomodel/internal/core" + "gomodel/internal/llmclient" ) // factoryMockProvider is a test implementation of core.Provider @@ -41,44 +42,35 @@ func (m *factoryMockProvider) StreamResponses(ctx context.Context, req *core.Res return nil, nil } -func TestRegister(t *testing.T) { - // Save current registry and restore after test - originalRegistry := registry - defer func() { registry = originalRegistry }() - - // Create a fresh registry for testing - registry = make(map[string]Builder) +func TestProviderFactory_Register(t *testing.T) { + factory := NewProviderFactory() // Test registering a new provider type - mockBuilder := func(cfg config.ProviderConfig) (core.Provider, error) { - return nil, nil - } - - Register("test-provider", mockBuilder) + factory.Register(Registration{ + Type: "test-provider", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + return &factoryMockProvider{} + }, + }) - if _, exists := registry["test-provider"]; !exists { - t.Error("expected 'test-provider' to be registered") + registered := factory.RegisteredTypes() + if len(registered) != 1 { + t.Errorf("expected 1 registered provider, got %d", len(registered)) } - - if len(registry) != 1 { - t.Errorf("expected registry to have 1 entry, got %d", len(registry)) + if registered[0] != "test-provider" { + t.Errorf("expected 'test-provider', got %q", registered[0]) } } -func TestCreate_UnknownType(t *testing.T) { - // Save current registry and restore after test - originalRegistry := registry - defer func() { registry = originalRegistry }() - - // Create a fresh registry for testing - registry = make(map[string]Builder) +func TestProviderFactory_Create_UnknownType(t *testing.T) { + factory := NewProviderFactory() cfg := config.ProviderConfig{ Type: "unknown-type", APIKey: "test-key", } - _, err := Create(cfg) + _, err := factory.Create(cfg) if err == nil { t.Error("expected error for unknown provider type, got nil") } @@ -89,17 +81,15 @@ func TestCreate_UnknownType(t *testing.T) { } } -func TestCreate_Success(t *testing.T) { - // Save current registry and restore after test - originalRegistry := registry - defer func() { registry = originalRegistry }() - - // Create a fresh registry for testing - registry = make(map[string]Builder) +func TestProviderFactory_Create_Success(t *testing.T) { + factory := NewProviderFactory() // Register a mock builder - Register("mock", func(cfg config.ProviderConfig) (core.Provider, error) { - return &factoryMockProvider{}, nil + factory.Register(Registration{ + Type: "mock", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + return &factoryMockProvider{} + }, }) cfg := config.ProviderConfig{ @@ -107,7 +97,7 @@ func TestCreate_Success(t *testing.T) { APIKey: "test-key", } - provider, err := Create(cfg) + provider, err := factory.Create(cfg) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -117,20 +107,20 @@ func TestCreate_Success(t *testing.T) { } } -func TestListRegistered(t *testing.T) { - // Save current registry and restore after test - originalRegistry := registry - defer func() { registry = originalRegistry }() - - // Create a fresh registry for testing - registry = make(map[string]Builder) +func TestProviderFactory_RegisteredTypes(t *testing.T) { + factory := NewProviderFactory() // Register some test providers - Register("provider1", func(cfg config.ProviderConfig) (core.Provider, error) { return nil, nil }) - Register("provider2", func(cfg config.ProviderConfig) (core.Provider, error) { return nil, nil }) - Register("provider3", func(cfg config.ProviderConfig) (core.Provider, error) { return nil, nil }) + for _, name := range []string{"provider1", "provider2", "provider3"} { + factory.Register(Registration{ + Type: name, + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + return &factoryMockProvider{} + }, + }) + } - registered := ListRegistered() + registered := factory.RegisteredTypes() if len(registered) != 3 { t.Errorf("expected 3 registered providers, got %d", len(registered)) @@ -150,25 +140,23 @@ func TestListRegistered(t *testing.T) { } } -func TestCreate_WithBaseURL(t *testing.T) { - // This test verifies that the factory pattern allows providers - // to use custom base URLs from configuration - // (Actual provider implementations are tested in their own packages) - - // Save current registry and restore after test - originalRegistry := registry - defer func() { registry = originalRegistry }() - - // Create a fresh registry for testing - registry = make(map[string]Builder) +func TestProviderFactory_Create_WithBaseURL(t *testing.T) { + factory := NewProviderFactory() customBaseURL := "https://custom.api.endpoint.com/v1" - var capturedBaseURL string - // Register a mock builder that captures the base URL - Register("custom", func(cfg config.ProviderConfig) (core.Provider, error) { - capturedBaseURL = cfg.BaseURL - return nil, nil + // Create a mock provider + type mockWithBaseURL struct { + factoryMockProvider + } + mockProvider := &mockWithBaseURL{} + + // Register a mock builder + factory.Register(Registration{ + Type: "custom", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + return mockProvider + }, }) cfg := config.ProviderConfig{ @@ -177,12 +165,117 @@ func TestCreate_WithBaseURL(t *testing.T) { BaseURL: customBaseURL, } - _, err := Create(cfg) + // The factory only calls SetBaseURL if the provider implements it. + // Our mock doesn't implement it, so we're just testing that Create succeeds. + provider, err := factory.Create(cfg) if err != nil { t.Errorf("unexpected error: %v", err) } - if capturedBaseURL != customBaseURL { - t.Errorf("expected base URL '%s', got '%s'", customBaseURL, capturedBaseURL) + if provider == nil { + t.Error("expected provider to be created, got nil") + } +} + +func TestProviderFactory_SetHooks(t *testing.T) { + factory := NewProviderFactory() + + // Create mock hooks with identifiable callbacks + mockHooks := llmclient.Hooks{ + OnRequestStart: func(ctx context.Context, info llmclient.RequestInfo) context.Context { + return ctx + }, + } + factory.SetHooks(mockHooks) + + // Verify hooks were set by creating a provider and checking it received them + var receivedHooks llmclient.Hooks + factory.Register(Registration{ + Type: "test", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + receivedHooks = hooks + return &factoryMockProvider{} + }, + }) + + cfg := config.ProviderConfig{ + Type: "test", + APIKey: "test-key", + } + + _, err := factory.Create(cfg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Verify hooks were passed by checking callback exists + if receivedHooks.OnRequestStart == nil { + t.Error("expected hooks to be passed to builder") + } +} + +func TestProviderFactory_HooksPassedToBuilder(t *testing.T) { + factory := NewProviderFactory() + + // Create mock hooks + mockHooks := llmclient.Hooks{ + OnRequestStart: func(ctx context.Context, info llmclient.RequestInfo) context.Context { + return ctx + }, + } + factory.SetHooks(mockHooks) + + var receivedHooks llmclient.Hooks + + factory.Register(Registration{ + Type: "test", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + receivedHooks = hooks + return &factoryMockProvider{} + }, + }) + + cfg := config.ProviderConfig{ + Type: "test", + APIKey: "test-key", + } + + _, err := factory.Create(cfg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Verify hooks were passed by checking callback exists + if receivedHooks.OnRequestStart == nil { + t.Error("expected hooks to be passed to builder") + } +} + +func TestProviderFactory_ZeroHooks(t *testing.T) { + factory := NewProviderFactory() + + var receivedHooks llmclient.Hooks + + factory.Register(Registration{ + Type: "test", + New: func(apiKey string, hooks llmclient.Hooks) core.Provider { + receivedHooks = hooks + return &factoryMockProvider{} + }, + }) + + cfg := config.ProviderConfig{ + Type: "test", + APIKey: "test-key", + } + + _, err := factory.Create(cfg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Without SetHooks being called, hooks should be zero value (empty callbacks) + if receivedHooks.OnRequestStart != nil || receivedHooks.OnRequestEnd != nil { + t.Error("expected zero hooks when SetHooks not called") } } diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index b9cfa52c..82fd375b 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -15,6 +15,12 @@ import ( "gomodel/internal/providers" ) +// Registration provides factory registration for the Gemini provider. +var Registration = providers.Registration{ + Type: "gemini", + New: New, +} + const ( // Gemini provides an OpenAI-compatible endpoint defaultOpenAICompatibleBaseURL = "https://generativelanguage.googleapis.com/v1beta/openai" @@ -22,40 +28,40 @@ const ( defaultModelsBaseURL = "https://generativelanguage.googleapis.com/v1beta" ) -func init() { - // Self-register with the factory - providers.RegisterProvider("gemini", New) -} - // Provider implements the core.Provider interface for Google Gemini type Provider struct { client *llmclient.Client + hooks llmclient.Hooks apiKey string modelsURL string } -// New creates a new Gemini provider -func New(apiKey string) *Provider { +// New creates a new Gemini provider. +func New(apiKey string, hooks llmclient.Hooks) core.Provider { p := &Provider{ apiKey: apiKey, + hooks: hooks, modelsURL: defaultModelsBaseURL, } cfg := llmclient.DefaultConfig("gemini", defaultOpenAICompatibleBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.New(cfg, p.setHeaders) return p } -// NewWithHTTPClient creates a new Gemini provider with a custom HTTP client -func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { +// NewWithHTTPClient creates a new Gemini provider with a custom HTTP client. +// If httpClient is nil, http.DefaultClient is used. +func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { + if httpClient == nil { + httpClient = http.DefaultClient + } p := &Provider{ apiKey: apiKey, + hooks: hooks, modelsURL: defaultModelsBaseURL, } cfg := llmclient.DefaultConfig("gemini", defaultOpenAICompatibleBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) return p } @@ -126,8 +132,7 @@ func (p *Provider) ListModels(ctx context.Context) (*core.ModelsResponse, error) // Use the native Gemini API to list models // We need to create a separate client for the models endpoint since it uses a different URL modelsCfg := llmclient.DefaultConfig("gemini", p.modelsURL) - // Apply global hooks if available - modelsCfg.Hooks = providers.GetGlobalHooks() + modelsCfg.Hooks = p.hooks modelsClient := llmclient.New( modelsCfg, func(req *http.Request) { diff --git a/internal/providers/gemini/gemini_test.go b/internal/providers/gemini/gemini_test.go index a0fb01b0..8be94a02 100644 --- a/internal/providers/gemini/gemini_test.go +++ b/internal/providers/gemini/gemini_test.go @@ -10,11 +10,13 @@ import ( "testing" "gomodel/internal/core" + "gomodel/internal/llmclient" ) func TestNew(t *testing.T) { apiKey := "test-api-key" - provider := New(apiKey) + // Use NewWithHTTPClient to get concrete type for internal testing + provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -27,6 +29,14 @@ func TestNew(t *testing.T) { } } +func TestNew_ReturnsProvider(t *testing.T) { + provider := New("test-api-key", llmclient.Hooks{}) + + if provider == nil { + t.Error("provider should not be nil") + } +} + func TestChatCompletion(t *testing.T) { tests := []struct { name string @@ -127,7 +137,7 @@ func TestChatCompletion(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -209,7 +219,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -332,7 +342,7 @@ func TestListModels(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.modelsURL = server.URL resp, err := provider.ListModels(context.Background()) @@ -360,7 +370,7 @@ func TestChatCompletionWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -404,7 +414,7 @@ func TestResponses(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -438,7 +448,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ diff --git a/internal/providers/groq/groq.go b/internal/providers/groq/groq.go index a8f5c970..48a3b79c 100644 --- a/internal/providers/groq/groq.go +++ b/internal/providers/groq/groq.go @@ -14,37 +14,40 @@ import ( "gomodel/internal/providers" ) +// Registration provides factory registration for the Groq provider. +var Registration = providers.Registration{ + Type: "groq", + New: New, +} + const ( defaultBaseURL = "https://api.groq.com/openai/v1" ) -func init() { - // Self-register with the factory - providers.RegisterProvider("groq", New) -} - // Provider implements the core.Provider interface for Groq type Provider struct { client *llmclient.Client apiKey string } -// New creates a new Groq provider -func New(apiKey string) *Provider { +// New creates a new Groq provider. +func New(apiKey string, hooks llmclient.Hooks) core.Provider { p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("groq", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.New(cfg, p.setHeaders) return p } -// NewWithHTTPClient creates a new Groq provider with a custom HTTP client -func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { +// NewWithHTTPClient creates a new Groq provider with a custom HTTP client. +// If httpClient is nil, http.DefaultClient is used. +func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { + if httpClient == nil { + httpClient = http.DefaultClient + } p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("groq", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) return p } diff --git a/internal/providers/groq/groq_test.go b/internal/providers/groq/groq_test.go index 8e69592a..03dd5868 100644 --- a/internal/providers/groq/groq_test.go +++ b/internal/providers/groq/groq_test.go @@ -10,12 +10,14 @@ import ( "testing" "gomodel/internal/core" + "gomodel/internal/llmclient" "gomodel/internal/providers" ) func TestNew(t *testing.T) { apiKey := "test-api-key" - provider := New(apiKey) + // Use NewWithHTTPClient to get concrete type for internal testing + provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -25,6 +27,14 @@ func TestNew(t *testing.T) { } } +func TestNew_ReturnsProvider(t *testing.T) { + provider := New("test-api-key", llmclient.Hooks{}) + + if provider == nil { + t.Error("provider should not be nil") + } +} + func TestChatCompletion(t *testing.T) { tests := []struct { name string @@ -127,7 +137,7 @@ func TestChatCompletion(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -211,7 +221,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -323,7 +333,7 @@ func TestListModels(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) resp, err := provider.ListModels(context.Background()) @@ -352,7 +362,7 @@ func TestChatCompletionWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -401,7 +411,7 @@ func TestResponses(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -492,7 +502,7 @@ func TestResponsesWithArrayInput(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -545,7 +555,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -587,7 +597,7 @@ func TestResponsesWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -973,7 +983,7 @@ func TestNewWithHTTPClient(t *testing.T) { customClient := &http.Client{} apiKey := "test-api-key" - provider := NewWithHTTPClient(apiKey, customClient) + provider := NewWithHTTPClient(apiKey, customClient, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -984,7 +994,7 @@ func TestNewWithHTTPClient(t *testing.T) { } func TestSetBaseURL(t *testing.T) { - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) customURL := "https://custom.groq.api.com/v1" provider.SetBaseURL(customURL) diff --git a/internal/providers/init.go b/internal/providers/init.go index d5428730..3a3cb546 100644 --- a/internal/providers/init.go +++ b/internal/providers/init.go @@ -18,6 +18,7 @@ type InitResult struct { Registry *ModelRegistry Router *Router Cache cache.Cache + Factory *ProviderFactory // stopRefresh is called to stop the background refresh goroutine stopRefresh func() @@ -39,14 +40,18 @@ func (r *InitResult) Close() error { // InitConfig holds options for provider initialization. type InitConfig struct { // RefreshInterval is how often to refresh the model registry. - // Default: 5 minutes + // Default: 1 hour RefreshInterval time.Duration + + // Factory is the provider factory with registered providers. + // Hooks should be set on the factory before passing it here. + Factory *ProviderFactory } // DefaultInitConfig returns sensible defaults for initialization. func DefaultInitConfig() InitConfig { return InitConfig{ - RefreshInterval: 5 * time.Minute, + RefreshInterval: time.Hour, } } @@ -67,18 +72,25 @@ func Init(ctx context.Context, cfg *config.Config) (*InitResult, error) { // InitWithConfig initializes the provider infrastructure with custom options. func InitWithConfig(ctx context.Context, cfg *config.Config, initCfg InitConfig) (*InitResult, error) { + // Validate that factory is provided + if initCfg.Factory == nil { + return nil, fmt.Errorf("InitConfig.Factory is required") + } + // Initialize cache backend based on configuration modelCache, err := initCache(cfg) if err != nil { return nil, fmt.Errorf("failed to initialize cache: %w", err) } + factory := initCfg.Factory + // Create model registry with cache registry := NewModelRegistry() registry.SetCache(modelCache) - // Register providers - count, err := registerProviders(cfg, registry) + // Register providers using the factory + count, err := registerProviders(cfg, factory, registry) if err != nil { modelCache.Close() return nil, err @@ -100,7 +112,7 @@ func InitWithConfig(ctx context.Context, cfg *config.Config, initCfg InitConfig) // Start background refresh interval := initCfg.RefreshInterval if interval <= 0 { - interval = 5 * time.Minute + interval = time.Hour } stopRefresh := registry.StartBackgroundRefresh(interval) @@ -116,6 +128,7 @@ func InitWithConfig(ctx context.Context, cfg *config.Config, initCfg InitConfig) Registry: registry, Router: router, Cache: modelCache, + Factory: factory, stopRefresh: stopRefresh, }, nil } @@ -166,7 +179,7 @@ func initCache(cfg *config.Config) (cache.Cache, error) { // registerProviders creates and registers all configured providers. // Returns the count of successfully initialized providers. -func registerProviders(cfg *config.Config, registry *ModelRegistry) (int, error) { +func registerProviders(cfg *config.Config, factory *ProviderFactory, registry *ModelRegistry) (int, error) { // Sort provider names for deterministic initialization order providerNames := make([]string, 0, len(cfg.Providers)) for name := range cfg.Providers { @@ -177,7 +190,7 @@ func registerProviders(cfg *config.Config, registry *ModelRegistry) (int, error) var initializedCount int for _, name := range providerNames { pCfg := cfg.Providers[name] - p, err := Create(pCfg) + p, err := factory.Create(pCfg) if err != nil { slog.Error("failed to initialize provider", "name", name, diff --git a/internal/providers/openai/openai.go b/internal/providers/openai/openai.go index c1481c03..c9eaf5de 100644 --- a/internal/providers/openai/openai.go +++ b/internal/providers/openai/openai.go @@ -11,37 +11,40 @@ import ( "gomodel/internal/providers" ) +// Registration provides factory registration for the OpenAI provider. +var Registration = providers.Registration{ + Type: "openai", + New: New, +} + const ( defaultBaseURL = "https://api.openai.com/v1" ) -func init() { - // Self-register with the factory - providers.RegisterProvider("openai", New) -} - // Provider implements the core.Provider interface for OpenAI type Provider struct { client *llmclient.Client apiKey string } -// New creates a new OpenAI provider -func New(apiKey string) *Provider { +// New creates a new OpenAI provider. +func New(apiKey string, hooks llmclient.Hooks) core.Provider { p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("openai", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.New(cfg, p.setHeaders) return p } -// NewWithHTTPClient creates a new OpenAI provider with a custom HTTP client -func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { +// NewWithHTTPClient creates a new OpenAI provider with a custom HTTP client. +// If httpClient is nil, http.DefaultClient is used. +func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { + if httpClient == nil { + httpClient = http.DefaultClient + } p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("openai", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) return p } diff --git a/internal/providers/openai/openai_test.go b/internal/providers/openai/openai_test.go index 336b8a59..1548ca9e 100644 --- a/internal/providers/openai/openai_test.go +++ b/internal/providers/openai/openai_test.go @@ -10,11 +10,13 @@ import ( "testing" "gomodel/internal/core" + "gomodel/internal/llmclient" ) func TestNew(t *testing.T) { apiKey := "test-api-key" - provider := New(apiKey) + // Use NewWithHTTPClient to get concrete type for internal testing + provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -24,6 +26,15 @@ func TestNew(t *testing.T) { } } +func TestNew_ReturnsProvider(t *testing.T) { + apiKey := "test-api-key" + provider := New(apiKey, llmclient.Hooks{}) + + if provider == nil { + t.Error("provider should not be nil") + } +} + func TestChatCompletion(t *testing.T) { tests := []struct { name string @@ -126,7 +137,7 @@ func TestChatCompletion(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -210,7 +221,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -322,7 +333,7 @@ func TestListModels(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) resp, err := provider.ListModels(context.Background()) @@ -351,7 +362,7 @@ func TestChatCompletionWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -492,7 +503,7 @@ func TestResponses(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -561,7 +572,7 @@ func TestResponsesWithArrayInput(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -688,7 +699,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -722,7 +733,7 @@ func TestResponsesWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) diff --git a/internal/providers/registry.go b/internal/providers/registry.go index b7d6c7a7..5557e63d 100644 --- a/internal/providers/registry.go +++ b/internal/providers/registry.go @@ -286,33 +286,33 @@ func (r *ModelRegistry) IsInitialized() bool { } // GetProvider returns the provider for the given model, or nil if not found -func (r *ModelRegistry) GetProvider(modelID string) core.Provider { +func (r *ModelRegistry) GetProvider(model string) core.Provider { r.mu.RLock() defer r.mu.RUnlock() - if info, ok := r.models[modelID]; ok { + if info, ok := r.models[model]; ok { return info.Provider } return nil } -// GetModel returns the model info for the given model ID, or nil if not found -func (r *ModelRegistry) GetModel(modelID string) *ModelInfo { +// GetModel returns the model info for the given model, or nil if not found +func (r *ModelRegistry) GetModel(model string) *ModelInfo { r.mu.RLock() defer r.mu.RUnlock() - if info, ok := r.models[modelID]; ok { + if info, ok := r.models[model]; ok { return info } return nil } // Supports returns true if the registry has a provider for the given model -func (r *ModelRegistry) Supports(modelID string) bool { +func (r *ModelRegistry) Supports(model string) bool { r.mu.RLock() defer r.mu.RUnlock() - _, ok := r.models[modelID] + _, ok := r.models[model] return ok } @@ -341,13 +341,13 @@ func (r *ModelRegistry) ModelCount() int { return len(r.models) } -// GetProviderType returns the provider type string for the given model ID. +// GetProviderType returns the provider type string for the given model. // Returns empty string if the model is not found. -func (r *ModelRegistry) GetProviderType(modelID string) string { +func (r *ModelRegistry) GetProviderType(model string) string { r.mu.RLock() defer r.mu.RUnlock() - info, ok := r.models[modelID] + info, ok := r.models[model] if !ok { return "" } diff --git a/internal/providers/registry_cache_test.go b/internal/providers/registry_cache_test.go index 33b66e2b..7503424a 100644 --- a/internal/providers/registry_cache_test.go +++ b/internal/providers/registry_cache_test.go @@ -28,7 +28,7 @@ func TestCacheFile(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "openai", modelsResponse: &core.ModelsResponse{ Object: "list", @@ -103,11 +103,11 @@ func TestCacheFile(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - openaiMock := &mockProvider{ + openaiMock := ®istryMockProvider{ name: "openai", modelsResponse: &core.ModelsResponse{Object: "list"}, } - anthropicMock := &mockProvider{ + anthropicMock := ®istryMockProvider{ name: "anthropic", modelsResponse: &core.ModelsResponse{Object: "list"}, } @@ -172,7 +172,7 @@ func TestCacheFile(t *testing.T) { registry := NewModelRegistry() localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - openaiMock := &mockProvider{name: "openai"} + openaiMock := ®istryMockProvider{name: "openai"} registry.RegisterProviderWithType(openaiMock, "openai") loaded, err := registry.LoadFromCache(context.Background()) @@ -238,7 +238,7 @@ func TestCacheFile(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", modelsResponse: &core.ModelsResponse{ Object: "list", @@ -286,7 +286,7 @@ func TestInitializeAsync(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", listModelsDelay: 50 * time.Millisecond, // delay long enough for assertion to run modelsResponse: &core.ModelsResponse{ @@ -318,7 +318,7 @@ func TestInitializeAsync(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", modelsResponse: &core.ModelsResponse{ Object: "list", @@ -354,7 +354,7 @@ func TestInitializeAsync(t *testing.T) { localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", modelsResponse: &core.ModelsResponse{ Object: "list", @@ -401,7 +401,7 @@ func TestIsInitialized(t *testing.T) { t.Run("TrueAfterInitialize", func(t *testing.T) { registry := NewModelRegistry() - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", modelsResponse: &core.ModelsResponse{ Object: "list", @@ -441,7 +441,7 @@ func TestIsInitialized(t *testing.T) { registry := NewModelRegistry() localCache := cache.NewLocalCache(cacheFile) registry.SetCache(localCache) - mock := &mockProvider{name: "test"} + mock := ®istryMockProvider{name: "test"} registry.RegisterProviderWithType(mock, "test") _, _ = registry.LoadFromCache(context.Background()) @@ -456,7 +456,7 @@ func TestIsInitialized(t *testing.T) { func TestRegisterProviderWithType(t *testing.T) { registry := NewModelRegistry() - mock := &mockProvider{ + mock := ®istryMockProvider{ name: "test", modelsResponse: &core.ModelsResponse{ Object: "list", diff --git a/internal/providers/registry_test.go b/internal/providers/registry_test.go new file mode 100644 index 00000000..c3c729c4 --- /dev/null +++ b/internal/providers/registry_test.go @@ -0,0 +1,469 @@ +package providers + +import ( + "context" + "errors" + "io" + "sync/atomic" + "testing" + "time" + + "gomodel/internal/core" +) + +// registryMockProvider is a mock implementation of core.Provider for Registry testing. +// It includes all fields needed for testing the full registry lifecycle. +type registryMockProvider struct { + name string + chatResponse *core.ChatResponse + responsesResponse *core.ResponsesResponse + modelsResponse *core.ModelsResponse + err error + listModelsDelay time.Duration +} + +func (m *registryMockProvider) ChatCompletion(_ context.Context, _ *core.ChatRequest) (*core.ChatResponse, error) { + if m.err != nil { + return nil, m.err + } + return m.chatResponse, nil +} + +func (m *registryMockProvider) StreamChatCompletion(_ context.Context, _ *core.ChatRequest) (io.ReadCloser, error) { + if m.err != nil { + return nil, m.err + } + return io.NopCloser(nil), nil +} + +func (m *registryMockProvider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { + if m.listModelsDelay > 0 { + select { + case <-time.After(m.listModelsDelay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if m.err != nil { + return nil, m.err + } + return m.modelsResponse, nil +} + +func (m *registryMockProvider) Responses(_ context.Context, _ *core.ResponsesRequest) (*core.ResponsesResponse, error) { + if m.err != nil { + return nil, m.err + } + return m.responsesResponse, nil +} + +func (m *registryMockProvider) StreamResponses(_ context.Context, _ *core.ResponsesRequest) (io.ReadCloser, error) { + if m.err != nil { + return nil, m.err + } + return io.NopCloser(nil), nil +} + +func TestModelRegistry(t *testing.T) { + t.Run("RegisterProvider", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + + if registry.ProviderCount() != 1 { + t.Errorf("expected 1 provider, got %d", registry.ProviderCount()) + } + }) + + t.Run("Initialize", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model-1", Object: "model", OwnedBy: "test"}, + {ID: "test-model-2", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + + err := registry.Initialize(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if registry.ModelCount() != 2 { + t.Errorf("expected 2 models, got %d", registry.ModelCount()) + } + }) + + t.Run("GetProvider", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + _ = registry.Initialize(context.Background()) + + provider := registry.GetProvider("test-model") + if provider != mock { + t.Error("expected to get the registered provider") + } + + provider = registry.GetProvider("unknown-model") + if provider != nil { + t.Error("expected nil for unknown model") + } + }) + + t.Run("Supports", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + _ = registry.Initialize(context.Background()) + + if !registry.Supports("test-model") { + t.Error("expected Supports to return true for registered model") + } + + if registry.Supports("unknown-model") { + t.Error("expected Supports to return false for unknown model") + } + }) + + t.Run("GetModel", func(t *testing.T) { + registry := NewModelRegistry() + expectedModel := core.Model{ + ID: "test-model", + Object: "model", + OwnedBy: "test-provider", + Created: 1234567890, + } + mock := ®istryMockProvider{ + name: "test-provider", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{expectedModel}, + }, + } + registry.RegisterProvider(mock) + _ = registry.Initialize(context.Background()) + + modelInfo := registry.GetModel("test-model") + if modelInfo == nil { + t.Fatal("expected ModelInfo for registered model, got nil") + } + if modelInfo.Model.ID != expectedModel.ID { + t.Errorf("expected model ID %q, got %q", expectedModel.ID, modelInfo.Model.ID) + } + if modelInfo.Model.OwnedBy != expectedModel.OwnedBy { + t.Errorf("expected model OwnedBy %q, got %q", expectedModel.OwnedBy, modelInfo.Model.OwnedBy) + } + if modelInfo.Model.Created != expectedModel.Created { + t.Errorf("expected model Created %d, got %d", expectedModel.Created, modelInfo.Model.Created) + } + if modelInfo.Provider != mock { + t.Error("expected Provider to be the registered mock provider") + } + + unknownInfo := registry.GetModel("unknown-model") + if unknownInfo != nil { + t.Errorf("expected nil for unknown model, got %+v", unknownInfo) + } + }) + + t.Run("DuplicateModels", func(t *testing.T) { + registry := NewModelRegistry() + mock1 := ®istryMockProvider{ + name: "provider1", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "shared-model", Object: "model", OwnedBy: "provider1"}, + }, + }, + } + mock2 := ®istryMockProvider{ + name: "provider2", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "shared-model", Object: "model", OwnedBy: "provider2"}, + }, + }, + } + registry.RegisterProvider(mock1) + registry.RegisterProvider(mock2) + _ = registry.Initialize(context.Background()) + + if registry.ModelCount() != 1 { + t.Errorf("expected 1 model (deduplicated), got %d", registry.ModelCount()) + } + + provider := registry.GetProvider("shared-model") + if provider != mock1 { + t.Error("expected first provider to win for duplicate model") + } + }) + + t.Run("AllProvidersFail", func(t *testing.T) { + registry := NewModelRegistry() + mock1 := ®istryMockProvider{ + name: "provider1", + err: errors.New("provider1 error"), + } + mock2 := ®istryMockProvider{ + name: "provider2", + err: errors.New("provider2 error"), + } + registry.RegisterProvider(mock1) + registry.RegisterProvider(mock2) + + err := registry.Initialize(context.Background()) + if err == nil { + t.Error("expected error when all providers fail, got nil") + } + + expectedMsg := "failed to fetch models from any provider" + if err.Error() != expectedMsg { + t.Errorf("expected error message '%s', got '%s'", expectedMsg, err.Error()) + } + }) + + t.Run("ListModelsOrdering", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "zebra-model", Object: "model", OwnedBy: "test"}, + {ID: "alpha-model", Object: "model", OwnedBy: "test"}, + {ID: "middle-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + _ = registry.Initialize(context.Background()) + + for i := 0; i < 5; i++ { + models := registry.ListModels() + if len(models) != 3 { + t.Fatalf("expected 3 models, got %d", len(models)) + } + + if models[0].ID != "alpha-model" { + t.Errorf("expected first model to be 'alpha-model', got '%s'", models[0].ID) + } + if models[1].ID != "middle-model" { + t.Errorf("expected second model to be 'middle-model', got '%s'", models[1].ID) + } + if models[2].ID != "zebra-model" { + t.Errorf("expected third model to be 'zebra-model', got '%s'", models[2].ID) + } + } + }) + + t.Run("RefreshDoesNotBlockReads", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProvider(mock) + _ = registry.Initialize(context.Background()) + + if !registry.Supports("test-model") { + t.Fatal("expected model to be available before refresh") + } + + err := registry.Refresh(context.Background()) + if err != nil { + t.Fatalf("unexpected refresh error: %v", err) + } + + if !registry.Supports("test-model") { + t.Error("expected model to be available after refresh") + } + }) + + t.Run("GetProviderType", func(t *testing.T) { + registry := NewModelRegistry() + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + registry.RegisterProviderWithType(mock, "openai") + _ = registry.Initialize(context.Background()) + + pType := registry.GetProviderType("test-model") + if pType != "openai" { + t.Errorf("expected provider type 'openai', got '%s'", pType) + } + + pType = registry.GetProviderType("unknown-model") + if pType != "" { + t.Errorf("expected empty provider type for unknown model, got '%s'", pType) + } + }) +} + +// countingRegistryMockProvider wraps registryMockProvider and counts ListModels calls +type countingRegistryMockProvider struct { + *registryMockProvider + listCount *atomic.Int32 +} + +func (c *countingRegistryMockProvider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { + c.listCount.Add(1) + return c.registryMockProvider.ListModels(ctx) +} + +func TestStartBackgroundRefresh(t *testing.T) { + t.Run("RefreshesAtInterval", func(t *testing.T) { + var refreshCount atomic.Int32 + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + + countingMock := &countingRegistryMockProvider{ + registryMockProvider: mock, + listCount: &refreshCount, + } + + registry := NewModelRegistry() + registry.RegisterProvider(countingMock) + _ = registry.Initialize(context.Background()) + + refreshCount.Store(0) + + interval := 50 * time.Millisecond + cancel := registry.StartBackgroundRefresh(interval) + defer cancel() + + time.Sleep(interval*3 + 25*time.Millisecond) + + count := refreshCount.Load() + if count < 2 { + t.Errorf("expected at least 2 refreshes, got %d", count) + } + }) + + t.Run("StopsOnCancel", func(t *testing.T) { + var refreshCount atomic.Int32 + mock := ®istryMockProvider{ + name: "test", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "test-model", Object: "model", OwnedBy: "test"}, + }, + }, + } + + countingMock := &countingRegistryMockProvider{ + registryMockProvider: mock, + listCount: &refreshCount, + } + + registry := NewModelRegistry() + registry.RegisterProvider(countingMock) + _ = registry.Initialize(context.Background()) + + refreshCount.Store(0) + + interval := 50 * time.Millisecond + cancel := registry.StartBackgroundRefresh(interval) + cancel() + + time.Sleep(interval * 3) + + count := refreshCount.Load() + if count > 1 { + t.Errorf("expected at most 1 refresh after cancel, got %d", count) + } + }) + + t.Run("HandlesRefreshErrors", func(t *testing.T) { + var refreshCount atomic.Int32 + mock := ®istryMockProvider{ + name: "failing", + err: errors.New("refresh error"), + } + + countingMock := &countingRegistryMockProvider{ + registryMockProvider: mock, + listCount: &refreshCount, + } + + registry := NewModelRegistry() + workingMock := ®istryMockProvider{ + name: "working", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + {ID: "working-model", Object: "model", OwnedBy: "working"}, + }, + }, + } + registry.RegisterProvider(workingMock) + registry.RegisterProvider(countingMock) + _ = registry.Initialize(context.Background()) + + refreshCount.Store(0) + + interval := 50 * time.Millisecond + cancel := registry.StartBackgroundRefresh(interval) + defer cancel() + + time.Sleep(interval*3 + 25*time.Millisecond) + + count := refreshCount.Load() + if count < 2 { + t.Errorf("expected at least 2 refresh attempts despite errors, got %d", count) + } + }) +} + +// Verify ModelRegistry implements core.ModelLookup interface +var _ core.ModelLookup = (*ModelRegistry)(nil) diff --git a/internal/providers/router.go b/internal/providers/router.go index 93d946fd..549f38c8 100644 --- a/internal/providers/router.go +++ b/internal/providers/router.go @@ -12,50 +12,50 @@ import ( // ErrRegistryNotInitialized is returned when the router is used before the registry has any models. var ErrRegistryNotInitialized = fmt.Errorf("model registry has no models: ensure Initialize() or LoadFromCache() is called before using the router") -// Router routes requests to the appropriate provider based on the model registry. +// Router routes requests to the appropriate provider based on the model lookup. // It uses a dynamic model-to-provider mapping that is populated at startup // by fetching available models from each provider's /models endpoint. type Router struct { - registry *ModelRegistry + lookup core.ModelLookup } -// NewRouter creates a new provider router with a model registry. -// The registry must be initialized (via Initialize() or LoadFromCache()) before using the router. -// Returns an error if the registry is nil. -func NewRouter(registry *ModelRegistry) (*Router, error) { - if registry == nil { - return nil, fmt.Errorf("registry cannot be nil") +// NewRouter creates a new provider router with a model lookup. +// The lookup must be initialized (via Initialize() or LoadFromCache()) before using the router. +// Returns an error if the lookup is nil. +func NewRouter(lookup core.ModelLookup) (*Router, error) { + if lookup == nil { + return nil, fmt.Errorf("lookup cannot be nil") } return &Router{ - registry: registry, + lookup: lookup, }, nil } -// checkReady verifies the registry has models available. +// checkReady verifies the lookup has models available. // Returns ErrRegistryNotInitialized if no models are loaded. func (r *Router) checkReady() error { - if r.registry.ModelCount() == 0 { + if r.lookup.ModelCount() == 0 { return ErrRegistryNotInitialized } return nil } // Supports returns true if any provider supports the given model. -// Returns false if the registry has no models loaded. +// Returns false if the lookup has no models loaded. func (r *Router) Supports(model string) bool { - if r.registry.ModelCount() == 0 { + if r.lookup.ModelCount() == 0 { return false } - return r.registry.Supports(model) + return r.lookup.Supports(model) } // ChatCompletion routes the request to the appropriate provider. -// Returns ErrRegistryNotInitialized if the registry has no models loaded. +// Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*core.ChatResponse, error) { if err := r.checkReady(); err != nil { return nil, err } - provider := r.registry.GetProvider(req.Model) + provider := r.lookup.GetProvider(req.Model) if provider == nil { return nil, fmt.Errorf("no provider found for model: %s", req.Model) } @@ -63,25 +63,25 @@ func (r *Router) ChatCompletion(ctx context.Context, req *core.ChatRequest) (*co } // StreamChatCompletion routes the streaming request to the appropriate provider. -// Returns ErrRegistryNotInitialized if the registry has no models loaded. +// Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) StreamChatCompletion(ctx context.Context, req *core.ChatRequest) (io.ReadCloser, error) { if err := r.checkReady(); err != nil { return nil, err } - provider := r.registry.GetProvider(req.Model) + provider := r.lookup.GetProvider(req.Model) if provider == nil { return nil, fmt.Errorf("no provider found for model: %s", req.Model) } return provider.StreamChatCompletion(ctx, req) } -// ListModels returns all models from the registry. -// Returns ErrRegistryNotInitialized if the registry has no models loaded. +// ListModels returns all models from the lookup. +// Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) ListModels(_ context.Context) (*core.ModelsResponse, error) { if err := r.checkReady(); err != nil { return nil, err } - models := r.registry.ListModels() + models := r.lookup.ListModels() return &core.ModelsResponse{ Object: "list", Data: models, @@ -89,12 +89,12 @@ func (r *Router) ListModels(_ context.Context) (*core.ModelsResponse, error) { } // Responses routes the Responses API request to the appropriate provider. -// Returns ErrRegistryNotInitialized if the registry has no models loaded. +// Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) Responses(ctx context.Context, req *core.ResponsesRequest) (*core.ResponsesResponse, error) { if err := r.checkReady(); err != nil { return nil, err } - provider := r.registry.GetProvider(req.Model) + provider := r.lookup.GetProvider(req.Model) if provider == nil { return nil, fmt.Errorf("no provider found for model: %s", req.Model) } @@ -102,12 +102,12 @@ func (r *Router) Responses(ctx context.Context, req *core.ResponsesRequest) (*co } // StreamResponses routes the streaming Responses API request to the appropriate provider. -// Returns ErrRegistryNotInitialized if the registry has no models loaded. +// Returns ErrRegistryNotInitialized if the lookup has no models loaded. func (r *Router) StreamResponses(ctx context.Context, req *core.ResponsesRequest) (io.ReadCloser, error) { if err := r.checkReady(); err != nil { return nil, err } - provider := r.registry.GetProvider(req.Model) + provider := r.lookup.GetProvider(req.Model) if provider == nil { return nil, fmt.Errorf("no provider found for model: %s", req.Model) } @@ -117,5 +117,5 @@ func (r *Router) StreamResponses(ctx context.Context, req *core.ResponsesRequest // GetProviderType returns the provider type string for the given model. // Returns empty string if the model is not found. func (r *Router) GetProviderType(model string) string { - return r.registry.GetProviderType(model) + return r.lookup.GetProviderType(model) } diff --git a/internal/providers/router_test.go b/internal/providers/router_test.go index 5fad5b35..39813796 100644 --- a/internal/providers/router_test.go +++ b/internal/providers/router_test.go @@ -4,21 +4,59 @@ import ( "context" "errors" "io" - "sync/atomic" "testing" - "time" "gomodel/internal/core" ) +// mockModelLookup implements core.ModelLookup for fast, isolated Router testing. +// This is simpler and faster than using a full ModelRegistry with providers. +type mockModelLookup struct { + models map[string]core.Provider + providerTypes map[string]string + modelList []core.Model +} + +func newMockLookup() *mockModelLookup { + return &mockModelLookup{ + models: make(map[string]core.Provider), + providerTypes: make(map[string]string), + } +} + +func (m *mockModelLookup) addModel(model string, provider core.Provider, providerType string) { + m.models[model] = provider + m.providerTypes[model] = providerType + m.modelList = append(m.modelList, core.Model{ID: model, Object: "model"}) +} + +func (m *mockModelLookup) Supports(model string) bool { + _, ok := m.models[model] + return ok +} + +func (m *mockModelLookup) GetProvider(model string) core.Provider { + return m.models[model] +} + +func (m *mockModelLookup) GetProviderType(model string) string { + return m.providerTypes[model] +} + +func (m *mockModelLookup) ListModels() []core.Model { + return m.modelList +} + +func (m *mockModelLookup) ModelCount() int { + return len(m.models) +} + // mockProvider is a simple mock implementation of core.Provider for testing type mockProvider struct { name string chatResponse *core.ChatResponse responsesResponse *core.ResponsesResponse - modelsResponse *core.ModelsResponse err error - listModelsDelay time.Duration // optional delay before returning from ListModels } func (m *mockProvider) ChatCompletion(_ context.Context, _ *core.ChatRequest) (*core.ChatResponse, error) { @@ -35,18 +73,8 @@ func (m *mockProvider) StreamChatCompletion(_ context.Context, _ *core.ChatReque return io.NopCloser(nil), nil } -func (m *mockProvider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { - if m.listModelsDelay > 0 { - select { - case <-time.After(m.listModelsDelay): - case <-ctx.Done(): - return nil, ctx.Err() - } - } - if m.err != nil { - return nil, m.err - } - return m.modelsResponse, nil +func (m *mockProvider) ListModels(_ context.Context) (*core.ModelsResponse, error) { + return nil, nil } func (m *mockProvider) Responses(_ context.Context, _ *core.ResponsesRequest) (*core.ResponsesResponse, error) { @@ -63,31 +91,20 @@ func (m *mockProvider) StreamResponses(_ context.Context, _ *core.ResponsesReque return io.NopCloser(nil), nil } -// createTestRegistry creates a registry with mock providers for testing -func createTestRegistry(providers ...*mockProvider) *ModelRegistry { - registry := NewModelRegistry() - for _, p := range providers { - registry.RegisterProvider(p) - } - // Initialize the registry to populate models from providers - _ = registry.Initialize(context.Background()) - return registry -} - -func TestNewRouterValidation(t *testing.T) { - t.Run("NilRegistry", func(t *testing.T) { +func TestNewRouter(t *testing.T) { + t.Run("nil lookup returns error", func(t *testing.T) { router, err := NewRouter(nil) if err == nil { - t.Error("expected error for nil registry") + t.Error("expected error for nil lookup") } if router != nil { - t.Error("expected nil router for nil registry") + t.Error("expected nil router") } }) - t.Run("ValidRegistry", func(t *testing.T) { - registry := NewModelRegistry() - router, err := NewRouter(registry) + t.Run("valid lookup succeeds", func(t *testing.T) { + lookup := newMockLookup() + router, err := NewRouter(lookup) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -97,37 +114,25 @@ func TestNewRouterValidation(t *testing.T) { }) } -func TestRouterUninitializedRegistry(t *testing.T) { - // Create a router with an empty registry (no models loaded) - registry := NewModelRegistry() - router, err := NewRouter(registry) - if err != nil { - t.Fatalf("failed to create router: %v", err) - } +func TestRouterEmptyLookup(t *testing.T) { + lookup := newMockLookup() // Empty - no models + router, _ := NewRouter(lookup) t.Run("Supports returns false", func(t *testing.T) { if router.Supports("any-model") { - t.Error("expected Supports to return false for uninitialized registry") + t.Error("expected false for empty lookup") } }) t.Run("ChatCompletion returns error", func(t *testing.T) { - req := &core.ChatRequest{Model: "any-model"} - _, err := router.ChatCompletion(context.Background(), req) - if err == nil { - t.Error("expected error for uninitialized registry") - } + _, err := router.ChatCompletion(context.Background(), &core.ChatRequest{Model: "any"}) if !errors.Is(err, ErrRegistryNotInitialized) { t.Errorf("expected ErrRegistryNotInitialized, got: %v", err) } }) t.Run("StreamChatCompletion returns error", func(t *testing.T) { - req := &core.ChatRequest{Model: "any-model"} - _, err := router.StreamChatCompletion(context.Background(), req) - if err == nil { - t.Error("expected error for uninitialized registry") - } + _, err := router.StreamChatCompletion(context.Background(), &core.ChatRequest{Model: "any"}) if !errors.Is(err, ErrRegistryNotInitialized) { t.Errorf("expected ErrRegistryNotInitialized, got: %v", err) } @@ -135,31 +140,20 @@ func TestRouterUninitializedRegistry(t *testing.T) { t.Run("ListModels returns error", func(t *testing.T) { _, err := router.ListModels(context.Background()) - if err == nil { - t.Error("expected error for uninitialized registry") - } if !errors.Is(err, ErrRegistryNotInitialized) { t.Errorf("expected ErrRegistryNotInitialized, got: %v", err) } }) t.Run("Responses returns error", func(t *testing.T) { - req := &core.ResponsesRequest{Model: "any-model"} - _, err := router.Responses(context.Background(), req) - if err == nil { - t.Error("expected error for uninitialized registry") - } + _, err := router.Responses(context.Background(), &core.ResponsesRequest{Model: "any"}) if !errors.Is(err, ErrRegistryNotInitialized) { t.Errorf("expected ErrRegistryNotInitialized, got: %v", err) } }) t.Run("StreamResponses returns error", func(t *testing.T) { - req := &core.ResponsesRequest{Model: "any-model"} - _, err := router.StreamResponses(context.Background(), req) - if err == nil { - t.Error("expected error for uninitialized registry") - } + _, err := router.StreamResponses(context.Background(), &core.ResponsesRequest{Model: "any"}) if !errors.Is(err, ErrRegistryNotInitialized) { t.Errorf("expected ErrRegistryNotInitialized, got: %v", err) } @@ -167,620 +161,175 @@ func TestRouterUninitializedRegistry(t *testing.T) { } func TestRouterSupports(t *testing.T) { - openaiMock := &mockProvider{ - name: "openai", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "gpt-4o", Object: "model", OwnedBy: "openai"}, - }, - }, - } - anthropicMock := &mockProvider{ - name: "anthropic", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "claude-3-5-sonnet-20241022", Object: "model", OwnedBy: "anthropic"}, - }, - }, - } + openai := &mockProvider{name: "openai"} + anthropic := &mockProvider{name: "anthropic"} - registry := createTestRegistry(openaiMock, anthropicMock) - router, err := NewRouter(registry) - if err != nil { - t.Fatalf("failed to create router: %v", err) - } + lookup := newMockLookup() + lookup.addModel("gpt-4o", openai, "openai") + lookup.addModel("claude-3-5-sonnet", anthropic, "anthropic") + + router, _ := NewRouter(lookup) tests := []struct { model string expected bool }{ {"gpt-4o", true}, - {"claude-3-5-sonnet-20241022", true}, - {"unsupported-model", false}, + {"claude-3-5-sonnet", true}, + {"unsupported", false}, } for _, tt := range tests { t.Run(tt.model, func(t *testing.T) { - result := router.Supports(tt.model) - if result != tt.expected { - t.Errorf("Supports(%q) = %v, want %v", tt.model, result, tt.expected) + if got := router.Supports(tt.model); got != tt.expected { + t.Errorf("Supports(%q) = %v, want %v", tt.model, got, tt.expected) } }) } } func TestRouterChatCompletion(t *testing.T) { - openaiResp := &core.ChatResponse{ID: "openai-response", Model: "gpt-4o"} - anthropicResp := &core.ChatResponse{ID: "anthropic-response", Model: "claude-3-5-sonnet-20241022"} - - openaiMock := &mockProvider{ - name: "openai", - chatResponse: openaiResp, - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "gpt-4o", Object: "model", OwnedBy: "openai"}, - }, - }, - } - anthropicMock := &mockProvider{ - name: "anthropic", - chatResponse: anthropicResp, - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "claude-3-5-sonnet-20241022", Object: "model", OwnedBy: "anthropic"}, - }, - }, - } + openaiResp := &core.ChatResponse{ID: "openai-resp", Model: "gpt-4o"} + anthropicResp := &core.ChatResponse{ID: "anthropic-resp", Model: "claude-3-5-sonnet"} - registry := createTestRegistry(openaiMock, anthropicMock) - router, err := NewRouter(registry) - if err != nil { - t.Fatalf("failed to create router: %v", err) - } + openai := &mockProvider{name: "openai", chatResponse: openaiResp} + anthropic := &mockProvider{name: "anthropic", chatResponse: anthropicResp} + + lookup := newMockLookup() + lookup.addModel("gpt-4o", openai, "openai") + lookup.addModel("claude-3-5-sonnet", anthropic, "anthropic") + + router, _ := NewRouter(lookup) tests := []struct { - name string - model string - expectedResp *core.ChatResponse - expectedError bool + name string + model string + wantResp *core.ChatResponse + wantError bool }{ - { - name: "route to openai", - model: "gpt-4o", - expectedResp: openaiResp, - }, - { - name: "route to anthropic", - model: "claude-3-5-sonnet-20241022", - expectedResp: anthropicResp, - }, - { - name: "unsupported model", - model: "unsupported-model", - expectedError: true, - }, + {"routes to openai", "gpt-4o", openaiResp, false}, + {"routes to anthropic", "claude-3-5-sonnet", anthropicResp, false}, + {"unsupported model", "unknown", nil, true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - req := &core.ChatRequest{ - Model: tt.model, - Messages: []core.Message{ - {Role: "user", Content: "test"}, - }, - } - + req := &core.ChatRequest{Model: tt.model} resp, err := router.ChatCompletion(context.Background(), req) - if tt.expectedError { + if tt.wantError { if err == nil { t.Error("expected error, got nil") } - } else { - if err != nil { - t.Errorf("unexpected error: %v", err) - } - if resp != tt.expectedResp { - t.Errorf("got response ID %q, want %q", resp.ID, tt.expectedResp.ID) - } + return + } + + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if resp != tt.wantResp { + t.Errorf("got response %v, want %v", resp, tt.wantResp) } }) } } -func TestRouterListModels(t *testing.T) { - openaiMock := &mockProvider{ - name: "openai", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "gpt-4o", Object: "model", OwnedBy: "openai"}, - }, - }, - } - anthropicMock := &mockProvider{ - name: "anthropic", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "claude-3-5-sonnet-20241022", Object: "model", OwnedBy: "anthropic"}, - }, - }, - } +func TestRouterResponses(t *testing.T) { + expectedResp := &core.ResponsesResponse{ID: "resp-123"} + provider := &mockProvider{name: "openai", responsesResponse: expectedResp} - registry := createTestRegistry(openaiMock, anthropicMock) - router, err := NewRouter(registry) - if err != nil { - t.Fatalf("failed to create router: %v", err) - } + lookup := newMockLookup() + lookup.addModel("gpt-4o", provider, "openai") - resp, err := router.ListModels(context.Background()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if len(resp.Data) != 2 { - t.Errorf("expected 2 models, got %d", len(resp.Data)) - } + router, _ := NewRouter(lookup) - // Verify both providers' models are included - foundOpenAI := false - foundAnthropic := false - for _, model := range resp.Data { - if model.ID == "gpt-4o" { - foundOpenAI = true + t.Run("routes correctly", func(t *testing.T) { + req := &core.ResponsesRequest{Model: "gpt-4o"} + resp, err := router.Responses(context.Background(), req) + if err != nil { + t.Errorf("unexpected error: %v", err) } - if model.ID == "claude-3-5-sonnet-20241022" { - foundAnthropic = true + if resp != expectedResp { + t.Errorf("got %v, want %v", resp, expectedResp) } - } + }) - if !foundOpenAI { - t.Error("OpenAI model not found in combined list") - } - if !foundAnthropic { - t.Error("Anthropic model not found in combined list") - } + t.Run("unknown model returns error", func(t *testing.T) { + req := &core.ResponsesRequest{Model: "unknown"} + _, err := router.Responses(context.Background(), req) + if err == nil { + t.Error("expected error for unknown model") + } + }) } -func TestRouterListModelsWithError(t *testing.T) { - // Test that router continues even if one provider fails during initialization - openaiMock := &mockProvider{ - name: "openai", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "gpt-4o", Object: "model", OwnedBy: "openai"}, - }, - }, - } - anthropicMock := &mockProvider{ - name: "anthropic", - err: errors.New("provider error"), - } +func TestRouterListModels(t *testing.T) { + lookup := newMockLookup() + lookup.addModel("gpt-4o", &mockProvider{}, "openai") + lookup.addModel("claude-3-5-sonnet", &mockProvider{}, "anthropic") - registry := createTestRegistry(openaiMock, anthropicMock) - router, err := NewRouter(registry) - if err != nil { - t.Fatalf("failed to create router: %v", err) - } + router, _ := NewRouter(lookup) resp, err := router.ListModels(context.Background()) if err != nil { t.Fatalf("unexpected error: %v", err) } - // Should still get OpenAI models even though Anthropic failed during initialization - if len(resp.Data) != 1 { - t.Errorf("expected 1 model, got %d", len(resp.Data)) + if len(resp.Data) != 2 { + t.Errorf("expected 2 models, got %d", len(resp.Data)) } - if resp.Data[0].ID != "gpt-4o" { - t.Errorf("expected gpt-4o, got %s", resp.Data[0].ID) + if resp.Object != "list" { + t.Errorf("expected object 'list', got %q", resp.Object) } } -func TestModelRegistry(t *testing.T) { - t.Run("RegisterProvider", func(t *testing.T) { - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - - if registry.ProviderCount() != 1 { - t.Errorf("expected 1 provider, got %d", registry.ProviderCount()) - } - }) - - t.Run("Initialize", func(t *testing.T) { - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model-1", Object: "model", OwnedBy: "test"}, - {ID: "test-model-2", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - - err := registry.Initialize(context.Background()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if registry.ModelCount() != 2 { - t.Errorf("expected 2 models, got %d", registry.ModelCount()) - } - }) - - t.Run("GetProvider", func(t *testing.T) { - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - _ = registry.Initialize(context.Background()) +func TestRouterGetProviderType(t *testing.T) { + lookup := newMockLookup() + lookup.addModel("gpt-4o", &mockProvider{}, "openai") + lookup.addModel("claude-3-5-sonnet", &mockProvider{}, "anthropic") - provider := registry.GetProvider("test-model") - if provider != mock { - t.Error("expected to get the registered provider") - } - - provider = registry.GetProvider("unknown-model") - if provider != nil { - t.Error("expected nil for unknown model") - } - }) + router, _ := NewRouter(lookup) - t.Run("Supports", func(t *testing.T) { - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - _ = registry.Initialize(context.Background()) - - if !registry.Supports("test-model") { - t.Error("expected Supports to return true for registered model") - } - - if registry.Supports("unknown-model") { - t.Error("expected Supports to return false for unknown model") - } - }) - - t.Run("GetModel", func(t *testing.T) { - registry := NewModelRegistry() - expectedModel := core.Model{ - ID: "test-model", - Object: "model", - OwnedBy: "test-provider", - Created: 1234567890, - } - mock := &mockProvider{ - name: "test-provider", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{expectedModel}, - }, - } - registry.RegisterProvider(mock) - _ = registry.Initialize(context.Background()) - - // Test getting a registered model - modelInfo := registry.GetModel("test-model") - if modelInfo == nil { - t.Fatal("expected ModelInfo for registered model, got nil") - } - if modelInfo.Model.ID != expectedModel.ID { - t.Errorf("expected model ID %q, got %q", expectedModel.ID, modelInfo.Model.ID) - } - if modelInfo.Model.OwnedBy != expectedModel.OwnedBy { - t.Errorf("expected model OwnedBy %q, got %q", expectedModel.OwnedBy, modelInfo.Model.OwnedBy) - } - if modelInfo.Model.Created != expectedModel.Created { - t.Errorf("expected model Created %d, got %d", expectedModel.Created, modelInfo.Model.Created) - } - if modelInfo.Provider != mock { - t.Error("expected Provider to be the registered mock provider") - } - - // Test getting an unknown model - unknownInfo := registry.GetModel("unknown-model") - if unknownInfo != nil { - t.Errorf("expected nil for unknown model, got %+v", unknownInfo) - } - }) - - t.Run("DuplicateModels", func(t *testing.T) { - // Test that first provider wins when models have the same ID - registry := NewModelRegistry() - mock1 := &mockProvider{ - name: "provider1", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "shared-model", Object: "model", OwnedBy: "provider1"}, - }, - }, - } - mock2 := &mockProvider{ - name: "provider2", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "shared-model", Object: "model", OwnedBy: "provider2"}, - }, - }, - } - registry.RegisterProvider(mock1) - registry.RegisterProvider(mock2) - _ = registry.Initialize(context.Background()) - - // Should only have one model (first provider wins) - if registry.ModelCount() != 1 { - t.Errorf("expected 1 model (deduplicated), got %d", registry.ModelCount()) - } - - // First provider should be the one associated with the model - provider := registry.GetProvider("shared-model") - if provider != mock1 { - t.Error("expected first provider to win for duplicate model") - } - }) - - t.Run("AllProvidersFail", func(t *testing.T) { - // Test that Initialize returns an error when all providers fail - registry := NewModelRegistry() - mock1 := &mockProvider{ - name: "provider1", - err: errors.New("provider1 error"), - } - mock2 := &mockProvider{ - name: "provider2", - err: errors.New("provider2 error"), - } - registry.RegisterProvider(mock1) - registry.RegisterProvider(mock2) - - err := registry.Initialize(context.Background()) - if err == nil { - t.Error("expected error when all providers fail, got nil") - } - - expectedMsg := "failed to fetch models from any provider" - if err.Error() != expectedMsg { - t.Errorf("expected error message '%s', got '%s'", expectedMsg, err.Error()) - } - }) - - t.Run("ListModelsOrdering", func(t *testing.T) { - // Test that ListModels returns models in consistent sorted order - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "zebra-model", Object: "model", OwnedBy: "test"}, - {ID: "alpha-model", Object: "model", OwnedBy: "test"}, - {ID: "middle-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - _ = registry.Initialize(context.Background()) - - // Call ListModels multiple times and verify consistent ordering - for i := 0; i < 5; i++ { - models := registry.ListModels() - if len(models) != 3 { - t.Fatalf("expected 3 models, got %d", len(models)) - } + tests := []struct { + model string + expected string + }{ + {"gpt-4o", "openai"}, + {"claude-3-5-sonnet", "anthropic"}, + {"unknown", ""}, + } - // Verify sorted order - if models[0].ID != "alpha-model" { - t.Errorf("expected first model to be 'alpha-model', got '%s'", models[0].ID) - } - if models[1].ID != "middle-model" { - t.Errorf("expected second model to be 'middle-model', got '%s'", models[1].ID) - } - if models[2].ID != "zebra-model" { - t.Errorf("expected third model to be 'zebra-model', got '%s'", models[2].ID) + for _, tt := range tests { + t.Run(tt.model, func(t *testing.T) { + if got := router.GetProviderType(tt.model); got != tt.expected { + t.Errorf("GetProviderType(%q) = %q, want %q", tt.model, got, tt.expected) } - } - }) - - t.Run("RefreshDoesNotBlockReads", func(t *testing.T) { - // Test that Refresh builds new map without blocking concurrent reads - registry := NewModelRegistry() - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - registry.RegisterProvider(mock) - _ = registry.Initialize(context.Background()) - - // Verify model is available before refresh - if !registry.Supports("test-model") { - t.Fatal("expected model to be available before refresh") - } - - // During refresh, the model should still be accessible - // (testing the atomic swap behavior) - err := registry.Refresh(context.Background()) - if err != nil { - t.Fatalf("unexpected refresh error: %v", err) - } - - // Model should still be available after refresh - if !registry.Supports("test-model") { - t.Error("expected model to be available after refresh") - } - }) + }) + } } -func TestStartBackgroundRefresh(t *testing.T) { - t.Run("RefreshesAtInterval", func(t *testing.T) { - var refreshCount atomic.Int32 - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - - // Create a wrapper that counts refreshes - countingMock := &countingMockProvider{ - mockProvider: mock, - listCount: &refreshCount, - } - - registry := NewModelRegistry() - registry.RegisterProvider(countingMock) - _ = registry.Initialize(context.Background()) +func TestRouterProviderError(t *testing.T) { + providerErr := errors.New("provider error") + provider := &mockProvider{name: "failing", err: providerErr} - // Reset counter after initial initialization - refreshCount.Store(0) + lookup := newMockLookup() + lookup.addModel("failing-model", provider, "test") - // Start background refresh with a short interval - interval := 50 * time.Millisecond - cancel := registry.StartBackgroundRefresh(interval) - defer cancel() + router, _ := NewRouter(lookup) - // Wait for a few refresh cycles - time.Sleep(interval*3 + 25*time.Millisecond) - - count := refreshCount.Load() - if count < 2 { - t.Errorf("expected at least 2 refreshes, got %d", count) - } - }) - - t.Run("StopsOnCancel", func(t *testing.T) { - var refreshCount atomic.Int32 - mock := &mockProvider{ - name: "test", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "test-model", Object: "model", OwnedBy: "test"}, - }, - }, - } - - countingMock := &countingMockProvider{ - mockProvider: mock, - listCount: &refreshCount, - } - - registry := NewModelRegistry() - registry.RegisterProvider(countingMock) - _ = registry.Initialize(context.Background()) - - // Reset counter after initial initialization - refreshCount.Store(0) - - // Start and immediately cancel - interval := 50 * time.Millisecond - cancel := registry.StartBackgroundRefresh(interval) - cancel() - - // Wait a bit to ensure no more refreshes happen - time.Sleep(interval * 3) - - count := refreshCount.Load() - if count > 1 { - t.Errorf("expected at most 1 refresh after cancel, got %d", count) + t.Run("ChatCompletion propagates error", func(t *testing.T) { + req := &core.ChatRequest{Model: "failing-model"} + _, err := router.ChatCompletion(context.Background(), req) + if !errors.Is(err, providerErr) { + t.Errorf("expected provider error, got: %v", err) } }) - t.Run("HandlesRefreshErrors", func(t *testing.T) { - var refreshCount atomic.Int32 - mock := &mockProvider{ - name: "failing", - err: errors.New("refresh error"), - } - - countingMock := &countingMockProvider{ - mockProvider: mock, - listCount: &refreshCount, - } - - registry := NewModelRegistry() - // First add a working provider to initialize successfully - workingMock := &mockProvider{ - name: "working", - modelsResponse: &core.ModelsResponse{ - Object: "list", - Data: []core.Model{ - {ID: "working-model", Object: "model", OwnedBy: "working"}, - }, - }, - } - registry.RegisterProvider(workingMock) - registry.RegisterProvider(countingMock) - _ = registry.Initialize(context.Background()) - - // Reset counter - refreshCount.Store(0) - - // Start background refresh - should continue even with errors - interval := 50 * time.Millisecond - cancel := registry.StartBackgroundRefresh(interval) - defer cancel() - - // Wait for refresh attempts - time.Sleep(interval*3 + 25*time.Millisecond) - - // The background refresh should continue attempting even with errors - count := refreshCount.Load() - if count < 2 { - t.Errorf("expected at least 2 refresh attempts despite errors, got %d", count) + t.Run("Responses propagates error", func(t *testing.T) { + req := &core.ResponsesRequest{Model: "failing-model"} + _, err := router.Responses(context.Background(), req) + if !errors.Is(err, providerErr) { + t.Errorf("expected provider error, got: %v", err) } }) } - -// countingMockProvider wraps mockProvider and counts ListModels calls -type countingMockProvider struct { - *mockProvider - listCount *atomic.Int32 -} - -func (c *countingMockProvider) ListModels(ctx context.Context) (*core.ModelsResponse, error) { - c.listCount.Add(1) - return c.mockProvider.ListModels(ctx) -} diff --git a/internal/providers/xai/xai.go b/internal/providers/xai/xai.go index f4a9e2b0..e330aa5f 100644 --- a/internal/providers/xai/xai.go +++ b/internal/providers/xai/xai.go @@ -11,37 +11,40 @@ import ( "gomodel/internal/providers" ) +// Registration provides factory registration for the xAI provider. +var Registration = providers.Registration{ + Type: "xai", + New: New, +} + const ( defaultBaseURL = "https://api.x.ai/v1" ) -func init() { - // Self-register with the factory - providers.RegisterProvider("xai", New) -} - // Provider implements the core.Provider interface for xAI type Provider struct { client *llmclient.Client apiKey string } -// New creates a new xAI provider -func New(apiKey string) *Provider { +// New creates a new xAI provider. +func New(apiKey string, hooks llmclient.Hooks) core.Provider { p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("xai", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.New(cfg, p.setHeaders) return p } -// NewWithHTTPClient creates a new xAI provider with a custom HTTP client -func NewWithHTTPClient(apiKey string, httpClient *http.Client) *Provider { +// NewWithHTTPClient creates a new xAI provider with a custom HTTP client. +// If httpClient is nil, http.DefaultClient is used. +func NewWithHTTPClient(apiKey string, httpClient *http.Client, hooks llmclient.Hooks) *Provider { + if httpClient == nil { + httpClient = http.DefaultClient + } p := &Provider{apiKey: apiKey} cfg := llmclient.DefaultConfig("xai", defaultBaseURL) - // Apply global hooks if available - cfg.Hooks = providers.GetGlobalHooks() + cfg.Hooks = hooks p.client = llmclient.NewWithHTTPClient(httpClient, cfg, p.setHeaders) return p } diff --git a/internal/providers/xai/xai_test.go b/internal/providers/xai/xai_test.go index 90e1b7b3..2da7a2af 100644 --- a/internal/providers/xai/xai_test.go +++ b/internal/providers/xai/xai_test.go @@ -10,11 +10,13 @@ import ( "testing" "gomodel/internal/core" + "gomodel/internal/llmclient" ) func TestNew(t *testing.T) { apiKey := "test-api-key" - provider := New(apiKey) + // Use NewWithHTTPClient to get concrete type for internal testing + provider := NewWithHTTPClient(apiKey, nil, llmclient.Hooks{}) if provider.apiKey != apiKey { t.Errorf("apiKey = %q, want %q", provider.apiKey, apiKey) @@ -24,6 +26,14 @@ func TestNew(t *testing.T) { } } +func TestNew_ReturnsProvider(t *testing.T) { + provider := New("test-api-key", llmclient.Hooks{}) + + if provider == nil { + t.Error("provider should not be nil") + } +} + // customHeaderRoundTripper is a RoundTripper that injects a custom header type customHeaderRoundTripper struct { transport http.RoundTripper @@ -72,7 +82,7 @@ func TestNewWithHTTPClient(t *testing.T) { } // Create provider with custom HTTP client - provider := NewWithHTTPClient("test-api-key", customClient) + provider := NewWithHTTPClient("test-api-key", customClient, llmclient.Hooks{}) // Verify provider is non-nil if provider == nil { @@ -217,7 +227,7 @@ func TestChatCompletion(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -301,7 +311,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ChatRequest{ @@ -413,7 +423,7 @@ func TestListModels(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) resp, err := provider.ListModels(context.Background()) @@ -442,7 +452,7 @@ func TestChatCompletionWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background()) @@ -583,7 +593,7 @@ func TestResponses(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -708,7 +718,7 @@ data: [DONE] })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) req := &core.ResponsesRequest{ @@ -742,7 +752,7 @@ func TestResponsesWithContext(t *testing.T) { })) defer server.Close() - provider := New("test-api-key") + provider := NewWithHTTPClient("test-api-key", nil, llmclient.Hooks{}) provider.SetBaseURL(server.URL) ctx, cancel := context.WithCancel(context.Background())