regression_test_ensemble.py - main - Paff/declarant

regression_test_ensemble.py
274 строки11 KB
1pash1985-gif
feat: AI Ensemble - DeepSeek + Kimi parallel parsing with debates
27 июл 2026, 11:14
27 июл 2026, 11:145ea7945
"""Regression test: compare legacy single-model parsing vs ensemble parsing.

Usage:
    python regression_test_ensemble.py <trip_folder> [--sample N]

Example:
    python regression_test_ensemble.py "c:\\Users\\0\\Desktop\\Задача по логистике\\26.07\\26.07\\BOJ444" --sample 3
"""
from __future__ import annotations

import argparse
import io
import json
import os
import sys
import tempfile
import time
import zipfile
from pathlib import Path

# Speed up regression by limiting debate rounds and extending timeout.
# These defaults can be overridden via environment variables.
os.environ.setdefault('ENSEMBLE_DEBATE_MAX_ROUNDS', '1')
os.environ.setdefault('ENSEMBLE_TIMEOUT_SECONDS', '300')

# Ensure project root and web/ are importable
PROJECT_ROOT = Path(__file__).resolve().parent
WEB_ROOT = PROJECT_ROOT / 'web'
for p in (str(PROJECT_ROOT), str(WEB_ROOT)):
    if p not in sys.path:
        sys.path.insert(0, p)

from ai_parser import (
    extract_text_from_pdf, extract_text_from_excel,
    parse_invoice_with_ai, parse_packing_with_ai, parse_certificate_with_kimi,
)
from ai_ensemble import (
    parse_invoice_with_ensemble, parse_packing_with_ensemble,
    parse_certificate_with_ensemble,
)


def collect_source_files(trip: Path, sample: int | None = None) -> list[Path]:
    """Collect loose + archived PDF/Excel source files, skipping output templates."""
    skip_suffixes = {
        'таблица европа', 'таблица импорта', 'florunner-specification',
        '24.07.', '21.07.', '22.07.', '23.07.', '24-',
    }

    loose_files: list[Path] = []
    for ext in ('*.pdf', '*.xls', '*.xlsx', '*.xlsm'):
        for f in trip.rglob(ext):
            name_lower = f.name.lower()
            if any(s in name_lower for s in skip_suffixes):
                continue
            loose_files.append(f)

    # Extract a few files from RAR/ZIP archives
    archived_files: list[Path] = []
    for archive in trip.rglob('*'):
        if archive.suffix.lower() not in ('.rar', '.zip'):
            continue
        try:
            if zipfile.is_zipfile(str(archive)):
                zf = zipfile.ZipFile(str(archive))
                members = [m for m in zf.namelist() if m.lower().endswith(('.pdf', '.xls', '.xlsx', '.xlsm'))]
                selected = members[:2]  # up to 2 per archive
                tmp_dir = Path(tempfile.mkdtemp(prefix='regression_'))
                for m in selected:
                    try:
                        zf.extract(m, tmp_dir)
                        archived_files.append(tmp_dir / m)
                    except Exception:
                        pass
            else:
                import rarfile
                rf = rarfile.RarFile(str(archive))
                members = [m for m in rf.namelist() if m.lower().endswith(('.pdf', '.xls', '.xlsx', '.xlsm'))]
                selected = members[:2]
                tmp_dir = Path(tempfile.mkdtemp(prefix='regression_'))
                for m in selected:
                    try:
                        rf.extract(m, tmp_dir)
                        archived_files.append(tmp_dir / m)
                    except Exception:
                        pass
        except Exception as e:
            print(f'  [!] Cannot read archive {archive.name}: {e}')

    all_files = loose_files + archived_files
    # Prefer PDF invoices first
    all_files.sort(key=lambda f: (0 if f.suffix.lower() == '.pdf' else 1, str(f)))
    if sample:
        all_files = all_files[:sample]
    return all_files


def guess_doc_type(path: Path) -> str:
    name_lower = path.name.lower()
    if 'sertif' in name_lower or 'certif' in name_lower or 'coo' in name_lower or 'eur1' in name_lower:
        return 'certificate'
    # Excel files in subfolders are usually packing lists or invoices
    if path.suffix.lower() in ('.xls', '.xlsx', '.xlsm'):
        # If filename contains packing/упак - packing, else invoice
        if 'pack' in name_lower or 'упак' in name_lower:
            return 'packing'
        return 'invoice'
    # Default PDF to invoice
    return 'invoice'


def parse_legacy(text: str, doc_type: str) -> dict:
    if doc_type == 'certificate':
        return parse_certificate_with_kimi(text)
    if doc_type == 'packing':
        return parse_packing_with_ai(text)
    return parse_invoice_with_ai(text)


def parse_ensemble(text: str, doc_type: str) -> dict:
    if doc_type == 'certificate':
        return parse_certificate_with_ensemble(text)
    if doc_type == 'packing':
        return parse_packing_with_ensemble(text)
    return parse_invoice_with_ensemble(text)


def extract_text(path: Path) -> str:
    ext = path.suffix.lower()
    if ext == '.pdf':
        return extract_text_from_pdf(str(path))
    if ext in ('.xls', '.xlsx', '.xlsm'):
        return extract_text_from_excel(str(path))
    return ''


def compare_results(legacy: dict, ensemble: dict) -> dict:
    """Compare key fields between legacy and ensemble results."""
    diffs: list[dict] = []

    def check(label: str, a, b):
        # Normalize numbers
        if isinstance(a, (int, float)) and isinstance(b, (int, float)):
            if float(a) != float(b):
                diffs.append({'field': label, 'legacy': a, 'ensemble': b})
        elif a != b:
            diffs.append({'field': label, 'legacy': a, 'ensemble': b})

    # Top-level fields
    check('supplier', legacy.get('supplier'), ensemble.get('supplier'))
    check('marking', legacy.get('marking'), ensemble.get('marking'))
    check('country', legacy.get('country'), ensemble.get('country'))
    check('invoice_number', legacy.get('invoice_number'), ensemble.get('invoice_number'))
    check('total_stems', legacy.get('total_stems'), ensemble.get('total_stems'))
    check('total_boxes', legacy.get('total_boxes'), ensemble.get('total_boxes'))
    check('items_count', len(legacy.get('items') or []), len(ensemble.get('items') or []))

    # Item-level comparison (by index)
    legacy_items = legacy.get('items') or []
    ensemble_items = ensemble.get('items') or []
    item_diffs: list[dict] = []

    max_len = max(len(legacy_items), len(ensemble_items))
    for idx in range(max_len):
        if idx >= len(legacy_items):
            item_diffs.append({'index': idx, 'field': 'item_missing', 'legacy': None, 'ensemble': ensemble_items[idx]})
        elif idx >= len(ensemble_items):
            item_diffs.append({'index': idx, 'field': 'item_missing', 'legacy': legacy_items[idx], 'ensemble': None})
        else:
            li = legacy_items[idx]
            ei = ensemble_items[idx]
            # Compare key item fields
            for field in ('marking', 'boxes', 'stems', 'price', 'flower_type', 'color'):
                lv = li.get(field) if isinstance(li, dict) else None
                ev = ei.get(field) if isinstance(ei, dict) else None
                if isinstance(lv, (int, float)) and isinstance(ev, (int, float)):
                    if float(lv) != float(ev):
                        item_diffs.append({'index': idx, 'field': field, 'legacy': lv, 'ensemble': ev})
                elif lv != ev:
                    item_diffs.append({'index': idx, 'field': field, 'legacy': lv, 'ensemble': ev})

    ensemble_meta = ensemble.get('_ensemble_meta') or {}
    return {
        'legacy_error': legacy.get('error'),
        'ensemble_error': ensemble.get('error'),
        'ensemble_status': ensemble_meta.get('status'),
        'debate_rounds': ensemble_meta.get('debate_rounds', 0),
        'disputed_paths': ensemble_meta.get('disputed_paths', []),
        'diffs': diffs,
        'item_diffs': item_diffs,
    }


def main():
    parser = argparse.ArgumentParser(description='Regression test: legacy vs ensemble AI parsing')
    parser.add_argument('trip', help='Path to trip folder')
    parser.add_argument('--sample', type=int, default=5, help='Number of files to test (default 5)')
    args = parser.parse_args()

    trip = Path(args.trip)
    if not trip.is_dir():
        print(f'Folder not found: {trip}')
        sys.exit(1)

    print(f'=== Regression test for {trip.name} ===', flush=True)
    print(f'Collecting up to {args.sample} source files...', flush=True)
    files = collect_source_files(trip, args.sample)
    print(f'Found {len(files)} files to test\n', flush=True)

    results = []
    session_start = time.time()
    for i, f in enumerate(files, 1):
        file_start = time.time()
        rel = f.relative_to(trip) if trip in f.parents else f.name
        print(f'[{i}/{len(files)}] {rel}', flush=True)
        text = extract_text(f)
        if not text.strip():
            print('  [ ] Empty text, skipping', flush=True)
            continue
        doc_type = guess_doc_type(f)
        print(f'  doc_type={doc_type}, text_len={len(text)}', flush=True)

        legacy = parse_legacy(text, doc_type)
        print(f'  legacy done ({time.time()-file_start:.1f}s)', flush=True)
        ensemble = parse_ensemble(text, doc_type)
        print(f'  ensemble done ({time.time()-file_start:.1f}s)', flush=True)
        comparison = compare_results(legacy, ensemble)

        if comparison['legacy_error'] or comparison['ensemble_error']:
            print(f'  [X] legacy_error={comparison["legacy_error"]}, ensemble_error={comparison["ensemble_error"]}', flush=True)
        else:
            n_top = len(comparison['diffs'])
            n_item = len(comparison.get('item_diffs', []))
            if n_top or n_item:
                print(f'  [~] DIFFS: {n_top} top-level, {n_item} item-level', flush=True)
                for d in comparison['diffs']:
                    print(f'      {d["field"]}: legacy={d["legacy"]} → ensemble={d["ensemble"]}', flush=True)
                for d in comparison.get('item_diffs', []):
                    print(f'      items[{d["index"]}].{d["field"]}: legacy={d["legacy"]} → ensemble={d["ensemble"]}', flush=True)
            else:
                print(f'  [OK] Match, status={comparison["ensemble_status"]}', flush=True)

        results.append({
            'file': str(rel),
            'doc_type': doc_type,
            **comparison,
        })
        print(f'  [file done in {time.time()-file_start:.1f}s, total {time.time()-session_start:.1f}s]\n', flush=True)

    # Summary
    total = len(results)
    errors = sum(1 for r in results if r['legacy_error'] or r['ensemble_error'])
    diffs = sum(1 for r in results if r['diffs'] or r.get('item_diffs'))
    matches = total - errors - diffs
    total_top_diffs = sum(len(r.get('diffs', [])) for r in results)
    total_item_diffs = sum(len(r.get('item_diffs', [])) for r in results)

    print('=== SUMMARY ===', flush=True)
    print(f'  Total tested:    {total}', flush=True)
    print(f'  Matches:         {matches}', flush=True)
    print(f'  Files with diffs: {diffs}', flush=True)
    print(f'  Errors:          {errors}', flush=True)
    print(f'  Top-level diffs: {total_top_diffs}', flush=True)
    print(f'  Item-level diffs: {total_item_diffs}', flush=True)

    # Save detailed report
    report_path = Path('regression_report.json').resolve()
    with open(report_path, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f'\nDetailed report saved to: {report_path}')


if __name__ == '__main__':
    main()