56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
"""Transaction data extraction: LLM Vision + persistence."""
|
|
|
|
from app.models import Transaction
|
|
import app.models.database as db_module
|
|
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,
|
|
progress_hook=None,
|
|
) -> 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.
|
|
"""
|
|
if progress_hook:
|
|
await progress_hook("init", 5, "初始化识别上下文")
|
|
provider = get_llm_provider()
|
|
if progress_hook:
|
|
await progress_hook("provider_ready", 15, f"已加载模型提供商: {type(provider).__name__}")
|
|
if progress_hook:
|
|
await progress_hook("calling_model", 35, "调用视觉模型识别截图中交易")
|
|
items: list[TransactionExtractItem] = await provider.extract_from_image(image_bytes)
|
|
if progress_hook:
|
|
await progress_hook("model_returned", 70, f"模型返回 {len(items)} 条交易")
|
|
results: list[TransactionResponse] = []
|
|
async with db_module.async_session_maker() as session:
|
|
if progress_hook:
|
|
await progress_hook("db_writing", 85, "写入交易记录到数据库")
|
|
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()
|
|
if progress_hook:
|
|
await progress_hook("completed", 100, "识别完成")
|
|
return results
|