216 lines
8.1 KiB
Python
216 lines
8.1 KiB
Python
"""Pipeline blueprint — SSE-сравнение + classify."""
|
|
import json, re, os, threading
|
|
from flask import Blueprint, request, jsonify, Response, stream_with_context
|
|
from services.process import run_pipeline
|
|
from services import extractor as stateless_extractor
|
|
from services.grouping import build_dossiers_for_batch
|
|
from services.timeline_builder import build_dossier_timeline
|
|
from llm_prompt import build_prompt
|
|
from db import documents
|
|
|
|
pipeline_bp = Blueprint("pipeline", __name__)
|
|
|
|
# In-memory lock для classify (замена файлового lock)
|
|
_classify_locks: dict[str, threading.Thread] = {}
|
|
_extraction_jobs: dict[str, dict] = {}
|
|
|
|
|
|
@pipeline_bp.route("/process-v2", methods=["GET"])
|
|
def process_v2():
|
|
"""SSE-стриминг сравнения договоров."""
|
|
dossier_id = request.args.get("dossier_id")
|
|
if dossier_id:
|
|
if not re.fullmatch(
|
|
r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}',
|
|
dossier_id, re.I,
|
|
):
|
|
return jsonify(ok=False, error="invalid dossier_id"), 400
|
|
|
|
def generate_timeline():
|
|
import time as _time
|
|
started = _time.time()
|
|
yield ": ok\n\n"
|
|
yield "data: " + json.dumps({
|
|
"type": "extract_start", "supplement_id": dossier_id,
|
|
"filename": "Построение таймлайна",
|
|
}, ensure_ascii=False) + "\n\n"
|
|
try:
|
|
result = build_dossier_timeline(dossier_id)
|
|
if not result.get("ok"):
|
|
yield "data: " + json.dumps({
|
|
"type": "error", "message": result.get("error", "timeline build failed"),
|
|
}, ensure_ascii=False) + "\n\n"
|
|
return
|
|
yield "data: " + json.dumps({
|
|
"type": "timeline", "dossier_id": dossier_id,
|
|
"timeline": result.get("timeline", []),
|
|
"exceptions": result.get("exceptions", []),
|
|
"counts": result.get("counts", {}),
|
|
}, ensure_ascii=False, default=str) + "\n\n"
|
|
yield "data: " + json.dumps({
|
|
"type": "done", "dossier_id": dossier_id,
|
|
"total_time_s": round(_time.time() - started, 2),
|
|
}, ensure_ascii=False) + "\n\n"
|
|
except GeneratorExit:
|
|
return
|
|
except Exception as error:
|
|
yield "data: " + json.dumps({
|
|
"type": "error", "message": str(error),
|
|
}, ensure_ascii=False) + "\n\n"
|
|
|
|
return Response(
|
|
stream_with_context(generate_timeline()),
|
|
content_type="text/event-stream; charset=utf-8",
|
|
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
|
)
|
|
|
|
cid = request.args.get("contract_id")
|
|
if not cid or not re.fullmatch(
|
|
r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}', cid, re.I
|
|
):
|
|
return jsonify(ok=False, error="invalid contract_id"), 400
|
|
|
|
order_ids = request.args.get("order", "")
|
|
|
|
def generate():
|
|
# Heartbeat каждые 15 сек — держит соединение при медленном LLM
|
|
import time as _time
|
|
last_beat = _time.time()
|
|
|
|
yield ": ok\n\n"
|
|
try:
|
|
for event in run_pipeline(cid, order_ids, build_prompt):
|
|
now = _time.time()
|
|
if now - last_beat >= 15:
|
|
yield ": heartbeat\n\n"
|
|
last_beat = now
|
|
yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
|
|
except GeneratorExit:
|
|
return
|
|
except Exception as e:
|
|
yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"
|
|
|
|
return Response(
|
|
stream_with_context(generate()),
|
|
content_type="text/event-stream; charset=utf-8",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
|
|
@pipeline_bp.route("/api/classify-batch", methods=["POST"])
|
|
def classify_batch_route():
|
|
"""Выполнить stateless extraction и построить contract dossiers."""
|
|
body = request.get_json(silent=True) or {}
|
|
if not isinstance(body, dict):
|
|
return jsonify(ok=False, error="JSON object required"), 400
|
|
batch_id = body.get("batch_id")
|
|
if not isinstance(batch_id, str) or not batch_id.strip():
|
|
return jsonify(ok=False, error="batch_id required"), 400
|
|
|
|
# Guard: уже запущена?
|
|
if batch_id in _classify_locks:
|
|
t = _classify_locks[batch_id]
|
|
if t.is_alive():
|
|
return jsonify(ok=False, error="classify already running"), 409
|
|
else:
|
|
del _classify_locks[batch_id]
|
|
|
|
pending = documents.list_for_extraction(batch_id)
|
|
total = len(pending)
|
|
if total == 0:
|
|
return jsonify(ok=False, error="no parsed documents"), 400
|
|
|
|
# Keep small batches synchronous for the existing browser workflow.
|
|
if total <= 10:
|
|
result = _run_v3_extraction(batch_id, total)
|
|
return jsonify(result)
|
|
|
|
_extraction_jobs[batch_id] = {
|
|
"state": "running", "total": total, "done": 0,
|
|
"failed": 0, "classes": {},
|
|
}
|
|
|
|
def _run():
|
|
try:
|
|
_run_v3_extraction(batch_id, total, _extraction_progress)
|
|
_extraction_jobs[batch_id]["state"] = "done"
|
|
except Exception as error:
|
|
_extraction_jobs[batch_id].update(state="failed", error=str(error))
|
|
finally:
|
|
_classify_locks.pop(batch_id, None)
|
|
|
|
t = threading.Thread(target=_run, daemon=True)
|
|
_classify_locks[batch_id] = t
|
|
t.start()
|
|
|
|
return jsonify(ok=True, total=total, v3=True), 202
|
|
|
|
|
|
def _extraction_progress(batch_id, completed, result):
|
|
"""Обновить лёгкий in-memory прогресс для крупного batch."""
|
|
job = _extraction_jobs[batch_id]
|
|
job["done"] = completed
|
|
if not result.get("ok"):
|
|
job["failed"] += 1
|
|
else:
|
|
extraction = result.get("extraction") or {}
|
|
doc_class = extraction.get("doc_class", "unknown")
|
|
job["classes"][doc_class] = job["classes"].get(doc_class, 0) + 1
|
|
|
|
|
|
def _run_v3_extraction(batch_id, total, progress_callback=None):
|
|
"""Собрать extraction и dossier summaries в один ответ API."""
|
|
extraction = stateless_extractor.extract_batch(
|
|
batch_id,
|
|
progress_callback=(
|
|
(lambda completed, result: progress_callback(batch_id, completed, result))
|
|
if progress_callback else None
|
|
),
|
|
)
|
|
dossier_result = build_dossiers_for_batch(batch_id)
|
|
classes = extraction.get("classes", {})
|
|
math_checks_passed = sum(
|
|
1 for result in extraction.get("results", [])
|
|
if (result.get("extraction") or {}).get("math_validation", {}).get("status") == "VERIFIED"
|
|
)
|
|
return {
|
|
"ok": bool(extraction.get("results")) and dossier_result.get("ok", False),
|
|
"v3": True,
|
|
"done": total,
|
|
"total": total,
|
|
"failed": extraction.get("failed", 0),
|
|
"classes": classes,
|
|
"ignored_accounting": classes.get("ignored_accounting", 0),
|
|
"math_checks_passed": math_checks_passed,
|
|
"dossier_count": len(dossier_result.get("dossiers", [])),
|
|
"ambiguous": dossier_result.get("ambiguous", 0),
|
|
"unlinked": dossier_result.get("unlinked", []),
|
|
}
|
|
|
|
|
|
@pipeline_bp.route("/api/classify-progress")
|
|
def classify_progress_route():
|
|
"""Вернуть состояние асинхронного stateless extraction batch."""
|
|
batch_id = request.args.get("batch")
|
|
if not batch_id:
|
|
return jsonify(ok=False, error="batch required"), 400
|
|
job = _extraction_jobs.get(batch_id)
|
|
if not job:
|
|
return jsonify(ok=False, error="batch extraction not found"), 404
|
|
response = {
|
|
"ok": job["state"] != "failed",
|
|
"v3": True,
|
|
"finished": job["state"] in {"done", "failed"},
|
|
"state": job["state"],
|
|
"done": job["done"],
|
|
"total": job["total"],
|
|
"failed": job["failed"],
|
|
"classes": job["classes"],
|
|
}
|
|
if job.get("error"):
|
|
response["error"] = job["error"]
|
|
return jsonify(response)
|