import { describe, expect, it, vi } from 'vitest'; import { validateStructuredBatch, exportValidatedBatch } from '../../src/document-workbench/structured-batch'; import { structureDocumentBatch } from '../../src/document-workbench/structure-document-batch'; import { Inf0GatewayClient, type Inf0GatewayReceipt } from '../../src/document-workbench/inf0-gateway-client'; const fields = [{ name: 'name', type: 'string' as const, required: true }]; const sources = [{ id: 's1', text: 'Acme\nBeta' }]; const records = [ { value: { name: 'Acme' }, evidence: [{ field: 'name', sourceId: 's1', start: 0, end: 4, transform: 'exact' as const }] }, { value: { name: 'Beta' }, evidence: [{ field: 'name', sourceId: 's1', start: 5, end: 9, transform: 'exact' as const }] }, ]; const request = { context: { tenantId: 'a', projectId: 'p', actorId: 'u', capability: 'inf0.complete', correlationId: 'c' }, admission: { tenantId: 'a', projectId: 'p', jobId: 'j', modelId: 'local_only', dataClass: 'high', maxInputTokens: 200, maxOutputTokens: 100, maximumCredits: 5 }, policy: { tenantId: 'a', projectId: 'p', allowedModelIds: ['local_only'], allowedDataClasses: ['high'], destination: 'internal', externalTransferAllowed: false, maxInputTokens: 200, maxOutputTokens: 100 }, reservation: { tenantId: 'a', projectId: 'p', jobId: 'j', status: 'reserved', credits: 5 }, } satisfies Omit[0], 'messages' | 'inputTokens'>; const receipt: Inf0GatewayReceipt = { status: 'model_stopped', content: JSON.stringify({ records }), modelExecuted: 'm1', providerKind: 'local', promptTokens: 50, completionTokens: 20, totalTokens: 70 }; describe('evidence-bound document batches', () => { it('serializes every supported row in order to JSON and CSV', () => { const json = exportValidatedBatch({ format: 'json', fields, sources, records }); expect(json.status).toBe('ready'); if (json.status === 'ready') expect(JSON.parse(json.content)).toEqual([{ name: 'Acme' }, { name: 'Beta' }]); expect(exportValidatedBatch({ format: 'csv', fields, sources, records })).toEqual({ status: 'ready', mediaType: 'text/csv', content: '"name"\r\n"Acme"\r\n"Beta"\r\n', }); }); it('rejects the entire output when one row is unsupported or malformed', () => { const altered = [records[0], { ...records[1]!, value: { name: 'Invented' } }]; expect(exportValidatedBatch({ format: 'json', fields, sources, records: altered })).toMatchObject({ status: 'needs_review', issues: [{ row: 1, field: 'name', code: 'evidence_mismatch' }] }); expect(validateStructuredBatch({ fields, sources, records: [null] })).toMatchObject({ status: 'needs_review', issues: [{ row: 0, code: 'invalid_record' }] }); expect(validateStructuredBatch({ fields, sources, records: [] })).toMatchObject({ status: 'needs_review' }); expect(validateStructuredBatch({ fields, sources, records: Array(201).fill(records[0]) })).toMatchObject({ status: 'needs_review' }); }); it('blocks spreadsheet formulas in any row while JSON preserves the literal value', () => { const source = [{ id: 's1', text: 'Acme\n=1+1' }]; const unsafe = [records[0]!, { value: { name: '=1+1' }, evidence: [{ field: 'name', sourceId: 's1', start: 5, end: 9, transform: 'exact' }] }]; expect(exportValidatedBatch({ format: 'csv', fields, sources: source, records: unsafe })).toMatchObject({ status: 'needs_review', issues: [{ row: 1, code: 'csv_unsafe_cell' }] }); expect(exportValidatedBatch({ format: 'json', fields, sources: source, records: unsafe }).status).toBe('ready'); }); it('preserves quoted multiline headers and cells in a multi-row CSV', () => { const f = [{ name: 'company\r\nname', type: 'string' as const, required: true }]; const s = [{ id: 's', text: 'A\r\nB,C' }]; const r = [{ value: { 'company\r\nname': 'A\r\nB' }, evidence: [{ field: 'company\r\nname', sourceId: 's', start: 0, end: 4, transform: 'exact' }] }, { value: { 'company\r\nname': ',C' }, evidence: [{ field: 'company\r\nname', sourceId: 's', start: 4, end: 6, transform: 'exact' }] }]; expect(exportValidatedBatch({ format: 'csv', fields: f, sources: s, records: r })).toMatchObject({ status: 'ready', content: '"company\r\nname"\r\n"A\r\nB"\r\n",C"\r\n' }); }); it('extracts multiple records and keeps their evidence tied to the original source', async () => { const mutable = structuredClone(sources); const complete = vi.fn(async () => receipt); const countTokens = vi.fn(async (messages: Parameters[0]['messages']) => { expect(Object.isFrozen(messages[1])).toBe(true); mutable[0]!.text = 'Evil\nBeta'; return 50; }); const result = await structureDocumentBatch({ request, fields, sources: mutable, instructions: 'Extract companies', countTokens, gateway: { complete } }); expect(result).toMatchObject({ status: 'mechanically_validated', records }); expect(result.model).not.toHaveProperty('content'); expect(complete.mock.calls[0]?.[0]?.messages).toEqual(countTokens.mock.calls[0]?.[0]); }); it('runs a batch through the real INF0 client with scoped policy and usage checks', async () => { const fetchImpl = vi.fn(async () => Response.json({ content: receipt.content, modelExecuted: 'm1', finishReason: 'stop', routeDecision: { route: 'local_only', providerKind: 'local' }, cost: { promptTokens: 50, completionTokens: 20, totalTokens: 70 } })); const gateway = new Inf0GatewayClient({ baseUrl: 'https://inf0.example.test', serviceToken: 'test', fetchImpl }); const result = await structureDocumentBatch({ request, fields, sources, instructions: 'Extract companies', countTokens: async () => 50, gateway }); expect(result).toMatchObject({ status: 'mechanically_validated', records }); if (result.status !== 'mechanically_validated') throw new Error('Expected a verified batch'); const exported = exportValidatedBatch({ format: 'json', fields, sources, records: result.records }); expect(exported.status).toBe('ready'); if (exported.status === 'ready') expect(JSON.parse(exported.content)).toEqual([{ name: 'Acme' }, { name: 'Beta' }]); expect(fetchImpl).toHaveBeenCalledOnce(); }); it('requires review for CSV if any optional field is absent in any row', () => { const optional = [...fields, { name: 'reference', type: 'string' as const, required: false }]; expect(exportValidatedBatch({ format: 'json', fields: optional, sources, records }).status).toBe('ready'); expect(exportValidatedBatch({ format: 'csv', fields: optional, sources, records })).toMatchObject({ status: 'needs_review', issues: [{ row: 0, field: 'reference', code: 'csv_missing_value' }], }); }); it('routes truncated, empty, malformed and unsupported model results to review', async () => { for (const r of [ { ...receipt, status: 'needs_review' as const }, { ...receipt, content: '{' }, { ...receipt, content: JSON.stringify({ records: [] }) }, { ...receipt, content: JSON.stringify({ records: [{ ...records[0], value: { name: 'Wrong' } }] }) }, ]) { const result = await structureDocumentBatch({ request, fields, sources, instructions: 'Extract companies', countTokens: async () => 50, gateway: { complete: async () => r } }); expect(result.status).toBe('needs_review'); } }); });