Files
fund-tracer/backend/app/services/extractor.py

43 lines
1.6 KiB
Python
Raw Normal View History

2026-03-09 14:46:56 +08:00
"""Transaction data extraction: LLM Vision + persistence."""
from app.models import Transaction
from app.models.database import async_session_maker
from app.schemas.transaction import TransactionExtractItem, TransactionResponse
from app.services.llm import get_llm_provider
async def extract_and_save(
case_id: int,
screenshot_id: int,
image_bytes: bytes,
) -> list[TransactionResponse]:
"""
Run vision extraction on image and persist transactions to DB.
Returns list of created transactions; low-confidence items are still saved but flagged.
"""
provider = get_llm_provider()
items: list[TransactionExtractItem] = await provider.extract_from_image(image_bytes)
results: list[TransactionResponse] = []
async with async_session_maker() as session:
for it in items:
t = Transaction(
case_id=case_id,
screenshot_id=screenshot_id,
app_source=it.app_source,
transaction_type=it.transaction_type,
amount=it.amount,
currency=it.currency or "CNY",
counterparty_name=it.counterparty_name,
counterparty_account=it.counterparty_account,
order_number=it.order_number,
transaction_time=it.transaction_time,
remark=it.remark,
confidence=it.confidence if it.confidence in ("high", "medium", "low") else "medium",
raw_text=None,
)
session.add(t)
await session.flush()
results.append(TransactionResponse.model_validate(t))
await session.commit()
return results