@malloy-publisher/server 0.0.230 → 0.0.232

package/src/mcp/tools/get_context_eval.ts

@@ -1,22 +1,34 @@
 /**
- * Lexical-baseline eval for malloy_getContext.
+ * Retrieval eval for malloy_getContext: lexical baseline vs the optional
+ * embedding-backed semantic mode.
  *
  * Runs a labeled set of plain-English queries against the live MCP endpoint
- * and reports recall@K: whether the expected entity appears in the top-K results.
- * This establishes the lexical (lunr/BM25) baseline so retrieval quality is
- * measurable; an embeddings comparison is deferred to if/when an embedding
- * provider is added (see the EPC eval risk).
+ * and reports recall@K: whether the expected entity appears in the top-K
+ * results. Each case is labeled with the retrieval mode the server reported
+ * (`semantic` / `lexical`; servers without an embedding provider report no
+ * mode and run the lexical baseline).
+ *
+ * A/B: run once against a server started WITHOUT `EMBEDDING_API_KEY`
+ * (lexical baseline), then once against a server started with it, and
+ * compare the two summaries. The cases marked `gap` share no tokens with
+ * their target entity's name or docs, so the lexical run is expected to
+ * miss them and the semantic run to hit them.
  *
  *   bun run packages/server/src/mcp/tools/get_context_eval.ts [K]
  *
- * Targets MCP_URL (default http://localhost:4040/mcp) over the demo
- * environment (EVAL_ENV, default "samples"); start the server first.
+ * Targets MCP_URL (default http://localhost:4040/mcp); start the server
+ * first. Cases cover the bundled `examples` environment (storefront), so
+ * the default config works out of the box, plus the classic ecommerce/faa
+ * cases for a malloy-samples config (`samples` environment); cases whose
+ * environment is not served are skipped, not counted as misses.
  */
 
 interface EvalCase {
+   env: string;
    pkg: string;
    query: string;
    expect: string; // substring expected in a top-K result name (case-insensitive)
+   gap?: boolean; // true = no token overlap; lexical is expected to miss
 }
 
 interface ResultEntity {
@@ -25,34 +37,108 @@ interface ResultEntity {
    source?: string | null;
 }
 
-// Ground truth verified against the demo packages (ecommerce, faa).
+interface GetContextPayload {
+   retrieval?: string;
+   results?: ResultEntity[];
+   error?: string;
+}
+
+// Ground truth verified against the bundled examples (storefront) and the
+// malloy-samples demo packages (ecommerce, faa).
 const CASES: EvalCase[] = [
+   // examples/storefront: served by the default config, works out of the box.
+   {
+      env: "examples",
+      pkg: "storefront",
+      query: "revenue by product category",
+      expect: "category",
+   },
+   {
+      env: "examples",
+      pkg: "storefront",
+      query: "top selling products",
+      expect: "top_products",
+   },
+   {
+      env: "examples",
+      pkg: "storefront",
+      query: "monthly sales trend",
+      expect: "sales_by_month",
+   },
+   // Token-gap cases: no word here appears in the entity's name or #(doc)
+   // text, so a token-based index misses them; embeddings should not.
+   {
+      env: "examples",
+      pkg: "storefront",
+      query: "how much money did we make",
+      expect: "total_sales",
+      gap: true,
+   },
    {
+      env: "examples",
+      pkg: "storefront",
+      query: "refund percentage",
+      expect: "return_rate",
+      gap: true,
+   },
+   // malloy-samples (needs a config serving them as environment "samples").
+   {
+      env: "samples",
       pkg: "ecommerce",
       query: "revenue by product category",
       expect: "category",
    },
-   { pkg: "ecommerce", query: "total sales revenue", expect: "total_sales" },
-   { pkg: "ecommerce", query: "customer state location", expect: "state" },
-   { pkg: "ecommerce", query: "order count", expect: "order_count" },
-   { pkg: "ecommerce", query: "product brand", expect: "brand" },
-   { pkg: "faa", query: "flights by carrier", expect: "carrier" },
-   { pkg: "faa", query: "airport", expect: "airport" },
-   { pkg: "faa", query: "aircraft model", expect: "aircraft" },
-   // Known lexical gap: the field is "dep_delay", which does not share tokens with
-   // "departure delay", so a token-based index misses it. Kept to surface the gap.
-   { pkg: "faa", query: "departure delay", expect: "delay" },
+   {
+      env: "samples",
+      pkg: "ecommerce",
+      query: "total sales revenue",
+      expect: "total_sales",
+   },
+   {
+      env: "samples",
+      pkg: "ecommerce",
+      query: "customer state location",
+      expect: "state",
+   },
+   {
+      env: "samples",
+      pkg: "ecommerce",
+      query: "order count",
+      expect: "order_count",
+   },
+   {
+      env: "samples",
+      pkg: "ecommerce",
+      query: "product brand",
+      expect: "brand",
+   },
+   {
+      env: "samples",
+      pkg: "faa",
+      query: "flights by carrier",
+      expect: "carrier",
+   },
+   { env: "samples", pkg: "faa", query: "airport", expect: "airport" },
+   { env: "samples", pkg: "faa", query: "aircraft model", expect: "aircraft" },
+   // The original lexical-gap case: the field is "dep_delay", which shares
+   // no tokens with "departure delay".
+   {
+      env: "samples",
+      pkg: "faa",
+      query: "departure delay",
+      expect: "delay",
+      gap: true,
+   },
 ];
 
 const ENDPOINT = process.env.MCP_URL || "http://localhost:4040/mcp";
-const EVAL_ENV = process.env.EVAL_ENV || "samples";
 const K = Number(process.argv[2] || 5);
+const WARMUP_ATTEMPTS = 15;
+const WARMUP_DELAY_MS = 2_000;
 
-async function getContext(
-   pkg: string,
-   query: string,
-   limit: number,
-): Promise<ResultEntity[]> {
+async function callGetContext(
+   args: Record<string, unknown>,
+): Promise<GetContextPayload> {
    const res = await fetch(ENDPOINT, {
       method: "POST",
       headers: {
@@ -63,36 +149,91 @@ async function getContext(
          jsonrpc: "2.0",
          id: 1,
          method: "tools/call",
-         params: {
-            name: "malloy_getContext",
-            arguments: {
-               environmentName: EVAL_ENV,
-               packageName: pkg,
-               query,
-               limit,
-            },
-         },
+         params: { name: "malloy_getContext", arguments: args },
       }),
    });
    const text = await res.text();
    const dataLine = text.split("\n").find((l) => l.startsWith("data: "));
-   if (!dataLine) return [];
+   if (!dataLine) return { results: [] };
    const msg = JSON.parse(dataLine.slice(6)) as {
       result?: { content?: { resource?: { text?: string } }[] };
    };
    const payloadText = msg.result?.content?.[0]?.resource?.text;
-   if (!payloadText) return [];
-   const payload = JSON.parse(payloadText) as { results?: ResultEntity[] };
-   return payload.results ?? [];
+   if (!payloadText) return { results: [] };
+   return JSON.parse(payloadText) as GetContextPayload;
+}
+
+/** True when the environment exists on the server (tier-2 listing works). */
+async function environmentServed(env: string): Promise<boolean> {
+   const payload = await callGetContext({ environmentName: env });
+   return payload.error === undefined;
+}
+
+/**
+ * On a semantic-capable server the first query against a package kicks off
+ * the embedding sync and answers lexically; poll until the mode flips (or
+ * give up and report whatever mode the server settles on).
+ */
+async function warmUp(env: string, pkg: string): Promise<string | undefined> {
+   let mode: string | undefined;
+   for (let attempt = 0; attempt < WARMUP_ATTEMPTS; attempt++) {
+      const payload = await callGetContext({
+         environmentName: env,
+         packageName: pkg,
+         query: "warm up",
+         limit: 1,
+      });
+      mode = payload.retrieval;
+      // No marker: the server has no embedding provider configured, so
+      // there is nothing to warm up.
+      if (mode === undefined || mode === "semantic") return mode;
+      await new Promise((resolve) => setTimeout(resolve, WARMUP_DELAY_MS));
+   }
+   return mode;
 }
 
 async function main(): Promise<void> {
    console.log(
-      `get_context lexical baseline, recall@${K} (endpoint ${ENDPOINT})\n`,
+      `get_context retrieval eval, recall@${K} (endpoint ${ENDPOINT})\n`,
    );
+
+   const servedEnvs = new Map<string, boolean>();
+   for (const env of new Set(CASES.map((c) => c.env))) {
+      servedEnvs.set(env, await environmentServed(env));
+   }
+   for (const [env, served] of servedEnvs) {
+      if (!served) {
+         console.log(
+            `  [SKIP] environment "${env}" is not served; skipping its cases`,
+         );
+      }
+   }
+
+   const warmed = new Set<string>();
    let hits = 0;
+   let scored = 0;
+   const modesSeen = new Set<string>();
+
    for (const c of CASES) {
-      const results = await getContext(c.pkg, c.query, K);
+      if (!servedEnvs.get(c.env)) continue;
+
+      const pkgKey = `${c.env}/${c.pkg}`;
+      if (!warmed.has(pkgKey)) {
+         warmed.add(pkgKey);
+         await warmUp(c.env, c.pkg);
+      }
+
+      const payload = await callGetContext({
+         environmentName: c.env,
+         packageName: c.pkg,
+         query: c.query,
+         limit: K,
+      });
+      const results = payload.results ?? [];
+      const mode = payload.retrieval ?? "lexical (no provider)";
+      modesSeen.add(mode);
+
+      scored++;
       const rank = results.findIndex((r) =>
          r.name.toLowerCase().includes(c.expect.toLowerCase()),
       );
@@ -104,20 +245,28 @@ async function main(): Promise<void> {
             .map((r) => `${r.kind}:${r.name}`)
             .join(", ") || "(none)";
       const tag = hit ? `HIT@${rank + 1}` : "MISS  ";
-      console.log(`  [${tag}] ${c.pkg} / "${c.query}" (want ~${c.expect})`);
+      const gap = c.gap ? " [gap]" : "";
+      console.log(
+         `  [${tag}] [${mode}] ${c.env}/${c.pkg} / "${c.query}" (want ~${c.expect})${gap}`,
+      );
       console.log(`            top: ${top}`);
    }
-   const pct = ((hits / CASES.length) * 100).toFixed(0);
-   console.log(`\nrecall@${K}: ${hits}/${CASES.length} (${pct}%)`);
+
+   if (scored === 0) {
+      console.log("\nNo cases ran: no listed environment is served.");
+      process.exit(1);
+   }
+   const pct = ((hits / scored) * 100).toFixed(0);
    console.log(
-      "Misses occur where the field name does not share tokens with the query",
+      `\nrecall@${K}: ${hits}/${scored} (${pct}%)  mode(s): ${[...modesSeen].join(", ")}`,
    );
    console.log(
-      "(e.g. 'departure delay' vs a 'dep_delay' field). Closing that gap is the",
+      "The [gap] cases share no tokens with their target entity, so a lexical",
    );
    console.log(
-      "motivation for an optional embeddings provider; lexical ships as the v1 baseline.",
+      "run misses them by construction; a semantic run (EMBEDDING_API_KEY set)",
    );
+   console.log("is expected to close them. Compare the two summaries.");
 }
 
 main().catch((err) => {

package/src/mcp/tools/get_context_tool.spec.ts

@@ -1,6 +1,42 @@
-import { describe, expect, it } from "bun:test";
-import { docText, sanitize, registerGetContextTool } from "./get_context_tool";
+import {
+   afterAll,
+   afterEach,
+   beforeAll,
+   beforeEach,
+   describe,
+   expect,
+   it,
+} from "bun:test";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+import {
+   docOnlyText,
+   docText,
+   sanitize,
+   registerGetContextTool,
+} from "./get_context_tool";
+import { embeddingText } from "./embedding_index";
 import type { EnvironmentStore } from "../../service/environment_store";
+import { DuckDBConnection } from "../../storage/duckdb/DuckDBConnection";
+import { createEntityEmbeddingsTable } from "../../storage/duckdb/schema";
+import {
+   EmbeddingProvider,
+   _clearEmbeddingProviderForTests,
+   _setEmbeddingProviderForTests,
+} from "../../service/embedding_provider";
+import { _resetEmbeddingIndexStateForTests } from "./embedding_index";
+
+// Pin the baseline: every test in this file runs unconfigured (lexical)
+// unless it sets a provider itself, regardless of ambient EMBEDDING_* env
+// vars in the developer's shell.
+beforeEach(() => {
+   _setEmbeddingProviderForTests(null);
+   _resetEmbeddingIndexStateForTests();
+});
+afterAll(() => {
+   _clearEmbeddingProviderForTests();
+});
 
 describe("get_context docText", () => {
    it("extracts #(doc) text from annotation lines", () => {
@@ -31,6 +67,56 @@ describe("get_context docText", () => {
    });
 });
 
+describe("get_context docOnlyText (embedding-input safety)", () => {
+   it("extracts #(doc) text like docText", () => {
+      expect(docOnlyText(["#(doc) Total revenue."])).toBe("Total revenue.");
+      expect(docOnlyText([{ value: "#(doc) Customer state." }])).toBe(
+         "Customer state.",
+      );
+   });
+
+   it("does NOT fall back to raw lines: no #(doc) yields empty", () => {
+      // The security-critical difference from docText: predicate-bearing
+      // annotations must never become embedding input.
+      expect(docOnlyText(["# bar_chart"])).toBe("");
+      expect(
+         docOnlyText([
+            "#(authorize) \"$ROLE = 'admin'\"",
+            "#(malloy) drillable",
+         ]),
+      ).toBe("");
+   });
+
+   it("keeps only the #(doc) line when mixed with predicate annotations", () => {
+      expect(
+         docOnlyText([
+            "#(authorize) \"$TENANT = 'acme'\"",
+            "#(doc) Secured orders.",
+         ]),
+      ).toBe("Secured orders.");
+   });
+
+   it("an #(authorize)-only entity produces no embedding text beyond its name", () => {
+      // End-to-end: what embeddingText actually sends for a governed,
+      // undocumented entity is the humanized name only, never the predicate.
+      const embedDoc = docOnlyText([
+         "#(authorize) \"$ROLE = 'admin'\"",
+         "#(authorize) \"$TENANT = 'acme' or $TENANT = 'globex'\"",
+      ]);
+      const text = embeddingText({
+         kind: "source",
+         name: "orders_secured",
+         source: "orders_secured",
+         modelPath: "m.malloy",
+         embedDoc,
+      });
+      expect(text).toBe("orders secured");
+      expect(text).not.toContain("authorize");
+      expect(text).not.toContain("ROLE");
+      expect(text).not.toContain("acme");
+   });
+});
+
 describe("get_context sanitize", () => {
    it("strips lunr operator characters so a plain-English query never throws", () => {
       const cleaned = sanitize('revenue: +top ~10 "exact" -minus ^boost *wild');
@@ -236,4 +322,211 @@ describe("get_context discovery tiers", () => {
       );
       expect(results).toHaveLength(5);
    });
+
+   it("tier 4: without a provider the payload has no retrieval marker or scores", async () => {
+      const handler = captureHandler({
+         getEnvironment: async () =>
+            ({ getPackage: async () => mockPackage }) as never,
+      });
+      const payload = parse(
+         await handler({
+            environmentName: "malloy-samples",
+            packageName: "ecommerce",
+            query: "state",
+         }),
+      );
+      expect(payload.retrieval).toBeUndefined();
+      expect(Object.keys(payload)).toEqual(["results"]);
+      for (const r of payload.results) {
+         expect(r.score).toBeUndefined();
+      }
+   });
+});
+
+describe("get_context semantic retrieval", () => {
+   let tempDir: string;
+   let db: DuckDBConnection;
+
+   beforeAll(async () => {
+      tempDir = fs.mkdtempSync(
+         path.join(os.tmpdir(), "get-context-semantic-spec-"),
+      );
+      db = new DuckDBConnection(path.join(tempDir, "test.db"));
+      await db.initialize();
+      await createEntityEmbeddingsTable(db);
+   });
+
+   afterAll(async () => {
+      // Close before removing: Windows refuses to delete a directory
+      // holding DuckDB's open file handle.
+      await db.close();
+      fs.rmSync(tempDir, { recursive: true, force: true });
+   });
+
+   afterEach(() => {
+      _setEmbeddingProviderForTests(null);
+   });
+
+   // Entity texts are humanize(name) [+ doc]: "order items" and "state".
+   const VECTORS: Record<string, number[]> = {
+      "order items": [0, 1],
+      state: [1, 0],
+      "where do customers live": [1, 0],
+   };
+
+   function stubProvider(options: { fail?: boolean } = {}): EmbeddingProvider {
+      const fetchStub = (async (
+         _url: RequestInfo | URL,
+         init?: RequestInit,
+      ) => {
+         if (options.fail) return new Response("down", { status: 500 });
+         const body = JSON.parse(String(init?.body)) as { input: string[] };
+         const data = body.input.map((text, index) => {
+            const embedding = VECTORS[text];
+            if (!embedding) throw new Error(`no stub vector for "${text}"`);
+            return { index, embedding };
+         });
+         return new Response(JSON.stringify({ data }), { status: 200 });
+      }) as typeof fetch;
+      return new EmbeddingProvider(
+         {
+            apiKey: "test",
+            model: "stub-model",
+            baseUrl: "https://stub.example.com/v1",
+         },
+         fetchStub,
+      );
+   }
+
+   /** A store whose package is a fresh instance, backed by the temp DB. */
+   function semanticStore(): Partial<EnvironmentStore> {
+      const pkg = {
+         listModels: async () => [{ path: "ecommerce.malloy" }],
+         getModel: () => mockModel,
+      };
+      return {
+         getEnvironment: async () => ({ getPackage: async () => pkg }) as never,
+         storageManager: {
+            getDuckDbConnection: () => db,
+         } as never,
+      };
+   }
+
+   async function callUntilSemantic(
+      handler: Handler,
+      params: Record<string, unknown>,
+   ) {
+      for (let i = 0; i < 200; i++) {
+         const payload = parse(await handler(params));
+         if (payload.retrieval === "semantic") return payload;
+         await new Promise((resolve) => setTimeout(resolve, 5));
+      }
+      throw new Error("retrieval never became semantic");
+   }
+
+   it("ranks semantically once the index syncs, with scores and a marker", async () => {
+      _setEmbeddingProviderForTests(stubProvider());
+      const handler = captureHandler(semanticStore());
+      const params = {
+         environmentName: "specs",
+         packageName: "semantic-pkg",
+         query: "where do customers live",
+      };
+
+      // Cold start: the sync kicks off in the background and this call
+      // answers lexically, marked as such.
+      const first = parse(await handler(params));
+      expect(first.retrieval).toBe("lexical");
+
+      const payload = await callUntilSemantic(handler, params);
+      expect(payload.results[0].name).toBe("state");
+      expect(payload.results[0].kind).toBe("dimension");
+      expect(payload.results[0].score).toBeCloseTo(1.0, 3);
+      // "order items" is orthogonal to the query: below the similarity
+      // floor, so it must not pad the results.
+      expect(
+         payload.results.some(
+            (r: { name: string }) => r.name === "order_items",
+         ),
+      ).toBe(false);
+   });
+
+   it("falls back to lexical, marked, when the provider goes down after indexing", async () => {
+      // Index healthily first, so this pins the query-embed failure
+      // path, not just the cold start (which answers lexically anyway).
+      _setEmbeddingProviderForTests(stubProvider());
+      const handler = captureHandler(semanticStore());
+      const params = {
+         environmentName: "specs",
+         packageName: "failing-pkg",
+         query: "where do customers live",
+      };
+      await callUntilSemantic(handler, params);
+
+      // Same model and config, but the endpoint now returns 500s: the
+      // per-call query embed fails and the call degrades to marked
+      // lexical with no scores.
+      _setEmbeddingProviderForTests(stubProvider({ fail: true }));
+      const payload = parse(await handler({ ...params, query: "state" }));
+      expect(payload.retrieval).toBe("lexical");
+      expect(
+         payload.results.some((r: { name: string }) => r.name === "state"),
+      ).toBe(true);
+      for (const r of payload.results) {
+         expect(r.score).toBeUndefined();
+      }
+   });
+
+   it("degrades to lexical when the storage handle is unavailable", async () => {
+      _setEmbeddingProviderForTests(stubProvider());
+      const store = semanticStore();
+      delete (store as { storageManager?: unknown }).storageManager;
+      const handler = captureHandler(store);
+      const payload = parse(
+         await handler({
+            environmentName: "specs",
+            packageName: "no-storage-pkg",
+            query: "state",
+         }),
+      );
+      expect(payload.retrieval).toBe("lexical");
+      expect(
+         payload.results.some((r: { name: string }) => r.name === "state"),
+      ).toBe(true);
+   });
+
+   it("degrades to lexical (never throws) when the real config is malformed", async () => {
+      // Exercises the tool-path catch with REAL env parsing, not the
+      // _setEmbeddingProviderForTests override: a malformed base makes
+      // getEmbeddingProvider() throw, and tier 4 must swallow it and
+      // answer marked lexical (embeddingConfigured() is still true).
+      const saved = {
+         key: process.env.EMBEDDING_API_KEY,
+         base: process.env.EMBEDDING_API_BASE,
+      };
+      process.env.EMBEDDING_API_KEY = "k";
+      process.env.EMBEDDING_API_BASE = "not a url";
+      _clearEmbeddingProviderForTests();
+      try {
+         const handler = captureHandler(semanticStore());
+         const payload = parse(
+            await handler({
+               environmentName: "specs",
+               packageName: "malformed-cfg-pkg",
+               query: "state",
+            }),
+         );
+         expect(payload.retrieval).toBe("lexical");
+         expect(
+            payload.results.some((r: { name: string }) => r.name === "state"),
+         ).toBe(true);
+      } finally {
+         if (saved.key === undefined) delete process.env.EMBEDDING_API_KEY;
+         else process.env.EMBEDDING_API_KEY = saved.key;
+         if (saved.base === undefined) delete process.env.EMBEDDING_API_BASE;
+         else process.env.EMBEDDING_API_BASE = saved.base;
+         _clearEmbeddingProviderForTests();
+         _setEmbeddingProviderForTests(null);
+      }
+   });
 });