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 93699d42..2b440fe4 100644 --- a/src/lib/ocr/mistral-batch.ts +++ b/src/lib/ocr/mistral-batch.ts @@ -1,7 +1,22 @@ +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; @@ -26,7 +41,7 @@ interface BatchOutputLine { response?: { status_code?: number; body?: { - pages?: OcrPage[]; + pages?: RawBatchPage[]; }; }; error?: { @@ -34,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) { @@ -51,6 +72,9 @@ 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 e88bf54f..61c88cde 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,9 @@ 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-client.ts b/src/lib/pdf/mistral-ocr-client.ts index 5a489fb5..3072441a 100644 --- a/src/lib/pdf/mistral-ocr-client.ts +++ b/src/lib/pdf/mistral-ocr-client.ts @@ -1,3 +1,4 @@ +import { mergeFigureAnnotationsIntoMarkdown } from "@/lib/pdf/ocr-figure-inline"; import { logger } from "@/lib/utils/logger"; const MAX_ATTEMPTS = 4; const RATE_LIMIT_BASE_DELAY_MS = 20_000; @@ -6,6 +7,11 @@ const TIMEOUT_BASE_DELAY_MS = 8_000; const TIMEOUT_MAX_DELAY_MS = 45_000; const RETRY_JITTER_MS = 1_500; +interface MistralOcrResponseImage { + id?: string; + image_annotation?: string | null; +} + interface MistralOcrResponsePage { index?: number; markdown?: string; @@ -13,6 +19,8 @@ interface MistralOcrResponsePage { header?: string | null; hyperlinks?: unknown[]; tables?: unknown[]; + /** Present when bbox annotations are enabled; merged into markdown and not stored on OcrPage. */ + images?: MistralOcrResponseImage[]; } export interface MistralOcrResponse { @@ -41,9 +49,11 @@ export interface MistralOcrCallResult { export function mapPages(rawPages: MistralOcrResponsePage[]): MistralOcrPage[] { return rawPages.map((page, index) => { + 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..24426aac 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,9 @@ function buildBaseBody(document: Record): Record; + const short = + typeof obj.short_description === "string" ? obj.short_description.trim() : ""; + return short || 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..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,36 +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. - * Image and table placeholders are mapped to actual content when available. + * 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( @@ -634,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}`); } @@ -692,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}`); }