@malloy-publisher/server 0.0.230 → 0.0.232

package/src/service/embedding_provider.spec.ts

@@ -0,0 +1,329 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import {
+   EmbeddingProvider,
+   MAX_EMBED_BATCH_SIZE,
+   MAX_EMBED_INPUT_CHARS,
+   _clearEmbeddingProviderForTests,
+   embeddingConfigured,
+   getEmbeddingProvider,
+   prepareEmbeddingInput,
+} from "./embedding_provider";
+import { EmbeddingConfig } from "../config";
+
+const CONFIG: EmbeddingConfig = {
+   apiKey: "test-key",
+   model: "test-model",
+   baseUrl: "https://embeddings.example.com/v1",
+};
+
+interface CapturedRequest {
+   url: string;
+   body: Record<string, unknown>;
+   authorization: string | undefined;
+}
+
+/** A fetch stub that records requests and answers with unit vectors. */
+function stubFetch(
+   captured: CapturedRequest[],
+   respond?: (inputs: string[]) => Response,
+): typeof fetch {
+   return (async (url: RequestInfo | URL, init?: RequestInit) => {
+      const body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+      captured.push({
+         url: String(url),
+         body,
+         authorization: (init?.headers as Record<string, string>)?.[
+            "Authorization"
+         ],
+      });
+      const inputs = body.input as string[];
+      if (respond) return respond(inputs);
+      return new Response(
+         JSON.stringify({
+            data: inputs.map((_, index) => ({ index, embedding: [index, 1] })),
+         }),
+         { status: 200 },
+      );
+   }) as typeof fetch;
+}
+
+describe("prepareEmbeddingInput", () => {
+   it("collapses whitespace and newlines", () => {
+      expect(prepareEmbeddingInput("a\nb\t c  d")).toBe("a b c d");
+   });
+
+   it("caps input length", () => {
+      const long = "x".repeat(MAX_EMBED_INPUT_CHARS + 500);
+      expect(prepareEmbeddingInput(long).length).toBe(MAX_EMBED_INPUT_CHARS);
+   });
+
+   it("never returns an empty string (providers 400 on empty inputs)", () => {
+      expect(prepareEmbeddingInput("")).toBe("-");
+      expect(prepareEmbeddingInput("   ")).toBe("-");
+      expect(prepareEmbeddingInput("\n\t")).toBe("-");
+   });
+});
+
+describe("EmbeddingProvider", () => {
+   it("embeds a batch and returns vectors in input order", async () => {
+      const captured: CapturedRequest[] = [];
+      const provider = new EmbeddingProvider(CONFIG, stubFetch(captured));
+      const vectors = await provider.embedBatch(["a", "b"], 1000);
+      expect(vectors).toEqual([
+         [0, 1],
+         [1, 1],
+      ]);
+      expect(captured).toHaveLength(1);
+      expect(captured[0].url).toBe(
+         "https://embeddings.example.com/v1/embeddings",
+      );
+      expect(captured[0].body.model).toBe("test-model");
+      expect(captured[0].authorization).toBe("Bearer test-key");
+   });
+
+   it("places vectors by the response's index fields, not response order", async () => {
+      const captured: CapturedRequest[] = [];
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            captured,
+            () =>
+               new Response(
+                  JSON.stringify({
+                     data: [
+                        { index: 1, embedding: [2, 2] },
+                        { index: 0, embedding: [1, 1] },
+                     ],
+                  }),
+                  { status: 200 },
+               ),
+         ),
+      );
+      const vectors = await provider.embedBatch(["a", "b"], 1000);
+      expect(vectors).toEqual([
+         [1, 1],
+         [2, 2],
+      ]);
+   });
+
+   it("splits large inputs into batches", async () => {
+      const captured: CapturedRequest[] = [];
+      const provider = new EmbeddingProvider(CONFIG, stubFetch(captured));
+      const texts = Array.from(
+         { length: MAX_EMBED_BATCH_SIZE + 1 },
+         (_, i) => `t${i}`,
+      );
+      const vectors = await provider.embedBatch(texts, 1000);
+      expect(vectors).toHaveLength(texts.length);
+      expect(captured).toHaveLength(2);
+      expect((captured[0].body.input as string[]).length).toBe(
+         MAX_EMBED_BATCH_SIZE,
+      );
+      expect((captured[1].body.input as string[]).length).toBe(1);
+   });
+
+   it("sends the dimensions param only when configured", async () => {
+      const captured: CapturedRequest[] = [];
+      const without = new EmbeddingProvider(CONFIG, stubFetch(captured));
+      await without.embedBatch(["a"], 1000);
+      expect(captured[0].body.dimensions).toBeUndefined();
+
+      const withDims = new EmbeddingProvider(
+         { ...CONFIG, dimensions: 512 },
+         stubFetch(captured),
+      );
+      await withDims.embedBatch(["a"], 1000);
+      expect(captured[1].body.dimensions).toBe(512);
+   });
+
+   it("reports HTTP failures with a truncated body and no header echo", async () => {
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            [],
+            () => new Response("boom ".repeat(100), { status: 429 }),
+         ),
+      );
+      const err = await provider.embedBatch(["a"], 1000).then(
+         () => undefined,
+         (e: Error) => e,
+      );
+      expect(err).toBeDefined();
+      expect(err!.message).toContain("(429)");
+      expect(err!.message.length).toBeLessThan(400);
+      expect(err!.message).not.toContain("test-key");
+   });
+
+   it("drops auth-failure bodies entirely (they can reflect the key)", async () => {
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            [],
+            () =>
+               new Response("Incorrect API key provided: test-key", {
+                  status: 401,
+               }),
+         ),
+      );
+      const err = await provider.embedBatch(["a"], 1000).then(
+         () => undefined,
+         (e: Error) => e,
+      );
+      expect(err!.message).toContain("(401)");
+      expect(err!.message).toContain("authentication failed");
+      expect(err!.message).not.toContain("test-key");
+   });
+
+   it("scrubs literal key occurrences from non-auth error bodies", async () => {
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            [],
+            () =>
+               new Response(
+                  "proxy error while forwarding token test-key upstream",
+                  {
+                     status: 502,
+                  },
+               ),
+         ),
+      );
+      const err = await provider.embedBatch(["a"], 1000).then(
+         () => undefined,
+         (e: Error) => e,
+      );
+      expect(err!.message).toContain("(502)");
+      expect(err!.message).toContain("[REDACTED]");
+      expect(err!.message).not.toContain("test-key");
+   });
+
+   it("names a timeout in the failure message", async () => {
+      const timeoutFetch = (async () => {
+         const error = new Error("The operation timed out");
+         error.name = "TimeoutError";
+         throw error;
+      }) as unknown as typeof fetch;
+      const provider = new EmbeddingProvider(CONFIG, timeoutFetch);
+      const err = await provider.embedBatch(["a"], 1234).then(
+         () => undefined,
+         (e: Error) => e,
+      );
+      expect(err!.message).toContain("timed out after 1234ms");
+   });
+
+   it("rejects a response with the wrong number of embeddings", async () => {
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            [],
+            () =>
+               new Response(
+                  JSON.stringify({ data: [{ index: 0, embedding: [1] }] }),
+                  { status: 200 },
+               ),
+         ),
+      );
+      await expect(provider.embedBatch(["a", "b"], 1000)).rejects.toThrow(
+         "malformed",
+      );
+   });
+
+   it("rejects a vector with non-finite or non-numeric contents", async () => {
+      const provider = new EmbeddingProvider(
+         CONFIG,
+         stubFetch(
+            [],
+            () =>
+               new Response(
+                  JSON.stringify({
+                     data: [{ index: 0, embedding: [1, null, "x"] }],
+                  }),
+                  { status: 200 },
+               ),
+         ),
+      );
+      await expect(provider.embedBatch(["a"], 1000)).rejects.toThrow(
+         "malformed",
+      );
+   });
+});
+
+describe("getEmbeddingProvider / embeddingConfigured", () => {
+   const ENV_KEYS = [
+      "EMBEDDING_API_KEY",
+      "EMBEDDING_MODEL",
+      "EMBEDDING_API_BASE",
+      "EMBEDDING_DIMENSIONS",
+   ];
+   const saved: Record<string, string | undefined> = {};
+
+   beforeEach(() => {
+      for (const key of ENV_KEYS) {
+         saved[key] = process.env[key];
+         delete process.env[key];
+      }
+      _clearEmbeddingProviderForTests();
+   });
+
+   afterEach(() => {
+      for (const key of ENV_KEYS) {
+         if (saved[key] === undefined) delete process.env[key];
+         else process.env[key] = saved[key];
+      }
+      _clearEmbeddingProviderForTests();
+   });
+
+   it("returns null (feature off) without EMBEDDING_API_KEY", () => {
+      expect(getEmbeddingProvider()).toBeNull();
+      expect(embeddingConfigured()).toBe(false);
+   });
+
+   it("ignores an ambient OPENAI_API_KEY", () => {
+      const before = process.env.OPENAI_API_KEY;
+      process.env.OPENAI_API_KEY = "ambient";
+      try {
+         expect(getEmbeddingProvider()).toBeNull();
+         expect(embeddingConfigured()).toBe(false);
+      } finally {
+         if (before === undefined) delete process.env.OPENAI_API_KEY;
+         else process.env.OPENAI_API_KEY = before;
+      }
+   });
+
+   it("builds a provider with defaults when the key is set", () => {
+      process.env.EMBEDDING_API_KEY = "k";
+      const provider = getEmbeddingProvider();
+      expect(provider).not.toBeNull();
+      expect(provider!.model).toBe("text-embedding-3-small");
+      expect(embeddingConfigured()).toBe(true);
+   });
+
+   it("rebuilds the provider when the configuration changes", () => {
+      process.env.EMBEDDING_API_KEY = "k";
+      const first = getEmbeddingProvider();
+      expect(getEmbeddingProvider()).toBe(first);
+      process.env.EMBEDDING_MODEL = "other-model";
+      const second = getEmbeddingProvider();
+      expect(second).not.toBe(first);
+      expect(second!.model).toBe("other-model");
+   });
+
+   it("does not stick to null after the key appears", () => {
+      expect(getEmbeddingProvider()).toBeNull();
+      process.env.EMBEDDING_API_KEY = "k";
+      expect(getEmbeddingProvider()).not.toBeNull();
+   });
+
+   it("throws on a malformed base URL but still reports configured", () => {
+      process.env.EMBEDDING_API_KEY = "k";
+      process.env.EMBEDDING_API_BASE = "not a url";
+      expect(() => getEmbeddingProvider()).toThrow("EMBEDDING_API_BASE");
+      expect(embeddingConfigured()).toBe(true);
+   });
+
+   it("throws on a malformed EMBEDDING_DIMENSIONS", () => {
+      process.env.EMBEDDING_API_KEY = "k";
+      process.env.EMBEDDING_DIMENSIONS = "lots";
+      expect(() => getEmbeddingProvider()).toThrow("EMBEDDING_DIMENSIONS");
+   });
+});

package/src/service/embedding_provider.ts

@@ -0,0 +1,236 @@
+import { EmbeddingConfig, getEmbeddingConfig } from "../config";
+
+/** Timeout for bulk (index-build) embedding calls. */
+export const EMBEDDING_BATCH_TIMEOUT_MS = 30_000;
+/**
+ * Timeout for the single per-request query embedding. Much shorter than
+ * the bulk timeout: this call sits on the latency path of every semantic
+ * `malloy_getContext` call, and a slow endpoint must degrade to lexical
+ * quickly rather than stall the tool.
+ */
+export const EMBEDDING_QUERY_TIMEOUT_MS = 5_000;
+/** Per-input character cap (matches the hosted indexing pipeline). */
+export const MAX_EMBED_INPUT_CHARS = 1_024;
+/** Inputs per HTTP request; large packages embed in several requests. */
+export const MAX_EMBED_BATCH_SIZE = 512;
+
+type FetchFn = typeof fetch;
+
+interface EmbeddingResponseItem {
+   index?: number;
+   embedding?: number[];
+}
+
+/**
+ * Collapse whitespace (newlines confuse some embedding models) and cap
+ * the input length, mirroring the hosted indexing pipeline's input
+ * preparation. Never returns an empty string: providers reject empty
+ * input items with a 400, which would fail a whole package's batch, so
+ * text that collapses to nothing (a whitespace-only backtick identifier)
+ * becomes a placeholder token instead. This is the single choke point
+ * for that invariant; the index hashes this same function's output, so
+ * hash and sent text cannot drift. Exported for tests and for
+ * content-hashing: the hash must cover the text actually sent, so a
+ * truncation-boundary edit re-embeds.
+ */
+export function prepareEmbeddingInput(text: string): string {
+   const cleaned = text.replace(/\s+/g, " ").trim();
+   const capped =
+      cleaned.length > MAX_EMBED_INPUT_CHARS
+         ? cleaned.slice(0, MAX_EMBED_INPUT_CHARS)
+         : cleaned;
+   return capped || "-";
+}
+
+/**
+ * Minimal client for an OpenAI-compatible `/embeddings` endpoint, called
+ * with global fetch (no provider SDK; see fetchFromPublisherDataplane for
+ * the outbound-HTTP precedent this follows). The API key travels only in
+ * the Authorization header and must never be logged; error messages carry
+ * at most a 200-char body excerpt and never echo request headers.
+ */
+export class EmbeddingProvider {
+   constructor(
+      private config: EmbeddingConfig,
+      private fetchFn: FetchFn = fetch,
+   ) {}
+
+   get model(): string {
+      return this.config.model;
+   }
+
+   get dimensions(): number | undefined {
+      return this.config.dimensions;
+   }
+
+   /**
+    * Embed `texts` in order. Inputs are prepared (whitespace-collapsed,
+    * capped) and sent in batches of {@link MAX_EMBED_BATCH_SIZE}. Throws
+    * on any HTTP, timeout, or malformed-response failure; callers own the
+    * fallback-to-lexical decision.
+    */
+   async embedBatch(texts: string[], timeoutMs: number): Promise<number[][]> {
+      const vectors: number[][] = [];
+      for (let i = 0; i < texts.length; i += MAX_EMBED_BATCH_SIZE) {
+         const chunk = texts
+            .slice(i, i + MAX_EMBED_BATCH_SIZE)
+            .map(prepareEmbeddingInput);
+         vectors.push(...(await this.embedChunk(chunk, timeoutMs)));
+      }
+      return vectors;
+   }
+
+   private async embedChunk(
+      inputs: string[],
+      timeoutMs: number,
+   ): Promise<number[][]> {
+      const url = `${this.config.baseUrl}/embeddings`;
+      const body: Record<string, unknown> = {
+         model: this.config.model,
+         input: inputs,
+      };
+      if (this.config.dimensions !== undefined) {
+         body.dimensions = this.config.dimensions;
+      }
+
+      let response: Response;
+      try {
+         response = await this.fetchFn(url, {
+            method: "POST",
+            headers: {
+               "Content-Type": "application/json",
+               Authorization: `Bearer ${this.config.apiKey}`,
+            },
+            body: JSON.stringify(body),
+            signal: AbortSignal.timeout(timeoutMs),
+         });
+      } catch (error) {
+         const reason =
+            (error as Error)?.name === "TimeoutError"
+               ? `timed out after ${timeoutMs}ms`
+               : (error as Error).message;
+         throw new Error(`Embedding request to ${url} failed: ${reason}`);
+      }
+
+      if (!response.ok) {
+         // Auth-failure bodies commonly reflect the presented credential
+         // (OpenAI's 401 echoes partial key characters; a self-hosted
+         // proxy may echo the whole token), and this message is logged
+         // by callers. Drop those bodies entirely, and scrub any literal
+         // occurrence of the key from the rest.
+         let detail: string;
+         if (response.status === 401 || response.status === 403) {
+            detail = "authentication failed; check EMBEDDING_API_KEY";
+         } else {
+            const bodyText = await response.text().catch(() => "");
+            detail = bodyText
+               .split(this.config.apiKey)
+               .join("[REDACTED]")
+               .slice(0, 200);
+         }
+         throw new Error(
+            `Embedding request to ${url} failed (${response.status}): ${detail}`,
+         );
+      }
+
+      const json = (await response.json()) as {
+         data?: EmbeddingResponseItem[];
+      };
+      const data = json?.data;
+      if (!Array.isArray(data) || data.length !== inputs.length) {
+         throw new Error(
+            `Embedding response from ${url} malformed: expected ${inputs.length} embeddings, got ${Array.isArray(data) ? data.length : "none"}`,
+         );
+      }
+
+      // The response documents per-item `index` fields; place by index
+      // defensively rather than assuming order.
+      const vectors = new Array<number[] | undefined>(inputs.length);
+      for (let i = 0; i < data.length; i++) {
+         const item = data[i];
+         const idx = typeof item?.index === "number" ? item.index : i;
+         if (
+            !Array.isArray(item?.embedding) ||
+            item.embedding.length === 0 ||
+            // Numeric contents, not just shape: a provider returning
+            // [1, null, "x"] would otherwise be stored and misbehave
+            // silently (NaN scores) at search time.
+            !item.embedding.every(
+               (n: unknown) => typeof n === "number" && Number.isFinite(n),
+            ) ||
+            idx < 0 ||
+            idx >= inputs.length ||
+            vectors[idx] !== undefined
+         ) {
+            throw new Error(
+               `Embedding response from ${url} malformed at item ${i}`,
+            );
+         }
+         vectors[idx] = item.embedding;
+      }
+      return vectors as number[][];
+   }
+}
+
+// Cached on a config fingerprint, never on null: a call after the env
+// changes (tests, operator restarts with new vars are moot, but the
+// integration suite runs many specs in one process) always sees the
+// current configuration instead of a stale provider or a sticky "off".
+let cached: { fingerprint: string; provider: EmbeddingProvider } | null = null;
+let testOverride: { provider: EmbeddingProvider | null } | null = null;
+
+/**
+ * Whether the operator has turned the embedding feature on at all.
+ * Malformed companion vars still count as configured: the operator
+ * clearly intends the feature, so the retrieval marker and a warning
+ * must surface rather than silently reporting plain lexical.
+ */
+export function embeddingConfigured(): boolean {
+   if (testOverride) {
+      return testOverride.provider !== null;
+   }
+   try {
+      return getEmbeddingConfig() !== null;
+   } catch {
+      return true;
+   }
+}
+
+/**
+ * The process-wide provider for the current embedding configuration, or
+ * null when `EMBEDDING_API_KEY` is unset. Throws on malformed companion
+ * env vars (see getEmbeddingConfig); callers on the tool path catch and
+ * degrade.
+ */
+export function getEmbeddingProvider(): EmbeddingProvider | null {
+   if (testOverride) {
+      return testOverride.provider;
+   }
+   const config = getEmbeddingConfig();
+   if (!config) {
+      cached = null;
+      return null;
+   }
+   const fingerprint = [
+      config.baseUrl,
+      config.model,
+      config.dimensions ?? "",
+      config.apiKey,
+   ].join("\u0000");
+   if (!cached || cached.fingerprint !== fingerprint) {
+      cached = { fingerprint, provider: new EmbeddingProvider(config) };
+   }
+   return cached.provider;
+}
+
+/** Test seam: force the provider (or null). Undo with _clear...(). */
+export function _setEmbeddingProviderForTests(
+   provider: EmbeddingProvider | null,
+): void {
+   testOverride = { provider };
+}
+
+export function _clearEmbeddingProviderForTests(): void {
+   testOverride = null;
+   cached = null;
+}