141 lines
5.0 KiB
Python
141 lines
5.0 KiB
Python
"""Process service — SSE pipeline: reset → supplements → LLM → apply."""
|
||
import json, time, threading
|
||
from db import supplements, spec_current, spec_events
|
||
from .metrics import check_arithmetic
|
||
|
||
|
||
def run_pipeline(contract_id, order_ids, sse_send, build_prompt_fn):
|
||
"""Run full SSE comparison pipeline. sse_send is a callback(dict)."""
|
||
t0 = time.time()
|
||
|
||
# 0. Reset
|
||
spec_events.reset(contract_id)
|
||
|
||
# 1. Get supplements with parsed documents
|
||
supps = supplements.list_by_contract(contract_id)
|
||
if order_ids:
|
||
order_list = [x.strip() for x in order_ids.split(",") if x.strip()]
|
||
order_map = {oid: i for i, oid in enumerate(order_list)}
|
||
supps.sort(key=lambda s: order_map.get(s["id"], 999999))
|
||
else:
|
||
# Sort by doc_date (from classification), fallback to created_at
|
||
supps.sort(key=lambda s: (s.get("doc_date") or "9999-99-99", s.get("created_at", "")))
|
||
|
||
if not supps:
|
||
sse_send({"type": "error", "message": "Нет распарсенных файлов"})
|
||
return
|
||
|
||
from .llm import call_llm
|
||
|
||
for s in supps:
|
||
sid = s["id"]
|
||
|
||
# Current spec
|
||
cur = spec_current.list_by_contract(contract_id)
|
||
current_spec = []
|
||
for r in cur:
|
||
current_spec.append({
|
||
"hash": r["name_hash"],
|
||
"name": r["name"],
|
||
"price": float(r["price"]) if r.get("price") is not None else None,
|
||
"qty": float(r["qty"]) if r.get("qty") is not None else None,
|
||
"sum": float(r["sum"]) if r.get("sum") is not None else None,
|
||
"date_start": r["date_start"],
|
||
})
|
||
|
||
# Get elements_json
|
||
ej = spec_current.get_elements_json(s["document_id"])
|
||
if not ej:
|
||
continue
|
||
|
||
if isinstance(ej, dict) and "Value" in ej:
|
||
ej = ej["Value"]
|
||
if isinstance(ej, str):
|
||
elements = json.loads(ej)
|
||
elif isinstance(ej, list):
|
||
elements = ej
|
||
else:
|
||
elements = []
|
||
|
||
elements = [{k.lower(): v for k, v in el.items()} for el in elements]
|
||
doc_text = _elements_to_text(elements)
|
||
|
||
sse_send({"type": "extract_start", "supplement_id": sid, "filename": s["filename"]})
|
||
|
||
# LLM call in thread with keepalive
|
||
t1 = time.time()
|
||
done = threading.Event()
|
||
result = [None]
|
||
error = [None]
|
||
|
||
def do_llm():
|
||
try:
|
||
result[0] = call_llm(current_spec, doc_text, build_prompt_fn)
|
||
except Exception as e:
|
||
error[0] = str(e)
|
||
finally:
|
||
done.set()
|
||
|
||
t = threading.Thread(target=do_llm)
|
||
t.start()
|
||
|
||
while not done.wait(15):
|
||
sse_send({"type": "keepalive"})
|
||
|
||
t.join()
|
||
elapsed = round(time.time() - t1, 1)
|
||
|
||
if error[0]:
|
||
sse_send({"type": "extract_error", "supplement_id": sid,
|
||
"filename": s["filename"], "error": error[0], "time_s": elapsed})
|
||
continue
|
||
|
||
llm_result, prompt_id = result[0]
|
||
ops = llm_result.get("ops", [])
|
||
mode = llm_result.get("mode", "partial")
|
||
|
||
# Enrich ops
|
||
id_map = {f"r{i+1}": r["hash"] for i, r in enumerate(current_spec)}
|
||
spec_names = {r["hash"]: r["name"] for r in current_spec}
|
||
for op in ops:
|
||
tid = op.get("target_id")
|
||
if tid and tid in id_map:
|
||
op["target_hash"] = id_map[tid]
|
||
th = op.get("target_hash")
|
||
if th and th in spec_names and not op.get("new_row", {}).get("name"):
|
||
op.setdefault("new_row", {})["name"] = spec_names[th]
|
||
|
||
sse_send({"type": "llm_done", "supplement_id": sid, "filename": s["filename"],
|
||
"ops_count": len(ops), "mode": mode, "time_s": elapsed})
|
||
|
||
# Apply events
|
||
summary = spec_events.apply_ops(
|
||
contract_id, sid, s.get("document_id", ""), ops, prompt_id, llm_result
|
||
)
|
||
# Arithmetic quality check (free signal, no golden dataset needed)
|
||
arith_mismatches = check_arithmetic(ops)
|
||
if arith_mismatches:
|
||
summary["arithmetic_mismatches"] = len(arith_mismatches)
|
||
sse_send({"type": "applied", "supplement_id": sid, "summary": summary, "ops": ops})
|
||
|
||
total_time = round(time.time() - t0, 1)
|
||
sse_send({"type": "done", "total_time_s": total_time})
|
||
|
||
|
||
def _elements_to_text(elements):
|
||
lines = []
|
||
for el in elements:
|
||
if el.get("type") == "paragraph":
|
||
prefix = f"[{el['style']}] " if el.get("style") else ""
|
||
lines.append(prefix + el.get("text", ""))
|
||
elif el.get("type") == "table":
|
||
rows = el.get("rows", [])
|
||
if rows:
|
||
ncols = len(rows[0])
|
||
lines.append(f"--- Таблица ({len(rows)}×{ncols}) ---")
|
||
for row in rows:
|
||
cells = [str(c).replace("\n", " ").replace("|", "\\|") for c in row[:ncols]]
|
||
lines.append("| " + " | ".join(cells) + " |")
|
||
lines.append("")
|
||
return "\n".join(lines)
|