diff --git a/CLAUDE.md b/CLAUDE.md index 77d5d86..bb2cb90 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -128,10 +128,13 @@ their repo-root counterparts. Install-specific entrypoints live in `.codex/`, | Variable | Default | Description | | ------------------------ | ------------------------ | ------------------------------------------ | -| `LUMEN_BACKEND` | `ollama` | Embedding backend (`ollama` or `lmstudio`) | +| `LUMEN_BACKEND` | `ollama` | Embedding backend (`ollama`, `lmstudio`, or `openai`) | | `LUMEN_EMBED_MODEL` | see note ¹ | Embedding model (must be in registry) | | `OLLAMA_HOST` | `http://localhost:11434` | Ollama server URL | | `LM_STUDIO_HOST` | `http://localhost:1234` | LM Studio server URL | +| `OPENAI_BASE_URL` | — | OpenAI-compatible server URL (`openai` backend) | +| `OPENAI_API_KEY` | — | Bearer token for the `openai` backend | +| `LUMEN_EMBED_SKIP_HEALTH_CHECK` | `false` | Skip `/v1/models` probe (`openai` backend) | | `LUMEN_MAX_CHUNK_TOKENS` | `512` | Max tokens per chunk before splitting | | `LUMEN_VECTOR_STORAGE` | `int8` | Vector precision (`int8` or `float32`) | diff --git a/README.md b/README.md index 8d33e76..a22608a 100644 --- a/README.md +++ b/README.md @@ -273,9 +273,12 @@ All configuration is via environment variables: | Variable | Default | Description | | ------------------------ | ------------------------ | ------------------------------------------------------------- | | `LUMEN_EMBED_MODEL` | see note ¹ | Embedding model; use with `LUMEN_EMBED_DIMS` for unlisted models | -| `LUMEN_BACKEND` | `ollama` | Embedding backend (`ollama` or `lmstudio`) | +| `LUMEN_BACKEND` | `ollama` | Embedding backend (`ollama`, `lmstudio`, or `openai`) | | `OLLAMA_HOST` | `http://localhost:11434` | Ollama server URL | | `LM_STUDIO_HOST` | `http://localhost:1234` | LM Studio server URL | +| `OPENAI_BASE_URL` | — | OpenAI-compatible server URL (required for `openai` backend) | +| `OPENAI_API_KEY` | — | Bearer token for the `openai` backend; omit for gateways that trust network position | +| `LUMEN_EMBED_SKIP_HEALTH_CHECK` | `false` | Skip the `/v1/models` probe (`openai` backend); needed when a gateway doesn't expose that endpoint | | `LUMEN_MAX_CHUNK_TOKENS` | `512` | Max tokens per chunk before splitting | | `LUMEN_VECTOR_STORAGE` | `int8` | Vector precision (`int8` or `float32`) | | `LUMEN_EMBED_DIMS` | — | Override embedding dimensions (required for unlisted models) | @@ -304,7 +307,9 @@ of the database path hash, so different models never collide. > **Caveat**: the DB path hash includes the model name but not the backend. If > the same model name is configured on two backends (e.g. an Ollama and an LM > Studio entry both named `foo`), they share the same index — use distinct -> model names per backend to avoid collisions. +> model names per backend to avoid collisions. This applies to the `openai` +> backend too — a model name shared with an Ollama or LM Studio entry collides +> in the index cache. ### Selecting a server per invocation @@ -348,6 +353,37 @@ LUMEN_EMBED_DIMS=4096 LUMEN_EMBED_CTX=40960 # optional, defaults to 8192 ``` +### Remote / internal OpenAI-compatible servers + +The `openai` backend targets any service exposing an OpenAI-compatible +`/v1/embeddings` endpoint — OpenAI itself, or an internal gateway. Configure +it via `config.yaml`: + +```yaml +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-... # optional — omit for gateways that trust network position + skip_health_check: true # optional — set when the gateway doesn't expose /v1/models +``` + +Or via environment variables: + +```sh +LUMEN_BACKEND=openai +OPENAI_BASE_URL=https://api.example.com +OPENAI_API_KEY=sk-... +LUMEN_EMBED_MODEL=text-embedding-3-small +LUMEN_EMBED_DIMS=1536 +``` + +`skip_health_check` exists because many custom gateways proxy a different +provider underneath (Bedrock, Gemini, etc.) and don't implement `/v1/models` +in the OpenAI shape, or don't implement it at all. Without it, Lumen's health +probe would incorrectly mark a working server as unhealthy. + ## Controlling what gets indexed Lumen filters files through six layers: built-in directory and lock file skips → diff --git a/cmd/index.go b/cmd/index.go index bcd4fc8..66df211 100644 --- a/cmd/index.go +++ b/cmd/index.go @@ -36,7 +36,7 @@ import ( func init() { indexCmd.Flags().StringP("model", "m", "", "embedding model (default: $LUMEN_EMBED_MODEL or "+embedder.DefaultModel+")") - indexCmd.Flags().StringP("backend", "b", "", "embedding backend to select (\"ollama\" or \"lmstudio\"); disambiguates when --model is configured on multiple backends") + indexCmd.Flags().StringP("backend", "b", "", "embedding backend to select (\"ollama\", \"lmstudio\", or \"openai\"); disambiguates when --model is configured on multiple backends") indexCmd.Flags().BoolP("force", "f", false, "force full re-index") rootCmd.AddCommand(indexCmd) } @@ -195,9 +195,9 @@ func loadConfigWithFlags(cmd *cobra.Command) (*config.ConfigService, error) { if model == "" && backend == "" { return config.NewConfigService(path) } - if backend != "" && backend != config.BackendOllama && backend != config.BackendLMStudio { - return nil, fmt.Errorf("unknown backend %q (must be %q or %q)", - backend, config.BackendOllama, config.BackendLMStudio) + if backend != "" && backend != config.BackendOllama && backend != config.BackendLMStudio && backend != config.BackendOpenAI { + return nil, fmt.Errorf("unknown backend %q (must be %q, %q, or %q)", + backend, config.BackendOllama, config.BackendLMStudio, config.BackendOpenAI) } cfg, selErr := config.NewConfigService(path, config.WithServerSelection(model, backend)) diff --git a/cmd/search.go b/cmd/search.go index 67003b5..62ed70c 100644 --- a/cmd/search.go +++ b/cmd/search.go @@ -79,7 +79,7 @@ func init() { searchCmd.Flags().BoolP("force", "f", false, "force full re-index before searching") searchCmd.Flags().Bool("trace", false, "print per-phase timing to stderr") searchCmd.Flags().StringP("model", "m", "", "embedding model override") - searchCmd.Flags().StringP("backend", "b", "", "embedding backend to select (\"ollama\" or \"lmstudio\")") + searchCmd.Flags().StringP("backend", "b", "", "embedding backend to select (\"ollama\", \"lmstudio\", or \"openai\")") rootCmd.AddCommand(searchCmd) } diff --git a/cmd/stdio.go b/cmd/stdio.go index 30e02ee..97af91b 100644 --- a/cmd/stdio.go +++ b/cmd/stdio.go @@ -1132,6 +1132,10 @@ func (ic *indexerCache) handleHealthCheck(ctx context.Context, _ *mcp.CallToolRe host := srv.Host model := srv.Model + if srv.SkipHealthCheck { + return healthResult(backend, host, model, true, "skip_health_check is set; service was not probed"), nil, nil + } + probeCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() diff --git a/cmd/stdio_test.go b/cmd/stdio_test.go index 794ef32..c4a8e9f 100644 --- a/cmd/stdio_test.go +++ b/cmd/stdio_test.go @@ -487,6 +487,53 @@ servers: } } +func TestHandleHealthCheck_RespectsSkipHealthCheck(t *testing.T) { + for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX", "OLLAMA_HOST", "LM_STUDIO_HOST", "OPENAI_API_KEY", "OPENAI_BASE_URL", "LUMEN_EMBED_SKIP_HEALTH_CHECK"} { + t.Setenv(k, "") + } + + dir := t.TempDir() + cfgFile := filepath.Join(dir, "config.yaml") + + // /v1/models always 503s, simulating a gateway that doesn't implement it, + // even though /v1/embeddings itself works fine. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/models" { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + if err := os.WriteFile(cfgFile, []byte(fmt.Sprintf(` +servers: + - backend: openai + host: %s + model: remote-embed + dims: 3 + skip_health_check: true +`, srv.URL)), 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + svc, err := config.NewConfigService(cfgFile) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + fe := embedder.NewFailoverEmbedder(svc) + + ic := &indexerCache{embedder: fe, cfg: svc} + result, _, err := ic.handleHealthCheck(context.Background(), &mcp.CallToolRequest{}, HealthCheckInput{}) + if err != nil { + t.Fatalf("handleHealthCheck: %v", err) + } + text := mustTextResult(t, result) + if !strings.Contains(text, "Status: OK") { + t.Fatalf("expected skip_health_check to report OK without probing /v1/models, got: %s", text) + } +} + func TestHandleHealthCheck_ModelMissingIsError(t *testing.T) { for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX", "OLLAMA_HOST", "LM_STUDIO_HOST"} { t.Setenv(k, "") diff --git a/internal/config/config.go b/internal/config/config.go index f1d6553..54edb82 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -31,6 +31,10 @@ const ( BackendOllama = "ollama" // BackendLMStudio is the backend identifier for LM Studio. BackendLMStudio = "lmstudio" + // BackendOpenAI is the backend identifier for OpenAI-compatible remote + // embedding servers (OpenAI itself, or any internal gateway exposing the + // same /v1/embeddings wire format). + BackendOpenAI = "openai" ) // DBPathForProject returns the default int8/512-token collection path. New diff --git a/internal/config/service.go b/internal/config/service.go index 560aec1..8992750 100644 --- a/internal/config/service.go +++ b/internal/config/service.go @@ -2,6 +2,7 @@ package config import ( "fmt" + "net" "net/url" "os" "path/filepath" @@ -21,12 +22,14 @@ import ( // ServerConfig holds per-server configuration. type ServerConfig struct { - Backend string `koanf:"backend"` - Host string `koanf:"host"` - Model string `koanf:"model"` - Dims int `koanf:"dims"` - CtxLength int `koanf:"ctx_length"` - MinScore float64 `koanf:"min_score"` + Backend string `koanf:"backend"` + Host string `koanf:"host"` + Model string `koanf:"model"` + Dims int `koanf:"dims"` + CtxLength int `koanf:"ctx_length"` + MinScore float64 `koanf:"min_score"` + APIKey string `koanf:"api_key"` + SkipHealthCheck bool `koanf:"skip_health_check"` } // ConfigService wraps koanf and provides typed config access. @@ -50,12 +53,27 @@ func defaultServerForBackend(backend string) ServerConfig { Host: "http://localhost:1234", Model: models.DefaultLMStudioModel, } - default: + case BackendOpenAI: + // No sane localhost default for a remote/internal gateway — host and + // model must be set explicitly via env vars or config file. + return ServerConfig{ + Backend: BackendOpenAI, + } + case BackendOllama: return ServerConfig{ Backend: BackendOllama, Host: "http://localhost:11434", Model: models.DefaultOllamaModel, } + default: + // Preserve whatever was passed instead of silently normalizing an + // unrecognized value to Ollama. This function is reached from + // applyEnvOverrides with a raw, unvalidated LUMEN_BACKEND value — a + // typo like "OpenAI" (wrong case) must surface as validate()'s + // "unknown backend" error, not silently redirect the server to + // localhost Ollama while dropping an api_key that was only valid for + // the originally configured backend. + return ServerConfig{Backend: backend} } } @@ -131,14 +149,7 @@ func NewConfigService(configPath string, opts ...Option) (*ConfigService, error) _ = k.Unmarshal("servers", &servers) if len(servers) > 0 { servers[0].Model = svc.modelOverride - serverMaps := make([]map[string]any, len(servers)) - for i, s := range servers { - serverMaps[i] = map[string]any{ - "backend": s.Backend, "host": s.Host, "model": s.Model, - "dims": s.Dims, "ctx_length": s.CtxLength, "min_score": s.MinScore, - } - } - _ = k.Load(confmap.Provider(map[string]any{"servers": serverMaps}, "."), nil) + _ = k.Load(confmap.Provider(map[string]any{"servers": serverConfigMaps(servers)}, "."), nil) } } @@ -156,14 +167,7 @@ func NewConfigService(configPath string, opts ...Option) (*ConfigService, error) // Drop the existing list first — koanf merge would otherwise keep // stale entries beyond the filtered length. k.Delete("servers") - serverMaps := make([]map[string]any, len(filtered)) - for i, s := range filtered { - serverMaps[i] = map[string]any{ - "backend": s.Backend, "host": s.Host, "model": s.Model, - "dims": s.Dims, "ctx_length": s.CtxLength, "min_score": s.MinScore, - } - } - _ = k.Load(confmap.Provider(map[string]any{"servers": serverMaps}, "."), nil) + _ = k.Load(confmap.Provider(map[string]any{"servers": serverConfigMaps(filtered)}, "."), nil) } if err := svc.validate(); err != nil { @@ -235,9 +239,13 @@ func applyEnvOverrides(k *koanf.Koanf) { ctx := os.Getenv("LUMEN_EMBED_CTX") ollamaHost := os.Getenv("OLLAMA_HOST") lmStudioHost := os.Getenv("LM_STUDIO_HOST") + openAIHost := os.Getenv("OPENAI_BASE_URL") + openAIKey := os.Getenv("OPENAI_API_KEY") + skipHealthCheck := os.Getenv("LUMEN_EMBED_SKIP_HEALTH_CHECK") // Only apply if at least one server env var is explicitly set - hasOverride := backendEnv != "" || model != "" || dims != "" || ctx != "" || ollamaHost != "" || lmStudioHost != "" + hasOverride := backendEnv != "" || model != "" || dims != "" || ctx != "" || + ollamaHost != "" || lmStudioHost != "" || openAIHost != "" || openAIKey != "" || skipHealthCheck != "" if !hasOverride { return } @@ -250,10 +258,30 @@ func applyEnvOverrides(k *koanf.Koanf) { } srv := servers[0] - // If backend is explicitly overridden, reset server[0] to backend-specific - // defaults first to avoid mixed config (e.g. lmstudio backend with Ollama host/model). - if backendEnv != "" { + // If backend is explicitly overridden to a DIFFERENT backend than what's + // currently configured, reset server[0] to backend-specific defaults + // first to avoid mixed config (e.g. lmstudio backend with Ollama + // host/model). Skip the reset when backendEnv just confirms the backend + // that's already configured (e.g. LUMEN_BACKEND=openai set redundantly + // alongside a config.yaml that already has backend: openai) — otherwise + // this would silently wipe an explicitly configured host/model. For + // ollama/lmstudio that wipe is masked by their localhost defaults (it + // just silently redirects a remote host back to localhost); for openai, + // which has no usable default host, it turns into a hard validation + // failure at startup. SkipHealthCheck is always carried forward across an + // actual backend switch — it's a harmless no-op for any backend. APIKey is + // only carried forward when switching INTO openai: validate() rejects a + // non-empty api_key on any other backend, so carrying it into e.g. ollama + // would turn a same-invocation `LUMEN_BACKEND=ollama` override (with an + // openai server left configured in config.yaml) into a hard startup + // failure instead of the intended backend switch. + if backendEnv != "" && backendEnv != srv.Backend { + prevAPIKey, prevSkipHealthCheck := srv.APIKey, srv.SkipHealthCheck srv = defaultServerForBackend(backendEnv) + if backendEnv == BackendOpenAI { + srv.APIKey = prevAPIKey + } + srv.SkipHealthCheck = prevSkipHealthCheck } if model != "" { @@ -272,12 +300,30 @@ func applyEnvOverrides(k *koanf.Koanf) { if lmStudioHost != "" { srv.Host = lmStudioHost } + case BackendOpenAI: + if openAIHost != "" { + srv.Host = openAIHost + } default: if ollamaHost != "" { srv.Host = ollamaHost } } + // OPENAI_API_KEY only applies to the openai backend — otherwise an + // unrelated OPENAI_API_KEY left in the environment (e.g. for another + // tool) would silently attach itself to an ollama/lmstudio server + // config, where it's never sent and would trip the api_key validation + // below for no reason. + if openAIKey != "" && selectedBackend == BackendOpenAI { + srv.APIKey = openAIKey + } + if skipHealthCheck != "" { + if b, err := strconv.ParseBool(skipHealthCheck); err == nil { + srv.SkipHealthCheck = b + } + } + if dims != "" { if n, err := strconv.Atoi(dims); err == nil { srv.Dims = n @@ -291,14 +337,23 @@ func applyEnvOverrides(k *koanf.Koanf) { servers[0] = srv // Re-marshal servers back into koanf - serverMaps := make([]map[string]any, len(servers)) + _ = k.Load(confmap.Provider(map[string]any{"servers": serverConfigMaps(servers)}, "."), nil) +} + +// serverConfigMaps converts ServerConfig entries into the map literal shape +// koanf expects. All three call sites that rebuild the servers list (model +// override, server-selection filter, env override) must go through this +// helper so new fields aren't silently dropped on one of the paths. +func serverConfigMaps(servers []ServerConfig) []map[string]any { + out := make([]map[string]any, len(servers)) for i, s := range servers { - serverMaps[i] = map[string]any{ + out[i] = map[string]any{ "backend": s.Backend, "host": s.Host, "model": s.Model, "dims": s.Dims, "ctx_length": s.CtxLength, "min_score": s.MinScore, + "api_key": s.APIKey, "skip_health_check": s.SkipHealthCheck, } } - _ = k.Load(confmap.Provider(map[string]any{"servers": serverMaps}, "."), nil) + return out } func (s *ConfigService) MaxChunkTokens() int { @@ -436,6 +491,16 @@ func (s *ConfigService) ServersForModel(model string) ([]int, error) { return indices, nil } +// isLoopbackHost reports whether host is "localhost" or a loopback IP, +// exempting local dev/test gateways from the api_key-requires-https rule. +func isLoopbackHost(host string) bool { + if host == "localhost" { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + // validate checks the current configuration for correctness. // Must be called either on a ConfigService that has not yet been shared with // other goroutines, or on a temporary ConfigService (as in reload). Must not @@ -453,7 +518,7 @@ func (s *ConfigService) validate() error { if srv.Backend == "" { return fmt.Errorf("config: servers[%d]: backend is required", i) } - if srv.Backend != BackendOllama && srv.Backend != BackendLMStudio { + if srv.Backend != BackendOllama && srv.Backend != BackendLMStudio && srv.Backend != BackendOpenAI { return fmt.Errorf("config: servers[%d]: unknown backend %q", i, srv.Backend) } if srv.Model == "" { @@ -466,6 +531,12 @@ func (s *ConfigService) validate() error { if err != nil || (u.Scheme != "http" && u.Scheme != "https") { return fmt.Errorf("config: servers[%d]: host %q must be a valid http/https URL", i, srv.Host) } + if srv.APIKey != "" && srv.Backend != BackendOpenAI { + return fmt.Errorf("config: servers[%d]: api_key is only supported for the %q backend, got %q", i, BackendOpenAI, srv.Backend) + } + if srv.APIKey != "" && u.Scheme != "https" && !isLoopbackHost(u.Hostname()) { + return fmt.Errorf("config: servers[%d]: api_key requires an https host (or loopback for local testing), got %q", i, srv.Host) + } if s.serverDims(i) == 0 { return fmt.Errorf("config: servers[%d]: cannot resolve dims for model %q — set dims explicitly", i, srv.Model) } diff --git a/internal/config/service_test.go b/internal/config/service_test.go index 38af51d..f1069cc 100644 --- a/internal/config/service_test.go +++ b/internal/config/service_test.go @@ -146,6 +146,256 @@ func TestEnvServerMapping_LMStudio(t *testing.T) { } } +func TestEnvServerMapping_OpenAI(t *testing.T) { + for _, k := range []string{"LUMEN_EMBED_CTX", "OLLAMA_HOST", "LM_STUDIO_HOST"} { + t.Setenv(k, "") + } + t.Setenv("LUMEN_BACKEND", "openai") + t.Setenv("OPENAI_BASE_URL", "https://api.example.com") + t.Setenv("OPENAI_API_KEY", "sk-test-123") + t.Setenv("LUMEN_EMBED_MODEL", "text-embedding-3-small") + t.Setenv("LUMEN_EMBED_DIMS", "1536") + svc, err := NewConfigService("") + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + s := svc.Servers()[0] + if s.Backend != BackendOpenAI { + t.Errorf("Backend = %q, want %q", s.Backend, BackendOpenAI) + } + if s.Host != "https://api.example.com" { + t.Errorf("Host = %q, want https://api.example.com", s.Host) + } + if s.APIKey != "sk-test-123" { + t.Errorf("APIKey = %q, want sk-test-123", s.APIKey) + } +} + +func TestEnvServerMapping_SkipHealthCheck(t *testing.T) { + for _, k := range []string{"LUMEN_EMBED_CTX", "OLLAMA_HOST", "LM_STUDIO_HOST", "OPENAI_API_KEY"} { + t.Setenv(k, "") + } + t.Setenv("LUMEN_BACKEND", "openai") + t.Setenv("OPENAI_BASE_URL", "https://api.example.com") + t.Setenv("LUMEN_EMBED_MODEL", "text-embedding-3-small") + t.Setenv("LUMEN_EMBED_DIMS", "1536") + t.Setenv("LUMEN_EMBED_SKIP_HEALTH_CHECK", "true") + svc, err := NewConfigService("") + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + if !svc.Servers()[0].SkipHealthCheck { + t.Error("SkipHealthCheck = false, want true") + } +} + +func TestValidation_OpenAIBackendAccepted(t *testing.T) { + for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX"} { + t.Setenv(k, "") + } + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(`servers: [{backend: openai, host: "https://api.example.com", model: text-embedding-3-small, dims: 1536}]`), 0644) + if _, err := NewConfigService(f); err != nil { + t.Fatalf("NewConfigService: %v", err) + } +} + +// TestAPIKeyAndSkipHealthCheck_SurviveModelOverride guards against the +// serverMaps rebuild in the WithModelOverride path silently dropping +// api_key/skip_health_check — every serverMaps literal must carry both fields. +func TestAPIKeyAndSkipHealthCheck_SurviveModelOverride(t *testing.T) { + for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX", "OPENAI_API_KEY", "OPENAI_BASE_URL"} { + t.Setenv(k, "") + } + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-yaml-key + skip_health_check: true +`), 0644) + svc, err := NewConfigService(f, WithModelOverride("text-embedding-3-large")) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + s := svc.Servers()[0] + if s.APIKey != "sk-yaml-key" { + t.Errorf("APIKey = %q, want sk-yaml-key (dropped by WithModelOverride rebuild)", s.APIKey) + } + if !s.SkipHealthCheck { + t.Error("SkipHealthCheck = false, want true (dropped by WithModelOverride rebuild)") + } +} + +// TestAPIKeyAndSkipHealthCheck_SurviveServerSelection guards the same rebuild +// in the WithServerSelection filter path. +func TestAPIKeyAndSkipHealthCheck_SurviveServerSelection(t *testing.T) { + clearServerEnv(t) + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-yaml-key + skip_health_check: true +`), 0644) + svc, err := NewConfigService(f, WithServerSelection("text-embedding-3-small", "")) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + s := svc.Servers()[0] + if s.APIKey != "sk-yaml-key" { + t.Errorf("APIKey = %q, want sk-yaml-key (dropped by WithServerSelection rebuild)", s.APIKey) + } + if !s.SkipHealthCheck { + t.Error("SkipHealthCheck = false, want true (dropped by WithServerSelection rebuild)") + } +} + +// TestAPIKeyAndSkipHealthCheck_SurviveReload guards the reload() path, which +// reuses applyEnvOverrides and must not drop api_key/skip_health_check either. +func TestAPIKeyAndSkipHealthCheck_SurviveReload(t *testing.T) { + clearServerEnv(t) + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-yaml-key + skip_health_check: true +`), 0644) + svc, err := NewConfigService(f) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + svc.reload() + s := svc.Servers()[0] + if s.APIKey != "sk-yaml-key" { + t.Errorf("APIKey = %q, want sk-yaml-key (dropped on reload)", s.APIKey) + } + if !s.SkipHealthCheck { + t.Error("SkipHealthCheck = false, want true (dropped on reload)") + } +} + +// TestAPIKeyAndSkipHealthCheck_SurviveBackendEnvReset guards against +// applyEnvOverrides' backend-switch reset: when LUMEN_BACKEND is set, +// server[0] is rebuilt from defaultServerForBackend to avoid mixed host/model +// config (e.g. a stale Ollama host surviving a switch to lmstudio). Host, +// Model and Dims are expected to be re-supplied via their own env vars in +// that case (as the README's env-var example does) — but api_key has no env +// var in this codepath's "supply everything via env" pattern other than +// OPENAI_API_KEY, so a deployment that intentionally keeps only the secret +// in config.yaml (env vars for everything else, api_key from a mounted +// file/secret) must not have it silently dropped by the reset. +func TestAPIKeyAndSkipHealthCheck_SurviveBackendEnvReset(t *testing.T) { + clearServerEnv(t) + t.Setenv("OPENAI_API_KEY", "") + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-yaml-key + skip_health_check: true +`), 0644) + // Host/Model/Dims are re-supplied via env, matching the documented + // "Or via environment variables" pattern — only api_key is left to come + // from the YAML file's secret. + t.Setenv("LUMEN_BACKEND", "openai") + t.Setenv("OPENAI_BASE_URL", "https://api.example.com") + t.Setenv("LUMEN_EMBED_MODEL", "text-embedding-3-small") + t.Setenv("LUMEN_EMBED_DIMS", "1536") + svc, err := NewConfigService(f) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + s := svc.Servers()[0] + if s.APIKey != "sk-yaml-key" { + t.Errorf("APIKey = %q, want sk-yaml-key (dropped by LUMEN_BACKEND reset)", s.APIKey) + } + if !s.SkipHealthCheck { + t.Error("SkipHealthCheck = false, want true (dropped by LUMEN_BACKEND reset)") + } +} + +// TestBackendEnvSwitch_AwayFromOpenAI_DropsAPIKey covers an actual backend +// switch (unlike TestAPIKeyAndSkipHealthCheck_SurviveBackendEnvReset above, +// where LUMEN_BACKEND merely confirms the backend already configured in +// config.yaml and the reset branch never runs). Here config.yaml configures +// an openai server with api_key set, and LUMEN_BACKEND=ollama switches to a +// genuinely different backend for this invocation — the documented +// "Selecting a server per invocation" pattern. validate() rejects a non-empty +// api_key on any backend other than openai, so the carried-over API key must +// not survive the switch, or NewConfigService fails outright. +func TestBackendEnvSwitch_AwayFromOpenAI_DropsAPIKey(t *testing.T) { + clearServerEnv(t) + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-yaml-key +`), 0644) + t.Setenv("LUMEN_BACKEND", "ollama") + svc, err := NewConfigService(f) + if err != nil { + t.Fatalf("NewConfigService: %v (a carried-over api_key must not fail validation on a backend switch)", err) + } + s := svc.Servers()[0] + if s.Backend != BackendOllama { + t.Errorf("Backend = %q, want %q", s.Backend, BackendOllama) + } + if s.APIKey != "" { + t.Errorf("APIKey = %q, want empty (api_key must not carry over to a non-openai backend)", s.APIKey) + } +} + +// TestBackendEnvSwitch_UnrecognizedValueErrors guards defaultServerForBackend +// against silently normalizing a mistyped LUMEN_BACKEND value to Ollama. +// Before this test's fix, defaultServerForBackend's `default:` arm returned +// Ollama defaults for ANY string it didn't recognize as lmstudio or openai — +// including a case-typo like "OpenAI" instead of "openai" — which made +// NewConfigService succeed with backend silently switched to ollama (and, +// with an openai server configured, its api_key silently dropped) instead of +// reporting an error for the unrecognized backend. +func TestBackendEnvSwitch_UnrecognizedValueErrors(t *testing.T) { + clearServerEnv(t) + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(` +servers: + - backend: openai + host: https://api.example.com + model: text-embedding-3-small + dims: 1536 + api_key: sk-prod-secret +`), 0644) + // Case typo: "OpenAI" instead of "openai". + t.Setenv("LUMEN_BACKEND", "OpenAI") + _, err := NewConfigService(f) + if err == nil { + t.Fatal("NewConfigService: want error for unrecognized LUMEN_BACKEND value, got nil (silently fell back to another backend)") + } +} + func TestHostConflict_BothSet(t *testing.T) { for _, k := range []string{"LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX"} { t.Setenv(k, "") @@ -345,6 +595,18 @@ func TestValidation_InvalidHost(t *testing.T) { } } +func TestValidation_APIKeyRequiresHTTPS(t *testing.T) { + for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "OPENAI_API_KEY", "OPENAI_BASE_URL", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX"} { + t.Setenv(k, "") + } + dir := t.TempDir() + f := filepath.Join(dir, "config.yaml") + _ = os.WriteFile(f, []byte(`servers: [{backend: openai, host: "http://insecure:1234", model: text-embedding-3-small, dims: 1536, api_key: sk-test}]`), 0644) + if _, err := NewConfigService(f); err == nil { + t.Fatal("expected error: api_key over plain http must fail validation") + } +} + func TestValidation_MissingModel(t *testing.T) { for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX"} { t.Setenv(k, "") @@ -457,7 +719,10 @@ func writeThreeServerYAML(t *testing.T) string { func clearServerEnv(t *testing.T) { t.Helper() - for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX"} { + for _, k := range []string{ + "LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "OLLAMA_HOST", "LM_STUDIO_HOST", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX", + "OPENAI_API_KEY", "OPENAI_BASE_URL", "LUMEN_EMBED_SKIP_HEALTH_CHECK", + } { t.Setenv(k, "") } } diff --git a/internal/embedder/failover.go b/internal/embedder/failover.go index 26e48c0..2a95c86 100644 --- a/internal/embedder/failover.go +++ b/internal/embedder/failover.go @@ -19,6 +19,7 @@ import ( "errors" "fmt" "log/slog" + "net/http" "sync" "time" @@ -139,6 +140,12 @@ func (f *FailoverEmbedder) Embed(ctx context.Context, texts []string) ([][]float return result, nil } + if ctxErr := ctx.Err(); ctxErr != nil { + // Caller cancellation or deadline expiration, not a server + // problem — don't mark the active server unhealthy or failover. + return nil, ctxErr + } + if !isTransientError(err) { return nil, err // 4xx = config error, don't failover } @@ -173,12 +180,14 @@ func (f *FailoverEmbedder) serversChanged() bool { return true } for i, srv := range current { - // Compare key fields — if any differ, servers have changed. if f.servers[i].emb == nil { continue // not initialized yet, can't compare } - cached := f.cachedConfigs[i] - if srv.Backend != cached.Backend || srv.Host != cached.Host || srv.Model != cached.Model { + // Compare the whole struct (all fields are comparable) so that any + // config change — including ones added after this comparison was + // first written, like APIKey or SkipHealthCheck — forces a + // re-init instead of silently keeping a stale cached embedder. + if f.cachedConfigs[i] != srv { return true } } @@ -239,6 +248,9 @@ func (f *FailoverEmbedder) probeHealth(ctx context.Context, i int) bool { return false } srv := servers[i] + if srv.SkipHealthCheck { + return true + } probeCtx, cancel := context.WithTimeout(ctx, healthCheckTimeout) defer cancel() if err := ProbeServer(probeCtx, srv); err != nil { @@ -268,6 +280,8 @@ func (f *FailoverEmbedder) ensureEmbedder(i int) error { emb, err = NewOllama(srv.Model, dims, ctxLen, srv.Host) case "lmstudio": emb, err = NewLMStudio(srv.Model, dims, srv.Host) + case "openai": + emb, err = NewOpenAI(srv.Model, dims, srv.Host, srv.APIKey) default: return fmt.Errorf("unknown backend %q", srv.Backend) } @@ -279,12 +293,13 @@ func (f *FailoverEmbedder) ensureEmbedder(i int) error { } // isTransientError returns true if the error represents a transient failure -// that warrants failover (5xx HTTP errors or network errors). Returns false -// for 4xx errors (configuration errors, no failover). +// that warrants failover (5xx HTTP errors, 429 rate limiting, or network +// errors). Returns false for other 4xx errors (configuration errors, no +// failover). func isTransientError(err error) bool { var ee *EmbedError if errors.As(err, &ee) { - return ee.StatusCode >= 500 + return ee.StatusCode >= 500 || ee.StatusCode == http.StatusTooManyRequests } return true // network errors are transient } diff --git a/internal/embedder/failover_test.go b/internal/embedder/failover_test.go index 78f6fa3..ca13b5e 100644 --- a/internal/embedder/failover_test.go +++ b/internal/embedder/failover_test.go @@ -68,11 +68,14 @@ func testConfigService(t *testing.T, servers ...config.ServerConfig) *config.Con t.Setenv("LUMEN_EMBED_CTX", "") t.Setenv("OLLAMA_HOST", "") t.Setenv("LM_STUDIO_HOST", "") + t.Setenv("OPENAI_API_KEY", "") + t.Setenv("OPENAI_BASE_URL", "") + t.Setenv("LUMEN_EMBED_SKIP_HEALTH_CHECK", "") y := "servers:\n" for _, s := range servers { - y += fmt.Sprintf(" - backend: %s\n host: %s\n model: %s\n dims: %d\n", - s.Backend, s.Host, s.Model, s.Dims) + y += fmt.Sprintf(" - backend: %s\n host: %s\n model: %s\n dims: %d\n api_key: %q\n skip_health_check: %t\n", + s.Backend, s.Host, s.Model, s.Dims, s.APIKey, s.SkipHealthCheck) } dir := t.TempDir() cfgFile := filepath.Join(dir, "config.yaml") @@ -329,6 +332,152 @@ func TestFailover_ReloadPicksUpNewServers(t *testing.T) { } } +func newTestOpenAIServer(t *testing.T, healthy bool, embedStatus int, wantAuth string) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if wantAuth != "" && r.Header.Get("Authorization") != wantAuth { + t.Errorf("unexpected Authorization header: %q, want %q", r.Header.Get("Authorization"), wantAuth) + } + switch { + case r.Method == "GET" && r.URL.Path == "/v1/models": + if healthy { + w.WriteHeader(200) + _, _ = fmt.Fprint(w, `{"data":[{"id":"test-openai"}]}`) + } else { + w.WriteHeader(503) + } + case r.Method == "POST" && r.URL.Path == "/v1/embeddings": + w.WriteHeader(embedStatus) + if embedStatus == 200 { + _, _ = fmt.Fprint(w, `{"data":[{"embedding":[0.1,0.2,0.3]}]}`) + } else { + _, _ = fmt.Fprintf(w, `{"error":"status %d"}`, embedStatus) + } + default: + w.WriteHeader(404) + } + })) +} + +func TestFailover_OpenAIBackend_Healthy(t *testing.T) { + srv := newTestOpenAIServer(t, true, 200, "Bearer sk-test") + defer srv.Close() + + cfg := testConfigService(t, + config.ServerConfig{Backend: "openai", Host: srv.URL, Model: "test-openai", Dims: 3, APIKey: "sk-test"}, + ) + fe := NewFailoverEmbedder(cfg) + _, err := fe.Embed(context.Background(), []string{"hello"}) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if fe.ActiveServerIndex() != 0 { + t.Errorf("active = %d, want 0", fe.ActiveServerIndex()) + } +} + +func TestFailover_OpenAIBackend_SkipHealthCheck(t *testing.T) { + // The server always returns 503 for /v1/models; if the health probe were + // actually invoked, this server would never be selected. SkipHealthCheck + // must bypass the probe entirely and go straight to embedding. + callCount := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/models" { + callCount++ + w.WriteHeader(503) + return + } + if r.Method == "POST" && r.URL.Path == "/v1/embeddings" { + _, _ = fmt.Fprint(w, `{"data":[{"embedding":[0.1,0.2,0.3]}]}`) + return + } + w.WriteHeader(404) + })) + defer srv.Close() + + cfg := testConfigService(t, + config.ServerConfig{Backend: "openai", Host: srv.URL, Model: "test-openai", Dims: 3, SkipHealthCheck: true}, + ) + fe := NewFailoverEmbedder(cfg) + _, err := fe.Embed(context.Background(), []string{"hello"}) + if err != nil { + t.Fatalf("Embed should succeed with skip_health_check, got: %v", err) + } + if callCount != 0 { + t.Errorf("expected /v1/models to never be called with skip_health_check, got %d calls", callCount) + } +} + +// TestFailover_ReloadPicksUpAPIKeyChange guards against serversChanged() +// comparing only backend/host/model: rotating api_key or skip_health_check +// via a config hot reload — with backend/host/model unchanged — must still +// force re-initialization, or the running FailoverEmbedder keeps using a +// stale embedder built with the old API key indefinitely. +func TestFailover_ReloadPicksUpAPIKeyChange(t *testing.T) { + for _, k := range []string{"LUMEN_BACKEND", "LUMEN_EMBED_MODEL", "LUMEN_EMBED_DIMS", "LUMEN_EMBED_CTX", "OLLAMA_HOST", "LM_STUDIO_HOST", "OPENAI_API_KEY", "OPENAI_BASE_URL"} { + t.Setenv(k, "") + } + + var lastAuth atomic.Value + lastAuth.Store("") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v1/embeddings" { + lastAuth.Store(r.Header.Get("Authorization")) + _, _ = fmt.Fprint(w, `{"data":[{"embedding":[0.1,0.2,0.3]}]}`) + return + } + w.WriteHeader(404) + })) + defer srv.Close() + + dir := t.TempDir() + cfgFile := filepath.Join(dir, "config.yaml") + writeCfg := func(apiKey string) { + t.Helper() + content := fmt.Sprintf(` +servers: + - backend: openai + host: %s + model: test-openai + dims: 3 + skip_health_check: true + api_key: %s +`, srv.URL, apiKey) + if err := os.WriteFile(cfgFile, []byte(content), 0644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + } + writeCfg("old-key") + + cfg, err := config.NewConfigService(cfgFile) + if err != nil { + t.Fatalf("NewConfigService: %v", err) + } + if err := cfg.Watch(); err != nil { + t.Fatalf("Watch: %v", err) + } + defer cfg.Stop() + + fe := NewFailoverEmbedder(cfg) + if _, err := fe.Embed(context.Background(), []string{"hello"}); err != nil { + t.Fatalf("Embed: %v", err) + } + if got := lastAuth.Load().(string); got != "Bearer old-key" { + t.Fatalf("Authorization = %q, want Bearer old-key", got) + } + + // Hot reload rotates only api_key — backend/host/model are unchanged. + writeCfg("new-key") + time.Sleep(500 * time.Millisecond) + + if _, err := fe.Embed(context.Background(), []string{"hello"}); err != nil { + t.Fatalf("Embed after reload: %v", err) + } + if got := lastAuth.Load().(string); got != "Bearer new-key" { + t.Fatalf("Authorization after reload = %q, want Bearer new-key (api_key rotation must force re-init)", got) + } +} + func newTestLMStudioServer(t *testing.T, healthy bool, embedStatus int) *httptest.Server { t.Helper() return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/internal/embedder/health.go b/internal/embedder/health.go index 3f3fde2..5e67ae6 100644 --- a/internal/embedder/health.go +++ b/internal/embedder/health.go @@ -30,14 +30,17 @@ import ( // as a healthy failover target. func ProbeServer(ctx context.Context, srv config.ServerConfig) error { endpoint := strings.TrimRight(srv.Host, "/") + "/api/tags" - if srv.Backend == config.BackendLMStudio { - endpoint = strings.TrimRight(srv.Host, "/") + "/v1/models" + if srv.Backend == config.BackendLMStudio || srv.Backend == config.BackendOpenAI { + endpoint = normalizeBaseURL(srv.Host) + "/v1/models" } req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return fmt.Errorf("create health request: %w", err) } + if srv.Backend == config.BackendOpenAI && srv.APIKey != "" { + req.Header.Set("Authorization", "Bearer "+srv.APIKey) + } resp, err := http.DefaultClient.Do(req) if err != nil { return fmt.Errorf("service unreachable: %w", err) @@ -63,7 +66,7 @@ func ProbeServer(ctx context.Context, srv config.ServerConfig) error { for _, model := range body.Models { available = append(available, model.Name, model.Model) } - case config.BackendLMStudio: + case config.BackendLMStudio, config.BackendOpenAI: var body struct { Data []struct { ID string `json:"id"` diff --git a/internal/embedder/health_test.go b/internal/embedder/health_test.go new file mode 100644 index 0000000..73796fc --- /dev/null +++ b/internal/embedder/health_test.go @@ -0,0 +1,39 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package embedder + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/ory/lumen/internal/config" +) + +func TestProbeServer_OpenAI_TrailingV1InHostNotDoubled(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("expected path /v1/models, got %q (host's /v1 must not be duplicated)", r.URL.Path) + } + _, _ = w.Write([]byte(`{"data":[{"id":"test-model"}]}`)) + })) + defer server.Close() + + srv := config.ServerConfig{Backend: config.BackendOpenAI, Host: server.URL + "/v1", Model: "test-model"} + if err := ProbeServer(context.Background(), srv); err != nil { + t.Fatalf("ProbeServer: %v", err) + } +} diff --git a/internal/embedder/lmstudio.go b/internal/embedder/lmstudio.go index d0d49fd..95e6faf 100644 --- a/internal/embedder/lmstudio.go +++ b/internal/embedder/lmstudio.go @@ -14,142 +14,19 @@ package embedder -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "slices" - "time" - - "github.com/sethvargo/go-retry" -) - // LMStudio implements the Embedder interface using an LM Studio server -// that exposes an OpenAI-compatible /v1/embeddings endpoint. +// that exposes an OpenAI-compatible /v1/embeddings endpoint. It delegates to +// OpenAI, which implements the same wire format, with no API key. type LMStudio struct { - model string - dimensions int - baseURL string - client *http.Client + *OpenAI } // NewLMStudio creates a new LMStudio embedder. // baseURL is the LM Studio server URL (e.g. "http://localhost:1234"). func NewLMStudio(model string, dimensions int, baseURL string) (*LMStudio, error) { - return &LMStudio{ - model: model, - dimensions: dimensions, - baseURL: baseURL, - client: &http.Client{ - Timeout: 10 * time.Minute, - }, - }, nil -} - -// Dimensions returns the embedding vector dimensionality. -func (l *LMStudio) Dimensions() int { - return l.dimensions -} - -// ModelName returns the model name used for embeddings. -func (l *LMStudio) ModelName() string { - return l.model -} - -// lmstudioEmbedRequest is the JSON body sent to /v1/embeddings. -type lmstudioEmbedRequest struct { - Model string `json:"model"` - Input []string `json:"input"` -} - -// lmstudioEmbedItem is a single embedding item in the response. -type lmstudioEmbedItem struct { - Embedding []float32 `json:"embedding"` - Index int `json:"index"` -} - -// lmstudioEmbedResponse is the JSON body returned from /v1/embeddings. -type lmstudioEmbedResponse struct { - Data []lmstudioEmbedItem `json:"data"` -} - -// Embed converts texts into embedding vectors, splitting into batches of 32. -func (l *LMStudio) Embed(ctx context.Context, texts []string) ([][]float32, error) { - if len(texts) == 0 { - return nil, nil - } - - var allVecs [][]float32 - for i := 0; i < len(texts); i += embedBatchSize { - batch := texts[i:min(i+embedBatchSize, len(texts))] - - vecs, err := l.embedBatch(ctx, batch) - if err != nil { - return nil, fmt.Errorf("embedding batch starting at %d: %w", i, err) - } - allVecs = append(allVecs, vecs...) - } - - return allVecs, nil -} - -// embedBatch sends a single batch of texts to the LM Studio /v1/embeddings endpoint. -// Retries up to embedMaxRetries times on transient errors (5xx, network failures), -// respecting context cancellation between attempts. -func (l *LMStudio) embedBatch(ctx context.Context, texts []string) ([][]float32, error) { - bodyBytes, err := json.Marshal(lmstudioEmbedRequest{ - Model: l.model, - Input: texts, - }) - if err != nil { - return nil, fmt.Errorf("marshalling request: %w", err) - } - - b := retry.NewExponential(100 * time.Millisecond) - - var embedResp lmstudioEmbedResponse - err = retry.Do(ctx, retry.WithMaxRetries(embedMaxRetries-1, b), func(ctx context.Context) error { - req, err := http.NewRequestWithContext(ctx, http.MethodPost, l.baseURL+"/v1/embeddings", bytes.NewReader(bodyBytes)) - if err != nil { - return fmt.Errorf("creating request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - - resp, err := l.client.Do(req) - if err != nil { - return retry.RetryableError(fmt.Errorf("request failed: %w", err)) - } - - body, readErr := io.ReadAll(resp.Body) - _ = resp.Body.Close() - - if resp.StatusCode >= 500 { - return retry.RetryableError(&EmbedError{StatusCode: resp.StatusCode, Message: string(body)}) - } - if resp.StatusCode != http.StatusOK { - return &EmbedError{StatusCode: resp.StatusCode, Message: string(body)} - } - if readErr != nil { - return fmt.Errorf("reading response body: %w", readErr) - } - - return json.Unmarshal(body, &embedResp) - }) + o, err := NewOpenAI(model, dimensions, baseURL, "") if err != nil { - return nil, fmt.Errorf("lmstudio embed: %w", err) - } - - // Sort by index — OpenAI spec allows out-of-order responses. - slices.SortFunc(embedResp.Data, func(a, b lmstudioEmbedItem) int { - return a.Index - b.Index - }) - - vecs := make([][]float32, len(embedResp.Data)) - for i, item := range embedResp.Data { - vecs[i] = item.Embedding + return nil, err } - return vecs, nil + return &LMStudio{OpenAI: o}, nil } diff --git a/internal/embedder/lmstudio_test.go b/internal/embedder/lmstudio_test.go index b484daf..4906097 100644 --- a/internal/embedder/lmstudio_test.go +++ b/internal/embedder/lmstudio_test.go @@ -23,12 +23,12 @@ import ( "time" ) -func makeLMStudioResponse(embeddings [][]float32) lmstudioEmbedResponse { - data := make([]lmstudioEmbedItem, len(embeddings)) +func makeLMStudioResponse(embeddings [][]float32) openaiEmbedResponse { + data := make([]openaiEmbedItem, len(embeddings)) for i, e := range embeddings { - data[i] = lmstudioEmbedItem{Embedding: e, Index: i} + data[i] = openaiEmbedItem{Embedding: e, Index: i} } - return lmstudioEmbedResponse{Data: data} + return openaiEmbedResponse{Data: data} } func TestLMStudioEmbedder_Embed(t *testing.T) { @@ -64,8 +64,8 @@ func TestLMStudioEmbedder_Embed(t *testing.T) { func TestLMStudioEmbedder_OrderingByIndex(t *testing.T) { // Mock returns items in reversed index order to verify sorting. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - resp := lmstudioEmbedResponse{ - Data: []lmstudioEmbedItem{ + resp := openaiEmbedResponse{ + Data: []openaiEmbedItem{ {Embedding: []float32{0.9, 0.9, 0.9, 0.9}, Index: 1}, {Embedding: []float32{0.1, 0.2, 0.3, 0.4}, Index: 0}, }, diff --git a/internal/embedder/openai.go b/internal/embedder/openai.go new file mode 100644 index 0000000..7213d9c --- /dev/null +++ b/internal/embedder/openai.go @@ -0,0 +1,193 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package embedder + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/sethvargo/go-retry" +) + +// OpenAI implements the Embedder interface using an OpenAI-compatible +// /v1/embeddings endpoint. It works with OpenAI itself and with any other +// service (internal gateways, proxies) exposing the same wire format. +type OpenAI struct { + model string + dimensions int + baseURL string + apiKey string + client *http.Client +} + +// NewOpenAI creates a new OpenAI-compatible embedder. +// baseURL is the API base URL (e.g. "https://api.openai.com", or +// "https://api.openai.com/v1" — a trailing "/v1" is stripped so callers +// following either convention don't end up with "/v1/v1/embeddings"). +// apiKey is the Bearer token for authentication; when empty, no +// Authorization header is sent, since some internal gateways trust network +// position rather than a token. +func NewOpenAI(model string, dimensions int, baseURL string, apiKey string) (*OpenAI, error) { + return &OpenAI{ + model: model, + dimensions: dimensions, + baseURL: normalizeBaseURL(baseURL), + apiKey: apiKey, + client: &http.Client{ + Timeout: 10 * time.Minute, + }, + }, nil +} + +// normalizeBaseURL strips a trailing slash and a trailing "/v1" so the +// embed/health-probe code can append "/v1/..." exactly once regardless of +// whether the configured host already includes the version prefix. +func normalizeBaseURL(raw string) string { + trimmed := strings.TrimRight(raw, "/") + trimmed = strings.TrimSuffix(trimmed, "/v1") + return strings.TrimRight(trimmed, "/") +} + +// Dimensions returns the embedding vector dimensionality. +func (o *OpenAI) Dimensions() int { + return o.dimensions +} + +// ModelName returns the model name used for embeddings. +func (o *OpenAI) ModelName() string { + return o.model +} + +// openaiEmbedRequest is the JSON body sent to /v1/embeddings. +type openaiEmbedRequest struct { + Model string `json:"model"` + Input []string `json:"input"` +} + +// openaiEmbedItem is a single embedding item in the response. +type openaiEmbedItem struct { + Embedding []float32 `json:"embedding"` + Index int `json:"index"` +} + +// openaiEmbedResponse is the JSON body returned from /v1/embeddings. +type openaiEmbedResponse struct { + Data []openaiEmbedItem `json:"data"` +} + +// Embed converts texts into embedding vectors, splitting into batches of 32. +func (o *OpenAI) Embed(ctx context.Context, texts []string) ([][]float32, error) { + if len(texts) == 0 { + return nil, nil + } + + var allVecs [][]float32 + for i := 0; i < len(texts); i += embedBatchSize { + batch := texts[i:min(i+embedBatchSize, len(texts))] + + vecs, err := o.embedBatch(ctx, batch) + if err != nil { + return nil, fmt.Errorf("embedding batch starting at %d: %w", i, err) + } + allVecs = append(allVecs, vecs...) + } + + return allVecs, nil +} + +// embedBatch sends a single batch of texts to the /v1/embeddings endpoint. +// Retries up to embedMaxRetries times on transient errors (5xx, 429 rate +// limits, network failures), respecting context cancellation between +// attempts. +func (o *OpenAI) embedBatch(ctx context.Context, texts []string) ([][]float32, error) { + bodyBytes, err := json.Marshal(openaiEmbedRequest{ + Model: o.model, + Input: texts, + }) + if err != nil { + return nil, fmt.Errorf("marshalling request: %w", err) + } + + b := retry.NewExponential(100 * time.Millisecond) + + var embedResp openaiEmbedResponse + err = retry.Do(ctx, retry.WithMaxRetries(embedMaxRetries-1, b), func(ctx context.Context) error { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/v1/embeddings", bytes.NewReader(bodyBytes)) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if o.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+o.apiKey) + } + + resp, err := o.client.Do(req) + if err != nil { + return retry.RetryableError(fmt.Errorf("request failed: %w", err)) + } + + body, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + + if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 { + return retry.RetryableError(&EmbedError{StatusCode: resp.StatusCode, Message: string(body)}) + } + if resp.StatusCode != http.StatusOK { + return &EmbedError{StatusCode: resp.StatusCode, Message: string(body)} + } + if readErr != nil { + return fmt.Errorf("reading response body: %w", readErr) + } + + return json.Unmarshal(body, &embedResp) + }) + if err != nil { + // Identify the failing server by base URL rather than a hardcoded + // "openai" label — this embedder also backs the LMStudio wrapper, so + // a static backend name here would mislabel LM Studio failures. + return nil, fmt.Errorf("embed request to %s: %w", o.baseURL, err) + } + + // The OpenAI spec allows out-of-order responses but guarantees one item + // per input, indexed 0..len(texts)-1. Validate that explicitly instead of + // trusting response order/length — a misbehaving gateway that drops, + // duplicates, or mis-sizes an item must fail loudly rather than silently + // misalign embeddings with their source texts. + if len(embedResp.Data) != len(texts) { + return nil, fmt.Errorf("embed response: got %d embeddings, want %d", len(embedResp.Data), len(texts)) + } + vecs := make([][]float32, len(texts)) + seen := make([]bool, len(texts)) + for _, item := range embedResp.Data { + if item.Index < 0 || item.Index >= len(texts) { + return nil, fmt.Errorf("embed response: index %d out of range [0,%d)", item.Index, len(texts)) + } + if seen[item.Index] { + return nil, fmt.Errorf("embed response: duplicate index %d", item.Index) + } + if len(item.Embedding) != o.dimensions { + return nil, fmt.Errorf("embed response: item %d has %d dimensions, want %d", item.Index, len(item.Embedding), o.dimensions) + } + seen[item.Index] = true + vecs[item.Index] = item.Embedding + } + return vecs, nil +} diff --git a/internal/embedder/openai_test.go b/internal/embedder/openai_test.go new file mode 100644 index 0000000..e041e48 --- /dev/null +++ b/internal/embedder/openai_test.go @@ -0,0 +1,250 @@ +// Copyright 2026 Aeneas Rekkas +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package embedder + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func makeOpenAIResponse(embeddings [][]float32) openaiEmbedResponse { + data := make([]openaiEmbedItem, len(embeddings)) + for i, e := range embeddings { + data[i] = openaiEmbedItem{Embedding: e, Index: i} + } + return openaiEmbedResponse{Data: data} +} + +func TestOpenAIEmbedder_Embed(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + t.Fatalf("unexpected path: %s", r.URL.Path) + } + resp := makeOpenAIResponse([][]float32{ + {0.1, 0.2, 0.3, 0.4}, + {0.5, 0.6, 0.7, 0.8}, + }) + _ = json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + e, err := NewOpenAI("text-embedding-3-small", 4, server.URL, "") + if err != nil { + t.Fatal(err) + } + + vecs, err := e.Embed(context.Background(), []string{"hello", "world"}) + if err != nil { + t.Fatal(err) + } + if len(vecs) != 2 { + t.Fatalf("expected 2 vectors, got %d", len(vecs)) + } + if len(vecs[0]) != 4 { + t.Fatalf("expected 4 dimensions, got %d", len(vecs[0])) + } +} + +func TestOpenAIEmbedder_OrderingByIndex(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + resp := openaiEmbedResponse{ + Data: []openaiEmbedItem{ + {Embedding: []float32{0.9, 0.9, 0.9, 0.9}, Index: 1}, + {Embedding: []float32{0.1, 0.2, 0.3, 0.4}, Index: 0}, + }, + } + _ = json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + e, _ := NewOpenAI("text-embedding-3-small", 4, server.URL, "") + vecs, err := e.Embed(context.Background(), []string{"first", "second"}) + if err != nil { + t.Fatal(err) + } + if vecs[0][0] != 0.1 { + t.Fatalf("expected vecs[0][0]=0.1 (index:0 item), got %v", vecs[0][0]) + } + if vecs[1][0] != 0.9 { + t.Fatalf("expected vecs[1][0]=0.9 (index:1 item), got %v", vecs[1][0]) + } +} + +func TestOpenAIEmbedder_Batching(t *testing.T) { + callCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + var req map[string]any + _ = json.NewDecoder(r.Body).Decode(&req) + input := req["input"].([]any) + + embeddings := make([][]float32, len(input)) + for i := range input { + embeddings[i] = []float32{0.1, 0.2, 0.3, 0.4} + } + _ = json.NewEncoder(w).Encode(makeOpenAIResponse(embeddings)) + })) + defer server.Close() + + e, _ := NewOpenAI("text-embedding-3-small", 4, server.URL, "") + texts := make([]string, 50) + for i := range texts { + texts[i] = "text" + } + + vecs, err := e.Embed(context.Background(), texts) + if err != nil { + t.Fatal(err) + } + if len(vecs) != 50 { + t.Fatalf("expected 50 vectors, got %d", len(vecs)) + } + if callCount != 2 { + t.Fatalf("expected 2 batch calls (32+18), got %d", callCount) + } +} + +func TestOpenAIEmbedder_Dimensions(t *testing.T) { + e, _ := NewOpenAI("text-embedding-3-small", 768, "http://localhost:8080", "") + if e.Dimensions() != 768 { + t.Fatalf("expected 768, got %d", e.Dimensions()) + } +} + +func TestOpenAIEmbedder_ModelName(t *testing.T) { + e, _ := NewOpenAI("text-embedding-3-small", 768, "http://localhost:8080", "") + if e.ModelName() != "text-embedding-3-small" { + t.Fatalf("expected text-embedding-3-small, got %s", e.ModelName()) + } +} + +func TestOpenAIEmbedder_ErrorHandling(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + + e, _ := NewOpenAI("text-embedding-3-small", 4, server.URL, "") + _, err := e.Embed(context.Background(), []string{"hello"}) + if err == nil { + t.Fatal("expected error for 500 response") + } +} + +func TestOpenAI_Embed_ContextCancelledStopsRetry(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + emb, _ := NewOpenAI("text-embedding-3-small", 4, srv.URL, "") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + start := time.Now() + _, err := emb.Embed(ctx, []string{"hello"}) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected error from cancelled context") + } + if elapsed > 500*time.Millisecond { + t.Fatalf("expected fast failure on pre-cancelled context, took %v", elapsed) + } +} + +func TestOpenAIEmbedder_NoAuthHeaderWhenKeyEmpty(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if auth := r.Header.Get("Authorization"); auth != "" { + t.Errorf("expected no Authorization header, got %q", auth) + } + _ = json.NewEncoder(w).Encode(makeOpenAIResponse([][]float32{{0.1, 0.2}})) + })) + defer server.Close() + + e, err := NewOpenAI("m", 2, server.URL, "") + if err != nil { + t.Fatal(err) + } + if _, err := e.Embed(context.Background(), []string{"hello"}); err != nil { + t.Fatal(err) + } +} + +func TestOpenAIEmbedder_BearerHeaderWhenKeySet(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if auth := r.Header.Get("Authorization"); auth != "Bearer sk-test-key" { + t.Errorf("expected Bearer sk-test-key, got %q", auth) + } + _ = json.NewEncoder(w).Encode(makeOpenAIResponse([][]float32{{0.1, 0.2}})) + })) + defer server.Close() + + e, err := NewOpenAI("m", 2, server.URL, "sk-test-key") + if err != nil { + t.Fatal(err) + } + if _, err := e.Embed(context.Background(), []string{"hello"}); err != nil { + t.Fatal(err) + } +} + +func TestOpenAIEmbedder_TrailingV1InBaseURLNotDoubled(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/embeddings" { + t.Errorf("expected path /v1/embeddings, got %q (base URL's /v1 must not be duplicated)", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(makeOpenAIResponse([][]float32{{0.1, 0.2}})) + })) + defer server.Close() + + e, err := NewOpenAI("m", 2, server.URL+"/v1", "") + if err != nil { + t.Fatal(err) + } + if _, err := e.Embed(context.Background(), []string{"hello"}); err != nil { + t.Fatal(err) + } +} + +func TestOpenAIEmbedder_RetriesOn429(t *testing.T) { + attempts := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts++ + if attempts < 2 { + w.WriteHeader(http.StatusTooManyRequests) + return + } + _ = json.NewEncoder(w).Encode(makeOpenAIResponse([][]float32{{0.1, 0.2}})) + })) + defer server.Close() + + e, _ := NewOpenAI("m", 2, server.URL, "") + vecs, err := e.Embed(context.Background(), []string{"hello"}) + if err != nil { + t.Fatalf("expected retry to succeed after 429, got: %v", err) + } + if len(vecs) != 1 { + t.Fatalf("expected 1 vector, got %d", len(vecs)) + } + if attempts < 2 { + t.Fatalf("expected at least 2 attempts, got %d", attempts) + } +}