feat(search): batch query highlights (firecrawl#3993)

Author: mogery

apps/api/src/config.ts

@@ -159,9 +159,8 @@ const configSchema = z.object({
   CLICKHOUSE_ANALYTICS_URL: z.string().optional(),
   CLICKHOUSE_ANALYTICS_DATABASE: z.string().optional(),
 
-  // Search highlights (beta): query-highlights model service. URL is the base
-  // (e.g. https://firecrawl--query-highlights.modal.run); TOKEN is the bearer
-  // token sent on every request.
+  // Search highlights (beta): highlighter service base URL. TOKEN is optional
+  // bearer auth for legacy/external services; the in-cluster service omits it.
   HIGHLIGHT_MODEL_URL: z.string().optional(),
   HIGHLIGHT_MODEL_TOKEN: z.string().optional(),
 

apps/api/src/search/highlight-model.test.ts

@@ -5,7 +5,8 @@ vi.mock("../config", () => ({
   },
 }));
 
-import { generateHighlights } from "./highlight-model";
+import { generateHighlightsBatch } from "./highlight-model";
+import { config } from "../config";
 
 const logger = {
   info: vi.fn(),
@@ -30,34 +31,81 @@ afterEach(() => {
   vi.clearAllMocks();
 });
 
-describe("generateHighlights", () => {
-  it("posts the query and full markdown to /highlight with the bearer token", async () => {
-    const fetchMock = mockFetchOnce({ highlights: [], markdown: "" });
+describe("generateHighlightsBatch", () => {
+  it("posts every page to one /batch_highlight call with the bearer token", async () => {
+    const fetchMock = mockFetchOnce({ pages: [] });
 
-    await generateHighlights("q1", "# Page\n\nbody text", { logger });
+    await generateHighlightsBatch(
+      "q1",
+      [
+        { id: "0", markdown: "# First\n\nbody text" },
+        { id: "1", markdown: "# Second\n\nmore text" },
+      ],
+      { logger },
+    );
 
     expect(fetchMock).toHaveBeenCalledTimes(1);
     const [url, init] = fetchMock.mock.calls[0];
-    expect(url).toBe("https://highlight.test/highlight");
+    expect(url).toBe("https://highlight.test/batch_highlight");
     expect(init.method).toBe("POST");
     expect(init.headers.Authorization).toBe("Bearer secret-token");
 
     const sent = JSON.parse(init.body);
-    expect(sent).toEqual({ query: "q1", markdown: "# Page\n\nbody text" });
+    expect(sent).toEqual({
+      query: "q1",
+      pages: [
+        { id: "0", markdown: "# First\n\nbody text" },
+        { id: "1", markdown: "# Second\n\nmore text" },
+      ],
+    });
   });
 
-  it("returns the highlights and reassembled markdown from the service", async () => {
+  it("omits bearer auth for the in-cluster service when no token is configured", async () => {
+    const token = config.HIGHLIGHT_MODEL_TOKEN;
+    config.HIGHLIGHT_MODEL_TOKEN = undefined;
+    const fetchMock = mockFetchOnce({ pages: [] });
+
+    try {
+      await generateHighlightsBatch("q1", [{ id: "0", markdown: "markdown" }], {
+        logger,
+      });
+    } finally {
+      config.HIGHLIGHT_MODEL_TOKEN = token;
+    }
+
+    const [, init] = fetchMock.mock.calls[0];
+    expect(init.headers).toEqual({ "Content-Type": "application/json" });
+  });
+
+  it("returns successful pages keyed by ID and leaves missing pages absent", async () => {
     mockFetchOnce({
-      highlights: [
-        { block_index: 0, score: 0.9 },
-        { block_index: 2, score: 0.5 },
+      pages: [
+        {
+          id: "1",
+          cache: "hit",
+          output: {
+            highlights: [
+              { block_index: 0, score: 0.9 },
+              { block_index: 2, score: 0.5 },
+            ],
+            markdown: "reassembled answer",
+          },
+        },
+        { id: "invalid-without-output" },
       ],
-      markdown: "reassembled answer",
     });
 
-    const out = await generateHighlights("q", "some markdown", { logger });
+    const out = await generateHighlightsBatch(
+      "q",
+      [
+        { id: "0", markdown: "missing" },
+        { id: "1", markdown: "some markdown" },
+      ],
+      { logger },
+    );
 
-    expect(out).toEqual({
+    expect(out?.get("0")).toBeUndefined();
+    expect(out?.get("1")).toEqual({
       highlights: [
         { block_index: 0, score: 0.9 },
         { block_index: 2, score: 0.5 },
@@ -66,37 +114,117 @@ describe("generateHighlights", () => {
     });
   });
 
-  it("debug-logs the highlights array", async () => {
+  it("debug-logs batch coverage and highlights", async () => {
     mockFetchOnce({
-      highlights: [{ block_index: 1, score: 0.8 }],
-      markdown: "x",
+      pages: [
+        {
+          id: "page-1",
+          output: {
+            highlights: [{ block_index: 1, score: 0.8 }],
+            markdown: "x",
+          },
+        },
+      ],
     });
 
-    await generateHighlights("q", "md", { logger });
+    await generateHighlightsBatch("q", [{ id: "page-1", markdown: "md" }], {
+      logger,
+    });
 
     expect(logger.debug).toHaveBeenCalledWith(
-      "query highlights",
+      "query highlights batch",
       expect.objectContaining({
         canonicalLog: "search/highlights",
-        highlights: [{ block_index: 1, score: 0.8 }],
+        requestedPages: 1,
+        returnedPages: 1,
+        pages: [
+          {
+            id: "page-1",
+            highlights: [{ block_index: 1, score: 0.8 }],
+          },
+        ],
       }),
     );
   });
 
-  it("defaults missing fields to empty highlights and markdown", async () => {
+  it("defaults missing pages to an empty result map", async () => {
     mockFetchOnce({});
 
-    const out = await generateHighlights("q", "md", { logger });
+    const out = await generateHighlightsBatch(
+      "q",
+      [{ id: "0", markdown: "md" }],
+      { logger },
+    );
 
-    expect(out).toEqual({ highlights: [], markdown: "" });
+    expect(out).toEqual(new Map());
+  });
+
+  it("does not call the service for an empty batch", async () => {
+    const fetchMock = vi.fn();
+    vi.stubGlobal("fetch", fetchMock);
+
+    const out = await generateHighlightsBatch("q", [], { logger });
+
+    expect(fetchMock).not.toHaveBeenCalled();
+    expect(out).toEqual(new Map());
   });
 
   it("returns null when the service errors", async () => {
     mockFetchOnce({ error: "boom" }, false, 500);
 
-    const out = await generateHighlights("q", "md", { logger });
+    const out = await generateHighlightsBatch(
+      "q",
+      [{ id: "0", markdown: "md" }],
+      { logger },
+    );
 
     expect(out).toBeNull();
     expect(logger.warn).toHaveBeenCalled();
   });
+
+  it("falls back to legacy per-page calls while the old service URL is configured", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce({
+        ok: false,
+        status: 400,
+        text: async () => "legacy batch contract",
+      })
+      .mockResolvedValueOnce({
+        ok: true,
+        status: 200,
+        json: async () => ({ markdown: "first legacy highlight" }),
+        text: async () => "",
+      })
+      .mockResolvedValueOnce({
+        ok: false,
+        status: 500,
+        json: async () => ({}),
+        text: async () => "page failed",
+      });
+    vi.stubGlobal("fetch", fetchMock);
+
+    const out = await generateHighlightsBatch(
+      "q",
+      [
+        { id: "0", markdown: "first" },
+        { id: "1", markdown: "second" },
+      ],
+      { logger },
+    );
+
+    expect(fetchMock).toHaveBeenCalledTimes(3);
+    expect(fetchMock.mock.calls.map(call => call[0])).toEqual([
+      "https://highlight.test/batch_highlight",
+      "https://highlight.test/highlight",
+      "https://highlight.test/highlight",
+    ]);
+    expect(out).toEqual(
+      new Map([["0", { highlights: [], markdown: "first legacy highlight" }]]),
+    );
+    expect(logger.warn).toHaveBeenCalledWith(
+      "legacy query highlight failed",
+      expect.objectContaining({ pageId: "1" }),
+    );
+  });
 });