From 2ce2ca4409d6ef0907a85a2c1d95961d084d58cf Mon Sep 17 00:00:00 2001 From: Molty Date: Wed, 2 Sep 2026 19:29:56 -0700 Subject: [PATCH] fix(embed): fall back to embeddings probe when models route is missing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OpenAI-compatible servers were probed for availability via GET /models and reported unavailable when that route did not exist. Some compatible providers serve embeddings but no models route at all — Voyage AI returns 404 there — leaving mnemon silently without semantic recall against a working endpoint. When the models route is missing (404/405/501), probe availability with a real embeddings round-trip instead, so availability reflects the endpoint the client actually depends on. Auth, quota, and server errors still report unavailable. The round-trip is shared with Embed via embedWithContext/decodeEmbedResponse so the probe and the real call cannot drift apart. Verified end-to-end against the Voyage AI API (voyage-3.5) and against a local Ollama instance. --- README.md | 14 ++- internal/memory/embed/ollama.go | 128 ++++++++++++++++++--------- internal/memory/embed/openai_test.go | 86 ++++++++++++++++++ 3 files changed, 185 insertions(+), 43 deletions(-) diff --git a/README.md b/README.md index a24f616c..96ff213e 100644 --- a/README.md +++ b/README.md @@ -468,7 +468,10 @@ reported by ID in the triggering command's `auto_pruned_ids` field. The embedding client speaks the Ollama API by default and the OpenAI-compatible embeddings API when the endpoint ends in `/v1` (or when -`MNEMON_EMBED_PROTOCOL=openai` is set). For example, a local server such as +`MNEMON_EMBED_PROTOCOL=openai` is set). OpenAI-compatible servers are +normally probed via their `models` route; servers that do not serve that +route (e.g. [Voyage AI](https://docs.voyageai.com)) are detected via an +embeddings round-trip instead. For example, a local server such as [oMLX](https://omlx.dev) can be configured with: ```bash @@ -478,6 +481,15 @@ export MNEMON_EMBED_API_KEY=sk-... # omit for keyless local servers mnemon embed --status ``` +A hosted provider such as Voyage AI needs only the endpoint, model, and key: + +```bash +export MNEMON_EMBED_ENDPOINT=https://api.voyageai.com/v1 +export MNEMON_EMBED_MODEL=voyage-3.5 +export MNEMON_EMBED_API_KEY=pa-... +mnemon embed --status +``` + ## Development ```bash diff --git a/internal/memory/embed/ollama.go b/internal/memory/embed/ollama.go index ff2f5a55..f1aef461 100644 --- a/internal/memory/embed/ollama.go +++ b/internal/memory/embed/ollama.go @@ -132,9 +132,15 @@ func (c *Client) endpointURL(route string) (string, error) { return endpointURL, nil } -// Available returns true if the embedding server's discovery endpoint -// responds successfully. Uses a 2s timeout to avoid blocking the CLI on -// unresponsive servers. +// Available returns true if the embedding server responds successfully. +// Uses a 2s timeout to avoid blocking the CLI on unresponsive servers. +// +// OpenAI-compatible servers are probed via GET /models, the +// conventional discovery route. Some compatible providers do not serve +// that route at all (e.g. Voyage AI returns 404 while /embeddings works); +// when the models route is missing (404/405/501) the probe falls back to +// a single embedding round-trip, so availability reflects the endpoint +// the client actually depends on. func (c *Client) Available() bool { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() @@ -145,21 +151,81 @@ func (c *Client) Available() bool { default: route = "api/tags" } + status, ok := c.probeStatus(ctx, route) + if ok { + return true + } + if c.protocol == ProtocolOpenAI && (status == 404 || status == 405 || status == 501) { + return c.probeEmbed(ctx) + } + return false +} + +// probeStatus issues a GET against a discovery route and reports the +// HTTP status code. Transport errors yield status 0, ok false. +func (c *Client) probeStatus(ctx context.Context, route string) (status int, ok bool) { endpointURL, err := c.endpointURL(route) if err != nil { - return false + return 0, false } req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpointURL, nil) if err != nil { - return false + return 0, false } c.applyAuth(req) resp, err := c.http.Do(req) + if err != nil { + return 0, false + } + defer resp.Body.Close() + return resp.StatusCode, resp.StatusCode == http.StatusOK +} + +// probeEmbed verifies availability with a real embedding round-trip and +// discards the vector. Only reached when the OpenAI models route does not +// exist, so auth or quota failures still report unavailable. +func (c *Client) probeEmbed(ctx context.Context) bool { + vec, err := c.embedWithContext(ctx, "availability probe") if err != nil { return false } + return len(vec) > 0 +} + +// embedWithContext is Embed with a caller-supplied context so the +// availability probe can enforce its 2s deadline. +func (c *Client) embedWithContext(ctx context.Context, text string) ([]float64, error) { + req := embedRequest{Model: c.model, Input: text} + if c.dims > 0 { + req.Dimensions = c.dims + } + body, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + endpointURL, err := c.endpointURL(c.embedRequestRoute()) + if err != nil { + return nil, err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + c.applyAuth(httpReq) + + resp, err := c.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("embed request: %w", err) + } defer resp.Body.Close() - return resp.StatusCode == http.StatusOK + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("embedding provider returned status %d", resp.StatusCode) + } + + return c.decodeEmbedResponse(resp) } // Model returns the configured model name. @@ -196,47 +262,25 @@ type openaiEmbedResponse struct { } `json:"data"` } +// embedRequestRoute returns the protocol-specific embeddings route. +func (c *Client) embedRequestRoute() string { + if c.protocol == ProtocolOpenAI { + return "embeddings" + } + return "api/embed" +} + // Embed generates an embedding vector for the given text. // The request body is identical for both protocols; only the endpoint // path and the response shape differ. func (c *Client) Embed(text string) ([]float64, error) { - req := embedRequest{Model: c.model, Input: text} - if c.dims > 0 { - req.Dimensions = c.dims - } - body, err := json.Marshal(req) - if err != nil { - return nil, fmt.Errorf("marshal request: %w", err) - } - - var route string - switch c.protocol { - case ProtocolOpenAI: - route = "embeddings" - default: - route = "api/embed" - } - endpointURL, err := c.endpointURL(route) - if err != nil { - return nil, err - } - httpReq, err := http.NewRequest(http.MethodPost, endpointURL, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("build request: %w", err) - } - httpReq.Header.Set("Content-Type", "application/json") - c.applyAuth(httpReq) - - resp, err := c.http.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("embed request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("embedding provider returned status %d", resp.StatusCode) - } + return c.embedWithContext(context.Background(), text) +} +// decodeEmbedResponse parses a successful embeddings response under the +// active protocol. Shared between Embed and the OpenAI availability +// fallback so the probe and the real call cannot drift apart. +func (c *Client) decodeEmbedResponse(resp *http.Response) ([]float64, error) { switch c.protocol { case ProtocolOpenAI: var result openaiEmbedResponse diff --git a/internal/memory/embed/openai_test.go b/internal/memory/embed/openai_test.go index b85fe22e..71d18c5b 100644 --- a/internal/memory/embed/openai_test.go +++ b/internal/memory/embed/openai_test.go @@ -170,3 +170,89 @@ func TestOpenAIEmbedEmptyResponse(t *testing.T) { t.Fatal("expected error for empty embedding response") } } + +func TestOpenAIAvailableFallsBackWithoutModelsRoute(t *testing.T) { + // OpenAI-compatible servers without a models route (e.g. Voyage AI) + // must still be reported available via an embeddings round-trip. + t.Setenv("MNEMON_EMBED_API_KEY", "sk-test") + var embedRequests int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models": + http.NotFound(w, r) + case "/v1/embeddings": + if r.Method != http.MethodPost { + t.Errorf("expected POST /v1/embeddings, got %s", r.Method) + } + if got := r.Header.Get("Authorization"); got != "Bearer sk-test" { + t.Errorf("expected Bearer sk-test on fallback probe, got %q", got) + } + embedRequests++ + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"embedding":[1.0,2.0]}]}`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + if !c.Available() { + t.Fatal("expected Available() true when /v1/models is 404 but /v1/embeddings works") + } + if embedRequests != 1 { + t.Fatalf("expected exactly one embedding probe, got %d", embedRequests) + } +} + +func TestOpenAIAvailableFallbackRejectsAuthFailure(t *testing.T) { + // A 404 models route plus a 401 embeddings route must report + // unavailable: availability follows the endpoint that matters. + t.Setenv("MNEMON_EMBED_API_KEY", "sk-bad") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models": + http.NotFound(w, r) + case "/v1/embeddings": + w.WriteHeader(http.StatusUnauthorized) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + if c.Available() { + t.Fatal("expected Available() false when fallback probe returns 401") + } +} + +func TestOpenAIAvailableNoFallbackOnServerError(t *testing.T) { + // Only a missing models route (404/405/501) triggers the fallback. + // A 500 models route must report unavailable without an embeddings call. + var embedRequests int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/v1/models": + w.WriteHeader(http.StatusInternalServerError) + case "/v1/embeddings": + embedRequests++ + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"embedding":[1.0]}]}`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+"/v1") + c := NewClient() + if c.Available() { + t.Fatal("expected Available() false for 500 models route") + } + if embedRequests != 0 { + t.Fatalf("expected no embedding probe after 500 models route, got %d", embedRequests) + } +}