Files
ht-docker-ai/test/test.bankstatements.nanonets.ts

531 lines
17 KiB
TypeScript
Raw Normal View History

2026-01-18 15:54:16 +00:00
/**
* Bank statement extraction using Nanonets-OCR-s + Qwen3 (two-stage pipeline)
*
* Stage 1: Nanonets-OCR-s converts document pages to markdown (its strength)
* Stage 2: Qwen3 extracts structured JSON from the combined markdown
*
* This leverages each model's strengths:
* - Nanonets: Document OCR with semantic tags
* - Qwen3: Text understanding and JSON extraction
*/
import { tap, expect } from '@git.zone/tstest/tapbundle';
import * as fs from 'fs';
import * as path from 'path';
import { execSync } from 'child_process';
import * as os from 'os';
import { ensureNanonetsOcr, ensureMiniCpm } from './helpers/docker.js';
const NANONETS_URL = 'http://localhost:8000/v1';
const NANONETS_MODEL = 'nanonets/Nanonets-OCR-s';
const OLLAMA_URL = 'http://localhost:11434';
const QWEN_MODEL = 'qwen3:8b';
interface ITransaction {
date: string;
counterparty: string;
amount: number;
}
// Nanonets-specific prompt for document OCR to markdown
const NANONETS_OCR_PROMPT = `Extract the text from the above document as if you were reading it naturally.
Return the tables in html format.
Return the equations in LaTeX representation.
If there is an image in the document and image caption is not present, add a small description inside <img></img> tag.
Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>.
Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number>.`;
// JSON extraction prompt for Qwen3
const JSON_EXTRACTION_PROMPT = `You are a financial data extractor. Below is a bank statement converted to text/markdown. Extract ALL transactions from it as a JSON array.
IMPORTANT RULES:
1. Each transaction has: date, description/counterparty, and an amount
2. Amount is NEGATIVE for money going OUT (debits, payments, withdrawals)
3. Amount is POSITIVE for money coming IN (credits, deposits, refunds)
4. Date format: YYYY-MM-DD
5. Do NOT include: opening balance, closing balance, subtotals, headers, or summary rows
6. Only include actual transactions with a specific date and amount
Return ONLY this JSON format, no explanation:
[
{"date": "2021-06-01", "counterparty": "COMPANY NAME", "amount": -25.99},
{"date": "2021-06-02", "counterparty": "DEPOSIT FROM", "amount": 100.00}
]
BANK STATEMENT TEXT:
`;
/**
* Convert PDF to PNG images using ImageMagick
*/
function convertPdfToImages(pdfPath: string): string[] {
const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'pdf-convert-'));
const outputPattern = path.join(tempDir, 'page-%d.png');
try {
// Use 150 DPI to keep images within model's context length
execSync(
`convert -density 150 -quality 90 "${pdfPath}" -background white -alpha remove "${outputPattern}"`,
{ stdio: 'pipe' }
);
const files = fs.readdirSync(tempDir).filter((f: string) => f.endsWith('.png')).sort();
const images: string[] = [];
for (const file of files) {
const imagePath = path.join(tempDir, file);
const imageData = fs.readFileSync(imagePath);
images.push(imageData.toString('base64'));
}
return images;
} finally {
fs.rmSync(tempDir, { recursive: true, force: true });
}
}
/**
* Stage 1: Convert a single page to markdown using Nanonets-OCR-s
*/
async function convertPageToMarkdown(image: string, pageNum: number): Promise<string> {
console.log(` [Nanonets] Converting page ${pageNum} to markdown...`);
const startTime = Date.now();
const response = await fetch(`${NANONETS_URL}/chat/completions`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': 'Bearer dummy',
},
body: JSON.stringify({
model: NANONETS_MODEL,
messages: [{
role: 'user',
content: [
{ type: 'image_url', image_url: { url: `data:image/png;base64,${image}` }},
{ type: 'text', text: NANONETS_OCR_PROMPT },
],
}],
max_tokens: 4096,
temperature: 0.0,
}),
});
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
if (!response.ok) {
const errorText = await response.text();
console.log(` [Nanonets] ERROR page ${pageNum}: ${response.status} - ${errorText}`);
throw new Error(`Nanonets API error: ${response.status}`);
}
const data = await response.json();
const content = (data.choices?.[0]?.message?.content || '').trim();
console.log(` [Nanonets] Page ${pageNum} converted (${elapsed}s, ${content.length} chars)`);
return content;
}
/**
* Stage 1: Convert all pages to markdown using Nanonets-OCR-s
*/
async function convertDocumentToMarkdown(images: string[]): Promise<string> {
console.log(` [Stage 1] Converting ${images.length} page(s) to markdown with Nanonets-OCR-s...`);
const markdownPages: string[] = [];
for (let i = 0; i < images.length; i++) {
const markdown = await convertPageToMarkdown(images[i], i + 1);
markdownPages.push(`--- PAGE ${i + 1} ---\n${markdown}`);
}
const fullMarkdown = markdownPages.join('\n\n');
console.log(` [Stage 1] Complete: ${fullMarkdown.length} chars total`);
return fullMarkdown;
}
/**
* Ensure Qwen3 model is available
*/
async function ensureQwen3(): Promise<boolean> {
try {
const response = await fetch(`${OLLAMA_URL}/api/tags`);
if (response.ok) {
const data = await response.json();
const models = data.models || [];
if (models.some((m: { name: string }) => m.name === QWEN_MODEL)) {
console.log(` [Ollama] Model available: ${QWEN_MODEL}`);
return true;
}
}
} catch {
return false;
}
console.log(` [Ollama] Pulling ${QWEN_MODEL}...`);
const pullResponse = await fetch(`${OLLAMA_URL}/api/pull`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ name: QWEN_MODEL, stream: false }),
});
return pullResponse.ok;
}
/**
* Stage 2: Extract transactions from markdown using Qwen3
*/
async function extractTransactionsFromMarkdown(markdown: string, queryId: string): Promise<ITransaction[]> {
console.log(` [${queryId}] Sending markdown to ${QWEN_MODEL}...`);
const startTime = Date.now();
const response = await fetch(`${OLLAMA_URL}/api/chat`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
model: QWEN_MODEL,
messages: [{
role: 'user',
content: JSON_EXTRACTION_PROMPT + markdown,
}],
stream: false,
options: {
num_predict: 8000,
temperature: 0.1,
},
}),
});
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
if (!response.ok) {
console.log(` [${queryId}] ERROR: ${response.status} (${elapsed}s)`);
throw new Error(`Ollama API error: ${response.status}`);
}
const data = await response.json();
const content = (data.message?.content || '').trim();
console.log(` [${queryId}] Response received (${elapsed}s, ${content.length} chars)`);
return parseJsonResponse(content, queryId);
}
/**
* Sanitize JSON string - fix common issues from LLM output
*/
function sanitizeJson(jsonStr: string): string {
let s = jsonStr;
// Fix +number (e.g., +93.80 -> 93.80) - JSON doesn't allow + prefix
s = s.replace(/"amount"\s*:\s*\+/g, '"amount": ');
s = s.replace(/:\s*\+(\d)/g, ': $1');
// Fix European number format with thousands separator
s = s.replace(/"amount"\s*:\s*(-?)(\d{1,3})\.(\d{3})\.(\d{2})\b/g, '"amount": $1$2$3.$4');
// Fix trailing commas before ] or }
s = s.replace(/,\s*([}\]])/g, '$1');
// Fix unescaped newlines/tabs inside strings
s = s.replace(/"([^"\\]*)\n([^"]*)"/g, '"$1 $2"');
s = s.replace(/"([^"\\]*)\t([^"]*)"/g, '"$1 $2"');
// Remove control characters
s = s.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F]/g, ' ');
return s;
}
/**
* Parse amount from various formats
*/
function parseAmount(value: unknown): number {
if (typeof value === 'number') return value;
if (typeof value !== 'string') return 0;
let s = value.replace(/[€$£\s]/g, '').replace('', '-').replace('', '-');
// European format: comma is decimal
if (s.includes(',') && s.indexOf(',') > s.lastIndexOf('.')) {
s = s.replace(/\./g, '').replace(',', '.');
} else {
s = s.replace(/,/g, '');
}
return parseFloat(s) || 0;
}
/**
* Parse JSON response into transactions
*/
function parseJsonResponse(response: string, queryId: string): ITransaction[] {
console.log(` [${queryId}] Parsing response...`);
// Remove thinking tags if present (Qwen3 may include <think>...</think>)
let cleanResponse = response.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
// Try to find JSON in markdown code block
const codeBlockMatch = cleanResponse.match(/```(?:json)?\s*([\s\S]*?)```/);
let jsonStr = codeBlockMatch ? codeBlockMatch[1].trim() : cleanResponse;
// Sanitize JSON
jsonStr = sanitizeJson(jsonStr);
try {
const parsed = JSON.parse(jsonStr);
if (Array.isArray(parsed)) {
const txs = parsed.map(tx => ({
date: String(tx.date || ''),
counterparty: String(tx.counterparty || tx.description || ''),
amount: parseAmount(tx.amount),
}));
console.log(` [${queryId}] Parsed ${txs.length} transactions`);
return txs;
}
} catch (e) {
console.log(` [${queryId}] Direct parse failed: ${(e as Error).message}`);
// Try to find JSON array pattern
const arrayMatch = jsonStr.match(/\[[\s\S]*\]/);
if (arrayMatch) {
try {
const parsed = JSON.parse(sanitizeJson(arrayMatch[0]));
if (Array.isArray(parsed)) {
const txs = parsed.map(tx => ({
date: String(tx.date || ''),
counterparty: String(tx.counterparty || tx.description || ''),
amount: parseAmount(tx.amount),
}));
console.log(` [${queryId}] Parsed ${txs.length} transactions (array match)`);
return txs;
}
} catch (e2) {
console.log(` [${queryId}] Array parse failed: ${(e2 as Error).message}`);
}
}
}
console.log(` [${queryId}] PARSE FAILED - returning empty array`);
return [];
}
/**
* Compare two transaction arrays for consensus
*/
function transactionArraysMatch(a: ITransaction[], b: ITransaction[]): boolean {
if (a.length !== b.length) return false;
for (let i = 0; i < a.length; i++) {
const dateMatch = a[i].date === b[i].date;
const amountMatch = Math.abs(a[i].amount - b[i].amount) < 0.01;
if (!dateMatch || !amountMatch) return false;
}
return true;
}
/**
* Stage 2: Extract transactions using Qwen3 with consensus
*/
async function extractWithConsensus(markdown: string): Promise<ITransaction[]> {
const MAX_ATTEMPTS = 3;
console.log(` [Stage 2] Extracting transactions with ${QWEN_MODEL} (consensus)...`);
for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
console.log(`\n [Stage 2] --- Attempt ${attempt}/${MAX_ATTEMPTS} ---`);
// Extract twice in parallel
const [txs1, txs2] = await Promise.all([
extractTransactionsFromMarkdown(markdown, `A${attempt}Q1`),
extractTransactionsFromMarkdown(markdown, `A${attempt}Q2`),
]);
console.log(` [Stage 2] Results: Q1=${txs1.length} txs, Q2=${txs2.length} txs`);
if (txs1.length > 0 && transactionArraysMatch(txs1, txs2)) {
console.log(` [Stage 2] CONSENSUS REACHED: ${txs1.length} transactions`);
return txs1;
}
console.log(` [Stage 2] NO CONSENSUS`);
}
// Fallback: use last response
console.log(`\n [Stage 2] === FALLBACK ===`);
const fallback = await extractTransactionsFromMarkdown(markdown, 'FALLBACK');
console.log(` [Stage 2] ~ FALLBACK RESULT: ${fallback.length} transactions`);
return fallback;
}
/**
* Full pipeline: PDF -> Images -> Markdown -> JSON
*/
async function extractTransactions(images: string[]): Promise<ITransaction[]> {
// Stage 1: Convert to markdown
const markdown = await convertDocumentToMarkdown(images);
// Stage 2: Extract transactions with consensus
const transactions = await extractWithConsensus(markdown);
// Log all transactions
console.log(`\n [Result] Extracted ${transactions.length} transactions:`);
for (let i = 0; i < transactions.length; i++) {
const tx = transactions[i];
console.log(` ${(i + 1).toString().padStart(2)}. ${tx.date} | ${tx.counterparty.substring(0, 30).padEnd(30)} | ${tx.amount >= 0 ? '+' : ''}${tx.amount.toFixed(2)}`);
}
return transactions;
}
/**
* Compare extracted transactions against expected
*/
function compareTransactions(
extracted: ITransaction[],
expected: ITransaction[]
): { matches: number; total: number; errors: string[]; variations: string[] } {
const errors: string[] = [];
const variations: string[] = [];
let matches = 0;
for (let i = 0; i < expected.length; i++) {
const exp = expected[i];
const ext = extracted[i];
if (!ext) {
errors.push(`Missing transaction ${i}: ${exp.date} ${exp.counterparty}`);
continue;
}
const dateMatch = ext.date === exp.date;
const amountMatch = Math.abs(ext.amount - exp.amount) < 0.01;
if (dateMatch && amountMatch) {
matches++;
if (ext.counterparty !== exp.counterparty) {
variations.push(`[${i}] "${exp.counterparty}" -> "${ext.counterparty}"`);
}
} else {
errors.push(`Mismatch at ${i}: expected ${exp.date}/${exp.amount}, got ${ext.date}/${ext.amount}`);
}
}
if (extracted.length > expected.length) {
errors.push(`Extra transactions: ${extracted.length - expected.length}`);
}
return { matches, total: expected.length, errors, variations };
}
/**
* Find all test cases (PDF + JSON pairs) in .nogit/
*/
function findTestCases(): Array<{ name: string; pdfPath: string; jsonPath: string }> {
const testDir = path.join(process.cwd(), '.nogit');
if (!fs.existsSync(testDir)) {
return [];
}
const files = fs.readdirSync(testDir);
const pdfFiles = files.filter((f: string) => f.endsWith('.pdf'));
const testCases: Array<{ name: string; pdfPath: string; jsonPath: string }> = [];
for (const pdf of pdfFiles) {
const baseName = pdf.replace('.pdf', '');
const jsonFile = `${baseName}.json`;
if (files.includes(jsonFile)) {
testCases.push({
name: baseName,
pdfPath: path.join(testDir, pdf),
jsonPath: path.join(testDir, jsonFile),
});
}
}
return testCases.sort((a, b) => a.name.localeCompare(b.name));
}
// Tests
tap.test('setup: ensure containers are running', async () => {
console.log('\n[Setup] Checking Docker containers...\n');
// Nanonets for OCR
const nanonetsOk = await ensureNanonetsOcr();
expect(nanonetsOk).toBeTrue();
// Ollama for Qwen3
const ollamaOk = await ensureMiniCpm();
expect(ollamaOk).toBeTrue();
// Qwen3 model
const qwenOk = await ensureQwen3();
expect(qwenOk).toBeTrue();
console.log('\n[Setup] All containers ready!\n');
});
tap.test('should have models available', async () => {
// Check Nanonets
const nanonetsResp = await fetch(`${NANONETS_URL}/models`);
expect(nanonetsResp.ok).toBeTrue();
// Check Qwen3
const ollamaResp = await fetch(`${OLLAMA_URL}/api/tags`);
expect(ollamaResp.ok).toBeTrue();
const data = await ollamaResp.json();
const modelNames = data.models.map((m: { name: string }) => m.name);
expect(modelNames.some((name: string) => name.includes('qwen3'))).toBeTrue();
});
const testCases = findTestCases();
console.log(`\nFound ${testCases.length} bank statement test cases (Nanonets + Qwen3)\n`);
let passedCount = 0;
let failedCount = 0;
for (const testCase of testCases) {
tap.test(`should extract: ${testCase.name}`, async () => {
const expected: ITransaction[] = JSON.parse(fs.readFileSync(testCase.jsonPath, 'utf-8'));
console.log(`\n=== ${testCase.name} ===`);
console.log(`Expected: ${expected.length} transactions`);
const images = convertPdfToImages(testCase.pdfPath);
console.log(` Pages: ${images.length}`);
const extracted = await extractTransactions(images);
console.log(` Extracted: ${extracted.length} transactions`);
const result = compareTransactions(extracted, expected);
const perfectMatch = result.matches === result.total && extracted.length === expected.length;
if (perfectMatch) {
passedCount++;
console.log(` Result: PASS (${result.matches}/${result.total})`);
} else {
failedCount++;
console.log(` Result: FAIL (${result.matches}/${result.total})`);
result.errors.slice(0, 10).forEach((e) => console.log(` - ${e}`));
}
if (result.variations.length > 0) {
console.log(` Counterparty variations (${result.variations.length}):`);
result.variations.slice(0, 5).forEach((v) => console.log(` ${v}`));
}
expect(result.matches).toEqual(result.total);
expect(extracted.length).toEqual(expected.length);
});
}
tap.test('summary', async () => {
const total = testCases.length;
console.log(`\n======================================================`);
console.log(` Bank Statement Summary (Nanonets + Qwen3 Pipeline)`);
console.log(`======================================================`);
console.log(` Stage 1: Nanonets-OCR-s (document -> markdown)`);
console.log(` Stage 2: Qwen3 8B (markdown -> JSON)`);
console.log(` Passed: ${passedCount}/${total}`);
console.log(` Failed: ${failedCount}/${total}`);
console.log(`======================================================\n`);
});
export default tap.start();