From 5adc7506b64a4a9c02e2e7909201431ba4a1e502 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:15:07 -0400 Subject: [PATCH 1/2] feat(ocr): inline Mistral figure short descriptions without base64 Request bbox annotations with short_description schema, merge into page markdown at ingest, and skip persisting separate image payloads. Made-with: Cursor --- src/lib/ocr/mistral-batch.ts | 3 + src/lib/ocr/mistral-direct.ts | 3 + src/lib/pdf/mistral-ocr-client.ts | 12 +++- src/lib/pdf/mistral-ocr.ts | 3 + src/lib/pdf/ocr-figure-inline.ts | 74 +++++++++++++++++++++++ src/lib/utils/format-workspace-context.ts | 5 +- 6 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 src/lib/pdf/ocr-figure-inline.ts diff --git a/src/lib/ocr/mistral-batch.ts b/src/lib/ocr/mistral-batch.ts index 93699d42..40a90345 100644 --- a/src/lib/ocr/mistral-batch.ts +++ b/src/lib/ocr/mistral-batch.ts @@ -1,3 +1,4 @@ +import { MISTRAL_BBOX_ANNOTATION_FORMAT } from "@/lib/pdf/ocr-figure-inline"; import { logger } from "@/lib/utils/logger"; import { getOcrConfig } from "./config"; import type { OcrCandidate, OcrItemFailureResult, OcrItemResult, OcrItemSuccessResult, OcrPage } from "./types"; @@ -51,6 +52,8 @@ function buildBatchRequest(candidate: OcrCandidate) { : { document_url: candidate.fileUrl }, extract_header: true, extract_footer: true, + bbox_annotation_format: MISTRAL_BBOX_ANNOTATION_FORMAT, + include_image_base64: false, }, }; } diff --git a/src/lib/ocr/mistral-direct.ts b/src/lib/ocr/mistral-direct.ts index e88bf54f..73ee63fc 100644 --- a/src/lib/ocr/mistral-direct.ts +++ b/src/lib/ocr/mistral-direct.ts @@ -1,3 +1,4 @@ +import { MISTRAL_BBOX_ANNOTATION_FORMAT } from "@/lib/pdf/ocr-figure-inline"; import { logger } from "@/lib/utils/logger"; import { callMistralOcr } from "@/lib/pdf/mistral-ocr-client"; import { getOcrConfig } from "./config"; @@ -11,6 +12,8 @@ function buildDocumentBody(itemType: OcrItemType, fileUrl: string): Record { + const rawMd = page.markdown ?? ""; + const markdown = mergeFigureAnnotationsIntoMarkdown(rawMd, page.images); const mapped: MistralOcrPage = { index: page.index ?? index, - markdown: page.markdown ?? "", + markdown, }; if (page.footer) mapped.footer = page.footer; if (page.header) mapped.header = page.header; diff --git a/src/lib/pdf/mistral-ocr.ts b/src/lib/pdf/mistral-ocr.ts index 900690b2..ecb40bca 100644 --- a/src/lib/pdf/mistral-ocr.ts +++ b/src/lib/pdf/mistral-ocr.ts @@ -2,6 +2,7 @@ * Mistral OCR integration for documents and images via the direct OCR API. */ +import { MISTRAL_BBOX_ANNOTATION_FORMAT } from "@/lib/pdf/ocr-figure-inline"; import { logger } from "@/lib/utils/logger"; import { callMistralOcr } from "./mistral-ocr-client"; import { getOcrConfig } from "@/lib/ocr/config"; @@ -17,6 +18,8 @@ function buildBaseBody(document: Record): Record; + const short = + typeof obj.short_description === "string" ? obj.short_description.trim() : ""; + const summary = typeof obj.summary === "string" ? obj.summary.trim() : ""; + const body = short || summary; + return body || null; + } catch { + return null; + } +} + +/** + * Replaces `![alt](id)` image placeholders with inline figure descriptions. + * Unknown ids or missing annotations are left unchanged. + */ +export function mergeFigureAnnotationsIntoMarkdown( + markdown: string, + images?: ReadonlyArray<{ + id?: string; + image_annotation?: string | null; + }> | null, +): string { + if (!images?.length) return markdown; + + const byId = new Map(); + for (const img of images) { + const id = img.id; + if (!id) continue; + const text = formatAnnotationForInline(img.image_annotation); + if (text) byId.set(id, text); + } + if (byId.size === 0) return markdown; + + let out = markdown; + for (const [id, replacement] of byId) { + const pattern = new RegExp(`!\\[[^\\]]*\\]\\(${escapeRegex(id)}\\)`, "g"); + out = out.replace(pattern, `\n\n${replacement}\n\n`); + } + return out; +} diff --git a/src/lib/utils/format-workspace-context.ts b/src/lib/utils/format-workspace-context.ts index b19da3bb..f785931c 100644 --- a/src/lib/utils/format-workspace-context.ts +++ b/src/lib/utils/format-workspace-context.ts @@ -592,7 +592,10 @@ function escapeRegex(s: string): string { * If OCR pages are available, include them so the agent can reason about the PDF * without needing any separate file-processing tool. * OCR pages output markdown as proper lines (one line per line) instead of JSON blobs. - * Image and table placeholders are mapped to actual content when available. + * Figure descriptions are merged into page markdown at OCR time. Tables usually + * appear as markdown in `page.markdown` already; if the response uses separate + * table chunks with `[id](id)` placeholders, `replaceOcrPlaceholders` inlines + * the matching `tables[].content` (format depends on Mistral `table_format`). * Optionally filter by pageStart/pageEnd (1-indexed, inclusive). */ function formatPdfDetailsFull( From 9d45553ea604c56c8c1b399174cfdb775e27c2b7 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Mon, 6 Apr 2026 11:59:07 -0400 Subject: [PATCH 2/2] feat(ocr): align Mistral OCR with inline tables and merged figure text - Set table_format to null on all OCR requests so tables stay in page markdown. - Merge bbox figure annotations into markdown on the batch path (same as direct). - Remove table placeholder inlining from workspace context; rely on inline tables. - Tighten ocr-figure-inline (replacer-safe replace, drop dead summary fallback). - Add unit tests for mergeFigureAnnotationsIntoMarkdown and batch merge behavior. Made-with: Cursor --- src/lib/ocr/__tests__/mistral-batch.test.ts | 47 ++++++++++++ src/lib/ocr/mistral-batch.ts | 35 +++++++-- src/lib/ocr/mistral-direct.ts | 1 + .../pdf/__tests__/ocr-figure-inline.test.ts | 72 +++++++++++++++++++ src/lib/pdf/mistral-ocr.ts | 1 + src/lib/pdf/ocr-figure-inline.ts | 6 +- src/lib/utils/format-workspace-context.ts | 48 ++----------- 7 files changed, 156 insertions(+), 54 deletions(-) create mode 100644 src/lib/pdf/__tests__/ocr-figure-inline.test.ts diff --git a/src/lib/ocr/__tests__/mistral-batch.test.ts b/src/lib/ocr/__tests__/mistral-batch.test.ts index 32be9e21..8cf3346c 100644 --- a/src/lib/ocr/__tests__/mistral-batch.test.ts +++ b/src/lib/ocr/__tests__/mistral-batch.test.ts @@ -64,6 +64,53 @@ describe("mapBatchResults", () => { ]); }); + it("merges figure annotations into markdown like the direct OCR path", () => { + const candidates: OcrCandidate[] = [ + { itemId: "doc-1", itemType: "file", fileUrl: "https://example.com/a.pdf" }, + ]; + + const results = mapBatchResults( + candidates, + [ + { + custom_id: "doc-1", + response: { + status_code: 200, + body: { + pages: [ + { + markdown: "Text\n\n![f](fig-1)\n\nMore", + images: [ + { + id: "fig-1", + image_annotation: JSON.stringify({ + short_description: "A diagram of the system.", + }), + }, + ], + }, + ], + }, + }, + error: null, + }, + ] as any, + [], + ); + + expect(results[0]).toMatchObject({ + ok: true, + pages: [ + { + index: 0, + markdown: expect.stringContaining("A diagram of the system."), + }, + ], + }); + const page = (results[0] as { ok: true; pages: { markdown: string }[] }).pages[0]; + expect(page.markdown).not.toMatch(/!\[[^\]]*\]\(fig-1\)/); + }); + it("uses error file entries when present", () => { const candidates: OcrCandidate[] = [ { itemId: "img-2", itemType: "image", fileUrl: "https://example.com/d.png" }, diff --git a/src/lib/ocr/mistral-batch.ts b/src/lib/ocr/mistral-batch.ts index 40a90345..2b440fe4 100644 --- a/src/lib/ocr/mistral-batch.ts +++ b/src/lib/ocr/mistral-batch.ts @@ -1,8 +1,22 @@ -import { MISTRAL_BBOX_ANNOTATION_FORMAT } from "@/lib/pdf/ocr-figure-inline"; +import { + mergeFigureAnnotationsIntoMarkdown, + MISTRAL_BBOX_ANNOTATION_FORMAT, +} from "@/lib/pdf/ocr-figure-inline"; import { logger } from "@/lib/utils/logger"; import { getOcrConfig } from "./config"; import type { OcrCandidate, OcrItemFailureResult, OcrItemResult, OcrItemSuccessResult, OcrPage } from "./types"; +/** Raw page shape from Mistral batch OCR JSONL (includes `images` before normalization to `OcrPage`). */ +interface RawBatchPage { + index?: number; + markdown?: string; + footer?: string | null; + header?: string | null; + hyperlinks?: unknown[]; + tables?: unknown[]; + images?: Array<{ id?: string; image_annotation?: string | null }>; +} + const POLL_INTERVAL_MS = 2_000; const MAX_POLL_MS = 10 * 60 * 1000; @@ -27,7 +41,7 @@ interface BatchOutputLine { response?: { status_code?: number; body?: { - pages?: OcrPage[]; + pages?: RawBatchPage[]; }; }; error?: { @@ -35,11 +49,17 @@ interface BatchOutputLine { } | null; } -function normalizePages(pages: OcrPage[] | undefined): OcrPage[] { - return (pages ?? []).map((page, index) => ({ - ...page, - index, - })); +function normalizePages(pages: RawBatchPage[] | undefined): OcrPage[] { + return (pages ?? []).map((page, index) => { + const rawMd = page.markdown ?? ""; + const markdown = mergeFigureAnnotationsIntoMarkdown(rawMd, page.images); + const out: OcrPage = { index, markdown }; + if (page.footer) out.footer = page.footer; + if (page.header) out.header = page.header; + if (page.hyperlinks?.length) out.hyperlinks = page.hyperlinks; + if (page.tables?.length) out.tables = page.tables; + return out; + }); } function buildBatchRequest(candidate: OcrCandidate) { @@ -52,6 +72,7 @@ function buildBatchRequest(candidate: OcrCandidate) { : { document_url: candidate.fileUrl }, extract_header: true, extract_footer: true, + table_format: null, bbox_annotation_format: MISTRAL_BBOX_ANNOTATION_FORMAT, include_image_base64: false, }, diff --git a/src/lib/ocr/mistral-direct.ts b/src/lib/ocr/mistral-direct.ts index 73ee63fc..61c88cde 100644 --- a/src/lib/ocr/mistral-direct.ts +++ b/src/lib/ocr/mistral-direct.ts @@ -12,6 +12,7 @@ function buildDocumentBody(itemType: OcrItemType, fileUrl: string): Record { + it("replaces ![alt](id) placeholders with figure descriptions", () => { + const md = "Intro\n\n![fig](img-1)\n\nOutro"; + const images = [ + { + id: "img-1", + image_annotation: JSON.stringify({ short_description: "A bar chart of sales." }), + }, + ]; + expect(mergeFigureAnnotationsIntoMarkdown(md, images)).toBe( + "Intro\n\n\n\nA bar chart of sales.\n\n\n\nOutro", + ); + }); + + it("skips image entries without id", () => { + const md = "![x](keep-me)"; + expect( + mergeFigureAnnotationsIntoMarkdown(md, [ + { image_annotation: JSON.stringify({ short_description: "ignored" }) }, + ]), + ).toBe(md); + }); + + it("ignores invalid JSON in image_annotation without throwing", () => { + const md = "![a](id1)"; + expect( + mergeFigureAnnotationsIntoMarkdown(md, [{ id: "id1", image_annotation: "not-json{" }]), + ).toBe(md); + }); + + it("ignores annotations with empty short_description", () => { + const md = "![a](id1)"; + expect( + mergeFigureAnnotationsIntoMarkdown(md, [ + { id: "id1", image_annotation: JSON.stringify({ short_description: "" }) }, + ]), + ).toBe(md); + }); + + it("leaves markdown unchanged when no placeholder matches an id", () => { + const md = "No figures here."; + const images = [ + { + id: "missing-from-md", + image_annotation: JSON.stringify({ short_description: "orphan" }), + }, + ]; + expect(mergeFigureAnnotationsIntoMarkdown(md, images)).toBe(md); + }); + + it("inserts descriptions literally when they contain $ (no replacement pattern corruption)", () => { + const md = "![f](img-1)"; + const desc = "Costs are $5 and $10"; + const images = [ + { + id: "img-1", + image_annotation: JSON.stringify({ short_description: desc }), + }, + ]; + expect(mergeFigureAnnotationsIntoMarkdown(md, images)).toContain(desc); + }); + + it("returns original markdown when images is empty or undefined", () => { + const md = "x"; + expect(mergeFigureAnnotationsIntoMarkdown(md, [])).toBe(md); + expect(mergeFigureAnnotationsIntoMarkdown(md, undefined)).toBe(md); + expect(mergeFigureAnnotationsIntoMarkdown(md, null)).toBe(md); + }); +}); diff --git a/src/lib/pdf/mistral-ocr.ts b/src/lib/pdf/mistral-ocr.ts index ecb40bca..24426aac 100644 --- a/src/lib/pdf/mistral-ocr.ts +++ b/src/lib/pdf/mistral-ocr.ts @@ -18,6 +18,7 @@ function buildBaseBody(document: Record): Record; const short = typeof obj.short_description === "string" ? obj.short_description.trim() : ""; - const summary = typeof obj.summary === "string" ? obj.summary.trim() : ""; - const body = short || summary; - return body || null; + return short || null; } catch { return null; } @@ -68,7 +66,7 @@ export function mergeFigureAnnotationsIntoMarkdown( let out = markdown; for (const [id, replacement] of byId) { const pattern = new RegExp(`!\\[[^\\]]*\\]\\(${escapeRegex(id)}\\)`, "g"); - out = out.replace(pattern, `\n\n${replacement}\n\n`); + out = out.replace(pattern, () => `\n\n${replacement}\n\n`); } return out; } diff --git a/src/lib/utils/format-workspace-context.ts b/src/lib/utils/format-workspace-context.ts index f785931c..478e6921 100644 --- a/src/lib/utils/format-workspace-context.ts +++ b/src/lib/utils/format-workspace-context.ts @@ -551,11 +551,7 @@ export function formatOcrPagesAsMarkdown( const pageNum = page.index + 1; lines.push(`--- Page ${pageNum} ---`); if (page.header) lines.push(`Header: ${page.header}`); - const rawMd = page.markdown ?? ""; - const md = replaceOcrPlaceholders( - rawMd, - page.tables as Array<{ id?: string; content?: string }> | undefined, - ); + const md = page.markdown ?? ""; for (const line of md.split(/\r?\n/)) lines.push(line); if (page.footer) lines.push(`Footer: ${page.footer}`); lines.push(""); @@ -563,39 +559,13 @@ export function formatOcrPagesAsMarkdown( return lines.join("\n").trimEnd(); } -/** Replaces table placeholders [id](id) with table content. */ -function replaceOcrPlaceholders( - markdown: string, - tables?: Array<{ id?: string; content?: string }>, -): string { - let out = markdown; - - for (const tbl of tables ?? []) { - const id = tbl.id; - const content = tbl.content; - if (!id || !content) continue; - out = out.replace( - new RegExp(`\\[${escapeRegex(id)}\\]\\(${escapeRegex(id)}\\)`, "g"), - `\n\n[Table ${id}]\n${content}\n\n`, - ); - } - - return out; -} - -function escapeRegex(s: string): string { - return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); -} - /** * Formats PDF details with FULL content * If OCR pages are available, include them so the agent can reason about the PDF * without needing any separate file-processing tool. * OCR pages output markdown as proper lines (one line per line) instead of JSON blobs. - * Figure descriptions are merged into page markdown at OCR time. Tables usually - * appear as markdown in `page.markdown` already; if the response uses separate - * table chunks with `[id](id)` placeholders, `replaceOcrPlaceholders` inlines - * the matching `tables[].content` (format depends on Mistral `table_format`). + * Figure descriptions are merged into page markdown at OCR time. Tables are inline + * in `page.markdown` (`table_format: null` on OCR requests). * Optionally filter by pageStart/pageEnd (1-indexed, inclusive). */ function formatPdfDetailsFull( @@ -637,11 +607,7 @@ function formatPdfDetailsFull( const pageNum = page.index + 1; lines.push(` --- Page ${pageNum} ---`); if (page.header) lines.push(` Header: ${page.header}`); - const rawMd = page.markdown ?? ""; - const md = replaceOcrPlaceholders( - rawMd, - page.tables as Array<{ id?: string; content?: string }> | undefined, - ); + const md = page.markdown ?? ""; for (const line of md.split(/\r?\n/)) { lines.push(` ${line}`); } @@ -695,11 +661,7 @@ function formatImageDetailsFull(data: ImageData): string[] { ); for (const page of data.ocrPages) { if (page.header) lines.push(` Header: ${page.header}`); - const rawMd = page.markdown ?? ""; - const md = replaceOcrPlaceholders( - rawMd, - page.tables as Array<{ id?: string; content?: string }> | undefined, - ); + const md = page.markdown ?? ""; for (const line of md.split(/\r?\n/)) { lines.push(` ${line}`); }