Skip to content

Commit 58aa669

Browse files
committed
fix(knowledge): chunk a long PDF for Azure OCR instead of refusing it
Both OCR providers cap how many pages a single request may carry, and both were handling that cap differently: one split the document to fit, the other rejected any document over it. A long PDF could therefore be ingested on one provider and not at all on the other, for a limit that belongs to a request rather than to a document. The splitting, concurrency, ordering and partial-failure rule now live in one place that both providers call, so they cannot drift apart again. A chunk that fails is dropped rather than failing the document — losing one section of a long document beats losing all of it — and every chunk failing still throws. Also drops the unpdf mock from the triage tests. It was masking real behaviour: the page count now comes from the parse metadata, so the mock was no longer needed, and while it was in place a test asserting the old page-cap refusal passed against both the old and new code.
1 parent 47cd6e1 commit 58aa669

2 files changed

Lines changed: 161 additions & 87 deletions

File tree

‎apps/sim/lib/knowledge/documents/document-processor.ts‎

Lines changed: 115 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -594,42 +594,19 @@ async function parseWithAzureMistralOCR(
594594

595595
const fileBuffer = await downloadFileForBase64(fileUrl, userId)
596596

597-
if (mimeType === 'application/pdf') {
598-
const pageCount = await getPdfPageCount(fileBuffer)
599-
if (pageCount > MISTRAL_MAX_PAGES) {
600-
throw new Error(
601-
`PDF has ${pageCount} pages, exceeding the Azure OCR limit of ${MISTRAL_MAX_PAGES}`
602-
)
603-
}
604-
logger.info('Azure Mistral OCR: PDF page count resolved', { pageCount })
605-
}
606-
607-
const base64Data = fileBuffer.toString('base64')
608-
const dataUri = `data:${mimeType};base64,${base64Data}`
609-
610597
try {
611-
const response = await retryWithExponentialBackoff(
612-
() =>
613-
makeOCRRequest(
614-
env.OCR_AZURE_ENDPOINT!,
615-
{
616-
'Content-Type': 'application/json',
617-
Authorization: `Bearer ${env.OCR_AZURE_API_KEY}`,
618-
},
619-
{
620-
model: env.OCR_AZURE_MODEL_NAME!,
621-
document: {
622-
type: 'document_url',
623-
document_url: dataUri,
624-
},
625-
include_image_base64: false,
626-
}
627-
),
628-
{ maxRetries: 3, initialDelayMs: 1000, maxDelayMs: 10000 }
629-
)
630-
631-
const ocrResult = (await response.json()) as AzureOCRResponse
632-
const content = extractPageContent(ocrResult.pages || []) || JSON.stringify(ocrResult, null, 2)
598+
/**
599+
* A PDF is chunked to the provider's page cap rather than refused for
600+
* exceeding it, matching the other OCR provider. Refusing meant a long
601+
* document could not be ingested at all, and the cap applies to a single
602+
* request, not to the document.
603+
*/
604+
const content =
605+
mimeType === 'application/pdf'
606+
? await ocrPdfInChunks(fileBuffer, 'azure-mistral', (chunk) =>
607+
recognizeWithAzureOCR(chunk.buffer, mimeType)
608+
)
609+
: await recognizeWithAzureOCR(fileBuffer, mimeType)
633610

634611
if (!content.trim()) {
635612
throw new Error('Azure Mistral OCR returned empty content')
@@ -645,6 +622,34 @@ async function parseWithAzureMistralOCR(
645622
}
646623
}
647624

625+
/** Sends one document to Azure Mistral OCR inline, as a base64 data URI. */
626+
async function recognizeWithAzureOCR(buffer: Buffer, mimeType: string): Promise<string> {
627+
const dataUri = `data:${mimeType};base64,${buffer.toString('base64')}`
628+
629+
const response = await retryWithExponentialBackoff(
630+
() =>
631+
makeOCRRequest(
632+
env.OCR_AZURE_ENDPOINT!,
633+
{
634+
'Content-Type': 'application/json',
635+
Authorization: `Bearer ${env.OCR_AZURE_API_KEY}`,
636+
},
637+
{
638+
model: env.OCR_AZURE_MODEL_NAME!,
639+
document: {
640+
type: 'document_url',
641+
document_url: dataUri,
642+
},
643+
include_image_base64: false,
644+
}
645+
),
646+
{ maxRetries: 3, initialDelayMs: 1000, maxDelayMs: 10000 }
647+
)
648+
649+
const ocrResult = (await response.json()) as AzureOCRResponse
650+
return extractPageContent(ocrResult.pages || []) || JSON.stringify(ocrResult, null, 2)
651+
}
652+
648653
async function parseWithMistralOCR(
649654
fileUrl: string,
650655
filename: string,
@@ -812,63 +817,97 @@ async function processChunk(
812817
}
813818
}
814819

815-
async function processMistralOCRInBatches(
816-
filename: string,
817-
apiKey: string,
820+
/**
821+
* Runs a PDF through OCR a chunk at a time and stitches the pages back together.
822+
*
823+
* A provider that caps how many pages one request may carry needs the document
824+
* split, and both providers cap at the same limit — so the splitting, the
825+
* concurrency, the ordering and the partial-failure rule live here once rather
826+
* than being restated per provider, where they had already drifted into one
827+
* provider chunking and the other refusing anything over the cap.
828+
*
829+
* A chunk that fails yields `null` and is dropped: losing one section of a long
830+
* document is better than losing all of it. Every chunk failing is a genuine
831+
* failure and throws.
832+
*/
833+
async function ocrPdfInChunks(
818834
pdfBuffer: Buffer,
819-
userId?: string,
820-
cloudUrl?: string
821-
): Promise<{
822-
content: string
823-
processingMethod: 'mistral-ocr'
824-
cloudUrl?: string
825-
}> {
835+
provider: string,
836+
recognize: (
837+
chunk: { buffer: Buffer; startPage: number; endPage: number },
838+
chunkIndex: number,
839+
totalChunks: number
840+
) => Promise<string | null>
841+
): Promise<string> {
826842
const totalPages = await getPdfPageCount(pdfBuffer)
827-
logger.info(`Splitting PDF into chunks`, { totalPages, maxPagesPerChunk: MISTRAL_MAX_PAGES })
828-
829843
const pdfChunks = await splitPdfIntoChunks(pdfBuffer, MISTRAL_MAX_PAGES)
830-
logger.info(
831-
`Split into ${pdfChunks.length} chunks, processing with concurrency ${MAX_CONCURRENT_CHUNKS}`
832-
)
844+
logger.info('Splitting PDF for OCR', {
845+
provider,
846+
totalPages,
847+
chunks: pdfChunks.length,
848+
maxPagesPerChunk: MISTRAL_MAX_PAGES,
849+
concurrency: MAX_CONCURRENT_CHUNKS,
850+
})
833851

834852
const results: { index: number; content: string | null }[] = []
835853

836854
for (let i = 0; i < pdfChunks.length; i += MAX_CONCURRENT_CHUNKS) {
837855
const batch = pdfChunks.slice(i, i + MAX_CONCURRENT_CHUNKS)
838-
const batchPromises = batch.map((chunk, batchIndex) =>
839-
processChunk(chunk, i + batchIndex, pdfChunks.length, filename, apiKey, userId)
840-
)
841-
842-
const batchResults = await Promise.all(batchPromises)
843-
for (const result of batchResults) {
844-
results.push(result)
845-
}
846-
847-
logger.info(
848-
`Completed batch ${Math.floor(i / MAX_CONCURRENT_CHUNKS) + 1}/${Math.ceil(pdfChunks.length / MAX_CONCURRENT_CHUNKS)}`
856+
const batchResults = await Promise.all(
857+
batch.map((chunk, batchIndex) => {
858+
const index = i + batchIndex
859+
return recognize(chunk, index, pdfChunks.length).then(
860+
(content) => ({ index, content }),
861+
(error) => {
862+
logger.warn('OCR chunk failed', {
863+
provider,
864+
chunk: index + 1,
865+
error: toError(error).message,
866+
})
867+
return { index, content: null }
868+
}
869+
)
870+
})
849871
)
872+
results.push(...batchResults)
850873
}
851874

852-
const sortedResults = results
875+
const recovered = results
853876
.sort((a, b) => a.index - b.index)
854-
.filter((r) => r.content !== null)
855-
.map((r) => r.content as string)
877+
.map((r) => r.content)
878+
.filter((content): content is string => content !== null && content.trim().length > 0)
856879

857-
if (sortedResults.length === 0) {
858-
throw new Error(
859-
`OCR failed for all ${pdfChunks.length} chunks. ` +
860-
`Large PDFs require OCR - file parser fallback would produce poor results.`
861-
)
880+
if (recovered.length === 0) {
881+
throw new Error(`OCR failed for all ${pdfChunks.length} chunks of the document`)
862882
}
863883

864-
const combinedContent = sortedResults.join('\n\n')
865-
logger.info(`Successfully processed ${sortedResults.length}/${pdfChunks.length} chunks`)
866-
867-
return {
868-
content: combinedContent,
869-
processingMethod: 'mistral-ocr',
870-
cloudUrl,
884+
if (recovered.length < pdfChunks.length) {
885+
logger.warn('OCR recovered only part of the document', {
886+
provider,
887+
recovered: recovered.length,
888+
chunks: pdfChunks.length,
889+
})
871890
}
891+
892+
return recovered.join('\n\n')
893+
}
894+
895+
async function processMistralOCRInBatches(
896+
filename: string,
897+
apiKey: string,
898+
pdfBuffer: Buffer,
899+
userId?: string,
900+
cloudUrl?: string
901+
): Promise<{
902+
content: string
903+
processingMethod: 'mistral-ocr'
904+
cloudUrl?: string
905+
}> {
906+
const content = await ocrPdfInChunks(pdfBuffer, 'mistral', (chunk, index, total) =>
907+
processChunk(chunk, index, total, filename, apiKey, userId).then((r) => r.content)
908+
)
909+
910+
return { content, processingMethod: 'mistral-ocr', cloudUrl }
872911
}
873912

874913
/**

‎apps/sim/lib/knowledge/documents/pdf-ocr-triage.test.ts‎

Lines changed: 46 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,12 @@
88
*/
99
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
1010

11-
const { mockParseBuffer, mockDownload, mockGetDocumentProxy, mockToken, mockBaseUrl } = vi.hoisted(
12-
() => ({
13-
mockParseBuffer: vi.fn(),
14-
mockDownload: vi.fn(),
15-
mockGetDocumentProxy: vi.fn(),
16-
mockToken: vi.fn(),
17-
mockBaseUrl: vi.fn(),
18-
})
19-
)
11+
const { mockParseBuffer, mockDownload, mockToken, mockBaseUrl } = vi.hoisted(() => ({
12+
mockParseBuffer: vi.fn(),
13+
mockDownload: vi.fn(),
14+
mockToken: vi.fn(),
15+
mockBaseUrl: vi.fn(),
16+
}))
2017

2118
vi.mock('@/lib/auth/internal', () => ({ generateInternalToken: mockToken }))
2219
vi.mock('@/lib/core/utils/urls', async (importOriginal) => ({
@@ -29,7 +26,6 @@ vi.mock('@/lib/file-parsers', () => ({
2926
isSupportedFileType: (extension: string) => ['pdf'].includes(extension),
3027
}))
3128
vi.mock('@/lib/uploads/utils/file-utils.server', () => ({ downloadFileFromUrl: mockDownload }))
32-
vi.mock('unpdf', () => ({ getDocumentProxy: mockGetDocumentProxy }))
3329

3430
import { env } from '@/lib/core/config/env'
3531
import { processDocument } from '@/lib/knowledge/documents/document-processor'
@@ -39,6 +35,14 @@ import { runWithKnowledgeModelInputProvenance } from '@/lib/knowledge/model-inpu
3935
const PDF_URL = 'https://example.com/Contract.pdf'
4036
const typeset = 'The Supplier shall provide the Services described herein. '.repeat(60)
4137

38+
/** A real PDF, because splitting loads the document rather than trusting metadata. */
39+
async function pdfOfPages(count: number): Promise<Buffer> {
40+
const { PDFDocument } = await import('pdf-lib')
41+
const pdf = await PDFDocument.create()
42+
for (let i = 0; i < count; i++) pdf.addPage()
43+
return Buffer.from(await pdf.save())
44+
}
45+
4246
function parse() {
4347
return runWithKnowledgeModelInputProvenance(
4448
undefined,
@@ -52,7 +56,6 @@ describe('PDF OCR triage', () => {
5256
vi.clearAllMocks()
5357
Object.assign(env, { OCR_PROVIDER: 'mistral', MISTRAL_API_KEY: 'key' })
5458
mockDownload.mockResolvedValue(Buffer.from('%PDF-1.7'))
55-
mockGetDocumentProxy.mockResolvedValue({ numPages: 2 })
5659
mockToken.mockResolvedValue('internal-token')
5760
mockBaseUrl.mockReturnValue('http://sim.local')
5861
})
@@ -107,6 +110,7 @@ describe('PDF OCR triage', () => {
107110
const result = await parse()
108111

109112
expect(result.metadata.processingMethod).toBe('mistral-ocr')
113+
// 1001 pages against a 1000-page request cap: two chunks, two requests.
110114
expect(fetchMock).toHaveBeenCalled()
111115
})
112116

@@ -145,3 +149,34 @@ describe('PDF OCR triage', () => {
145149
expect(result.metadata.processingMethod).toBe('mistral-ocr')
146150
})
147151
})
152+
153+
describe('Azure OCR chunking', () => {
154+
/**
155+
* Both providers cap how many pages one OCR request may carry. Mistral split the
156+
* document to fit; Azure refused anything over the cap, so a long PDF could not
157+
* be ingested at all. The cap belongs to a request, not to a document.
158+
*/
159+
it('splits a PDF past the page cap instead of refusing it', async () => {
160+
Object.assign(env, {
161+
OCR_PROVIDER: 'azure-mistral',
162+
OCR_AZURE_API_KEY: 'key',
163+
OCR_AZURE_ENDPOINT: 'https://example.openai.azure.com',
164+
OCR_AZURE_MODEL_NAME: 'mistral-ocr',
165+
})
166+
mockParseBuffer.mockResolvedValue({ content: '', metadata: { pageCount: 2500 } })
167+
mockDownload.mockResolvedValue(await pdfOfPages(1001))
168+
const fetchMock = vi.fn().mockResolvedValue(
169+
new Response(JSON.stringify({ pages: [{ markdown: 'Recognised page' }] }), {
170+
status: 200,
171+
headers: { 'Content-Type': 'application/json' },
172+
})
173+
)
174+
vi.stubGlobal('fetch', fetchMock)
175+
176+
const result = await parse()
177+
178+
expect(result.metadata.processingMethod).toBe('mistral-ocr')
179+
// 1001 pages against a 1000-page request cap: two chunks, two requests.
180+
expect(fetchMock).toHaveBeenCalledTimes(2)
181+
})
182+
})

0 commit comments

Comments
 (0)