feat(tests): integrate SmartAi/DualAgentOrchestrator into extraction tests and add JSON self-validation
This commit is contained in:
@@ -11,7 +11,9 @@ import * as fs from 'fs';
|
||||
import * as path from 'path';
|
||||
import { execSync } from 'child_process';
|
||||
import * as os from 'os';
|
||||
import { ensureNanonetsOcr, ensureMiniCpm, removeContainer, isContainerRunning } from './helpers/docker.js';
|
||||
import { ensureNanonetsOcr, ensureMiniCpm, isContainerRunning } from './helpers/docker.js';
|
||||
import { SmartAi } from '@push.rocks/smartai';
|
||||
import { DualAgentOrchestrator, JsonValidatorTool } from '@push.rocks/smartagent';
|
||||
|
||||
const NANONETS_URL = 'http://localhost:8000/v1';
|
||||
const NANONETS_MODEL = 'nanonets/Nanonets-OCR2-3B';
|
||||
@@ -22,6 +24,22 @@ const EXTRACTION_MODEL = 'gpt-oss:20b';
|
||||
// Temp directory for storing markdown between stages
|
||||
const TEMP_MD_DIR = path.join(os.tmpdir(), 'nanonets-markdown');
|
||||
|
||||
// SmartAi instance for Ollama with optimized settings
|
||||
const smartAi = new SmartAi({
|
||||
ollama: {
|
||||
baseUrl: OLLAMA_URL,
|
||||
model: EXTRACTION_MODEL,
|
||||
defaultOptions: {
|
||||
num_ctx: 32768, // Larger context for long statements + thinking
|
||||
temperature: 0, // Deterministic for JSON extraction
|
||||
},
|
||||
defaultTimeout: 600000, // 10 minute timeout for large documents
|
||||
},
|
||||
});
|
||||
|
||||
// DualAgentOrchestrator for structured task execution
|
||||
let orchestrator: DualAgentOrchestrator;
|
||||
|
||||
interface ITransaction {
|
||||
date: string;
|
||||
counterparty: string;
|
||||
@@ -252,95 +270,107 @@ async function ensureExtractionModel(): Promise<boolean> {
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract transactions from markdown using GPT-OSS 20B (streaming)
|
||||
* Try to extract valid JSON from a response string
|
||||
*/
|
||||
function tryExtractJson(response: string): unknown[] | null {
|
||||
// Remove thinking tags
|
||||
let clean = response.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
|
||||
|
||||
// Try task_complete tags first
|
||||
const completeMatch = clean.match(/<task_complete>([\s\S]*?)<\/task_complete>/);
|
||||
if (completeMatch) {
|
||||
clean = completeMatch[1].trim();
|
||||
}
|
||||
|
||||
// Try code block
|
||||
const codeBlockMatch = clean.match(/```(?:json)?\s*([\s\S]*?)```/);
|
||||
const jsonStr = codeBlockMatch ? codeBlockMatch[1].trim() : clean;
|
||||
|
||||
try {
|
||||
const parsed = JSON.parse(jsonStr);
|
||||
if (Array.isArray(parsed)) return parsed;
|
||||
} catch {
|
||||
// Try to find JSON array
|
||||
const jsonMatch = jsonStr.match(/\[[\s\S]*\]/);
|
||||
if (jsonMatch) {
|
||||
try {
|
||||
const parsed = JSON.parse(sanitizeJson(jsonMatch[0]));
|
||||
if (Array.isArray(parsed)) return parsed;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract transactions from markdown using smartagent DualAgentOrchestrator
|
||||
* Validates JSON and retries if invalid
|
||||
*/
|
||||
async function extractTransactionsFromMarkdown(markdown: string, queryId: string): Promise<ITransaction[]> {
|
||||
const startTime = Date.now();
|
||||
|
||||
console.log(` [${queryId}] Statement: ${markdown.length} chars, Prompt: ${JSON_EXTRACTION_PROMPT.length} chars`);
|
||||
console.log(` [${queryId}] Statement: ${markdown.length} chars`);
|
||||
|
||||
const response = await fetch(`${OLLAMA_URL}/api/chat`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
model: EXTRACTION_MODEL,
|
||||
messages: [
|
||||
{ role: 'user', content: 'Hi there, how are you?' },
|
||||
{ role: 'assistant', content: 'Good, how can I help you today?' },
|
||||
{ role: 'user', content: `Here is a bank statement document:\n\n${markdown}` },
|
||||
{ role: 'assistant', content: 'I have read the bank statement document you provided. I can see all the transaction data. What would you like me to do with it?' },
|
||||
{ role: 'user', content: JSON_EXTRACTION_PROMPT },
|
||||
],
|
||||
stream: true,
|
||||
options: {
|
||||
num_ctx: 32768, // Larger context for long statements + thinking
|
||||
temperature: 0, // Deterministic for JSON extraction
|
||||
},
|
||||
}),
|
||||
signal: AbortSignal.timeout(600000), // 10 minute timeout
|
||||
});
|
||||
// Build the extraction task with document context
|
||||
const taskPrompt = `Extract all transactions from this bank statement document and output ONLY the JSON array:
|
||||
|
||||
if (!response.ok) {
|
||||
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
|
||||
console.log(` [${queryId}] ERROR: ${response.status} (${elapsed}s)`);
|
||||
throw new Error(`Ollama API error: ${response.status}`);
|
||||
}
|
||||
${markdown}
|
||||
|
||||
// Stream the response
|
||||
let content = '';
|
||||
let thinkingContent = '';
|
||||
let thinkingStarted = false;
|
||||
let outputStarted = false;
|
||||
const reader = response.body!.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
${JSON_EXTRACTION_PROMPT}
|
||||
|
||||
Before completing, validate your JSON using the json.validate tool:
|
||||
|
||||
<tool_call>
|
||||
<tool>json</tool>
|
||||
<action>validate</action>
|
||||
<params>{"jsonString": "YOUR_JSON_ARRAY_HERE"}</params>
|
||||
</tool_call>
|
||||
|
||||
Only complete after validation passes. Output the final JSON array in <task_complete></task_complete> tags.`;
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
const result = await orchestrator.run(taskPrompt);
|
||||
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
|
||||
console.log(` [${queryId}] Status: ${result.status}, Iterations: ${result.iterations} (${elapsed}s)`);
|
||||
|
||||
const chunk = decoder.decode(value, { stream: true });
|
||||
// Try to parse JSON from result
|
||||
let jsonData: unknown[] | null = null;
|
||||
let responseText = result.result || '';
|
||||
|
||||
// Each line is a JSON object
|
||||
for (const line of chunk.split('\n').filter(l => l.trim())) {
|
||||
try {
|
||||
const json = JSON.parse(line);
|
||||
if (result.success && responseText) {
|
||||
jsonData = tryExtractJson(responseText);
|
||||
}
|
||||
|
||||
// Stream thinking tokens
|
||||
const thinking = json.message?.thinking || '';
|
||||
if (thinking) {
|
||||
if (!thinkingStarted) {
|
||||
process.stdout.write(` [${queryId}] THINKING: `);
|
||||
thinkingStarted = true;
|
||||
}
|
||||
process.stdout.write(thinking);
|
||||
thinkingContent += thinking;
|
||||
}
|
||||
|
||||
// Stream content tokens
|
||||
const token = json.message?.content || '';
|
||||
if (token) {
|
||||
if (!outputStarted) {
|
||||
if (thinkingStarted) process.stdout.write('\n');
|
||||
process.stdout.write(` [${queryId}] OUTPUT: `);
|
||||
outputStarted = true;
|
||||
}
|
||||
process.stdout.write(token);
|
||||
content += token;
|
||||
}
|
||||
} catch {
|
||||
// Ignore parse errors for partial chunks
|
||||
}
|
||||
// Fallback: try parsing from history
|
||||
if (!jsonData && result.history?.length > 0) {
|
||||
const lastMessage = result.history[result.history.length - 1];
|
||||
if (lastMessage?.content) {
|
||||
responseText = lastMessage.content;
|
||||
jsonData = tryExtractJson(responseText);
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (thinkingStarted || outputStarted) process.stdout.write('\n');
|
||||
|
||||
if (!jsonData) {
|
||||
console.log(` [${queryId}] Failed to parse JSON`);
|
||||
return [];
|
||||
}
|
||||
|
||||
// Convert to transactions
|
||||
const txs = jsonData.map((tx: any) => ({
|
||||
date: String(tx.date || ''),
|
||||
counterparty: String(tx.counterparty || tx.description || ''),
|
||||
amount: parseAmount(tx.amount),
|
||||
}));
|
||||
console.log(` [${queryId}] Parsed ${txs.length} transactions`);
|
||||
return txs;
|
||||
} catch (error) {
|
||||
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
|
||||
console.log(` [${queryId}] ERROR: ${error} (${elapsed}s)`);
|
||||
throw error;
|
||||
}
|
||||
|
||||
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
|
||||
console.log(` [${queryId}] Done: ${thinkingContent.length} thinking chars, ${content.length} output chars (${elapsed}s)`);
|
||||
|
||||
return parseJsonResponse(content, queryId);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -591,6 +621,53 @@ tap.test('Stage 2: Setup Ollama + GPT-OSS 20B', async () => {
|
||||
|
||||
const extractionOk = await ensureExtractionModel();
|
||||
expect(extractionOk).toBeTrue();
|
||||
|
||||
// Initialize SmartAi and DualAgentOrchestrator
|
||||
console.log(' [SmartAgent] Starting SmartAi...');
|
||||
await smartAi.start();
|
||||
|
||||
console.log(' [SmartAgent] Creating DualAgentOrchestrator...');
|
||||
orchestrator = new DualAgentOrchestrator({
|
||||
smartAiInstance: smartAi,
|
||||
defaultProvider: 'ollama',
|
||||
guardianPolicyPrompt: `
|
||||
JSON EXTRACTION POLICY:
|
||||
- APPROVE all JSON extraction tasks
|
||||
- APPROVE all json.validate tool calls
|
||||
- This is a read-only operation - no file system or network access needed
|
||||
- The task is to extract structured transaction data from document text
|
||||
`,
|
||||
driverSystemMessage: `You are a precise JSON extraction assistant. Your only job is to extract transaction data from bank statements.
|
||||
|
||||
CRITICAL RULES:
|
||||
1. Output valid JSON array with the exact format requested
|
||||
2. Amounts should be NEGATIVE for debits/withdrawals, POSITIVE for credits/deposits
|
||||
3. IMPORTANT: Before completing, validate your JSON using the json.validate tool:
|
||||
|
||||
<tool_call>
|
||||
<tool>json</tool>
|
||||
<action>validate</action>
|
||||
<params>{"jsonString": "YOUR_JSON_ARRAY"}</params>
|
||||
</tool_call>
|
||||
|
||||
4. Only complete after validation passes
|
||||
|
||||
When done, wrap your JSON array in <task_complete></task_complete> tags.`,
|
||||
maxIterations: 5,
|
||||
// Enable streaming for real-time progress visibility
|
||||
onToken: (token, source) => {
|
||||
if (source === 'driver') {
|
||||
process.stdout.write(token);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
// Register JsonValidatorTool for self-validation
|
||||
orchestrator.registerTool(new JsonValidatorTool());
|
||||
|
||||
console.log(' [SmartAgent] Starting orchestrator...');
|
||||
await orchestrator.start();
|
||||
console.log(' [SmartAgent] Ready for extraction');
|
||||
});
|
||||
|
||||
let passedCount = 0;
|
||||
@@ -642,11 +719,19 @@ for (const tc of testCases) {
|
||||
}
|
||||
|
||||
tap.test('Summary', async () => {
|
||||
// Cleanup orchestrator and SmartAi
|
||||
if (orchestrator) {
|
||||
console.log('\n [SmartAgent] Stopping orchestrator...');
|
||||
await orchestrator.stop();
|
||||
}
|
||||
console.log(' [SmartAgent] Stopping SmartAi...');
|
||||
await smartAi.stop();
|
||||
|
||||
console.log(`\n======================================================`);
|
||||
console.log(` Bank Statement Summary (Nanonets + GPT-OSS 20B Sequential)`);
|
||||
console.log(` Bank Statement Summary (Nanonets + SmartAgent)`);
|
||||
console.log(`======================================================`);
|
||||
console.log(` Stage 1: Nanonets-OCR-s (document -> markdown)`);
|
||||
console.log(` Stage 2: GPT-OSS 20B (markdown -> JSON)`);
|
||||
console.log(` Stage 2: GPT-OSS 20B + SmartAgent (markdown -> JSON)`);
|
||||
console.log(` Passed: ${passedCount}/${testCases.length}`);
|
||||
console.log(` Failed: ${failedCount}/${testCases.length}`);
|
||||
console.log(`======================================================\n`);
|
||||
|
||||
Reference in New Issue
Block a user