355 lines
14 KiB
Python
355 lines
14 KiB
Python
"""Pipeline blueprint — SSE-сравнение + classify."""
|
|
import hashlib, hmac, json, re, os, threading, time
|
|
from flask import Blueprint, request, jsonify, Response, stream_with_context
|
|
from config import API_KEY
|
|
from services.process import run_pipeline
|
|
from services import extractor as stateless_extractor
|
|
from services.engine import analyze_batch
|
|
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, extraction_jobs
|
|
from db.connection import transaction
|
|
|
|
pipeline_bp = Blueprint("pipeline", __name__)
|
|
|
|
_engine_rate_lock = threading.Lock()
|
|
_engine_rate_hits: dict[str, list[float]] = {}
|
|
_ENGINE_MAX_BODY_BYTES = 10 * 1024 * 1024
|
|
_ENGINE_MAX_DOCUMENTS = 100
|
|
_ENGINE_MAX_ELEMENTS = 20_000
|
|
_ENGINE_MAX_DOCUMENT_BYTES = 2 * 1024 * 1024
|
|
_ENGINE_RATE_WINDOW_SECONDS = 60
|
|
_ENGINE_RATE_LIMIT = 10
|
|
|
|
|
|
@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
|
|
|
|
pending = documents.list_for_extraction(batch_id)
|
|
total = len(pending)
|
|
if total == 0:
|
|
return jsonify(ok=False, error="no parsed documents"), 400
|
|
owner = extraction_jobs.claim(batch_id, total)
|
|
if not owner:
|
|
return jsonify(ok=False, error="classify already running"), 409
|
|
|
|
# Keep small batches synchronous for the existing browser workflow.
|
|
if total <= 10:
|
|
result = _run_extraction_job(batch_id, pending, owner)
|
|
return jsonify(result)
|
|
|
|
_start_extraction_job(batch_id, pending, owner)
|
|
return jsonify(ok=True, total=total, v3=True), 202
|
|
|
|
|
|
@pipeline_bp.route("/api/engine/analyze", methods=["POST"])
|
|
def engine_analyze_route():
|
|
"""Accept parsed documents over a versioned, authenticated JSON boundary.
|
|
|
|
Validation is completed for the whole payload before opening the write
|
|
transaction, so malformed input cannot leave a partial batch. The caller
|
|
supplies parsed elements only; extraction passports and dossier decisions
|
|
are always produced by the in-process Engine and persisted by its adapters.
|
|
"""
|
|
if not API_KEY:
|
|
return jsonify(ok=False, error="engine API is not configured"), 503
|
|
supplied_key = request.headers.get("X-API-Key", "")
|
|
if not hmac.compare_digest(supplied_key.encode("utf-8"), API_KEY.encode("utf-8")):
|
|
return jsonify(ok=False, error="unauthorized"), 401
|
|
if _engine_rate_limited():
|
|
return jsonify(ok=False, error="rate limit exceeded"), 429, {"Retry-After": "60"}
|
|
|
|
raw_body = request.stream.read(_ENGINE_MAX_BODY_BYTES + 1)
|
|
if len(raw_body) > _ENGINE_MAX_BODY_BYTES:
|
|
return jsonify(ok=False, error="request body too large"), 413
|
|
try:
|
|
body = json.loads(raw_body) if request.is_json else None
|
|
except (ValueError, UnicodeDecodeError):
|
|
body = None
|
|
if not isinstance(body, dict):
|
|
return jsonify(ok=False, error="JSON object required"), 400
|
|
if type(body.get("schema_version")) is not int or body["schema_version"] != 1:
|
|
return jsonify(ok=False, error="unsupported schema_version"), 400
|
|
|
|
batch_id = body.get("batch_id")
|
|
if not isinstance(batch_id, str) or not re.fullmatch(r"[A-Za-z0-9._:-]{1,128}", batch_id):
|
|
return jsonify(ok=False, error="invalid batch_id"), 400
|
|
source_documents = body.get("documents")
|
|
if not isinstance(source_documents, list) or not source_documents:
|
|
return jsonify(ok=False, error="documents must be a non-empty array"), 400
|
|
if len(source_documents) > _ENGINE_MAX_DOCUMENTS:
|
|
return jsonify(ok=False, error="too many documents"), 413
|
|
validated = []
|
|
for index, source_document in enumerate(source_documents):
|
|
if not isinstance(source_document, dict):
|
|
return jsonify(ok=False, error=f"documents[{index}] must be an object"), 400
|
|
filename = source_document.get("filename")
|
|
elements = source_document.get("elements")
|
|
if not isinstance(filename, str) or not filename.strip() or len(filename) > 512:
|
|
return jsonify(ok=False, error=f"documents[{index}].filename is invalid"), 400
|
|
if not isinstance(elements, list) or len(elements) > _ENGINE_MAX_ELEMENTS:
|
|
return jsonify(ok=False, error=f"documents[{index}].elements is invalid"), 400
|
|
for element_index, element in enumerate(elements):
|
|
if not _valid_engine_element(element):
|
|
return jsonify(
|
|
ok=False, error=f"documents[{index}].elements[{element_index}] is invalid",
|
|
), 400
|
|
try:
|
|
elements_json = json.dumps(
|
|
elements, ensure_ascii=False, sort_keys=True, separators=(",", ":"),
|
|
allow_nan=False,
|
|
)
|
|
except (TypeError, ValueError):
|
|
return jsonify(ok=False, error=f"documents[{index}].elements is not valid JSON"), 400
|
|
encoded_elements = elements_json.encode("utf-8")
|
|
if len(encoded_elements) > _ENGINE_MAX_DOCUMENT_BYTES:
|
|
return jsonify(ok=False, error=f"documents[{index}] is too large"), 413
|
|
validated.append((filename.strip(), elements, elements_json, encoded_elements))
|
|
|
|
inserted = []
|
|
with transaction():
|
|
# Engine API owns a fresh batch per request. Preventing reuse inside
|
|
# BEGIN IMMEDIATE keeps old parsed rows out of this run and closes the
|
|
# race between two concurrent requests using the same batch ID.
|
|
if documents.list_by_batch(batch_id):
|
|
return jsonify(ok=False, error="batch_id already exists"), 409
|
|
owner = extraction_jobs.claim(batch_id, len(validated))
|
|
if not owner:
|
|
return jsonify(ok=False, error="classify already running"), 409
|
|
for filename, elements, elements_json, encoded_elements in validated:
|
|
content_hash = hashlib.sha256(encoded_elements).hexdigest()
|
|
document = documents.insert(
|
|
filename, "application/json", None, status="uploaded",
|
|
batch_id=batch_id, content_hash=content_hash,
|
|
)
|
|
documents.set_parsed(document["id"], elements)
|
|
inserted.append({"id": document["id"], "filename": filename})
|
|
|
|
pending = documents.list_for_extraction(batch_id)
|
|
total = len(pending)
|
|
if total > 10:
|
|
_start_extraction_job(batch_id, pending, owner)
|
|
return jsonify(ok=True, v3=True, total=total, documents=inserted), 202
|
|
|
|
result = _run_extraction_job(batch_id, pending, owner)
|
|
result["documents"] = inserted
|
|
return jsonify(result)
|
|
|
|
|
|
def _valid_engine_element(element):
|
|
"""Validate the parser fields consumed by structured-text formatting."""
|
|
if not isinstance(element, dict):
|
|
return False
|
|
if element.get("type") == "paragraph":
|
|
return isinstance(element.get("text"), str) and isinstance(element.get("style", ""), str)
|
|
if element.get("type") != "table":
|
|
return False
|
|
rows = element.get("rows")
|
|
context = element.get("preceding_paragraphs", [])
|
|
table_index = element.get("table_index")
|
|
return (
|
|
isinstance(rows, list)
|
|
and all(isinstance(row, list) and all(isinstance(cell, str) for cell in row) for row in rows)
|
|
and isinstance(context, list)
|
|
and all(isinstance(paragraph, str) for paragraph in context)
|
|
and (table_index is None or (type(table_index) is int and table_index > 0))
|
|
)
|
|
|
|
|
|
def _start_extraction_job(batch_id, pending, owner):
|
|
"""Run a claimed batch in the background with persistent progress."""
|
|
def _run():
|
|
try:
|
|
_run_extraction_job(batch_id, pending, owner)
|
|
except Exception:
|
|
pass
|
|
|
|
thread = threading.Thread(target=_run, daemon=True)
|
|
try:
|
|
thread.start()
|
|
except Exception as error:
|
|
extraction_jobs.finish(batch_id, owner, error=str(error))
|
|
raise
|
|
|
|
|
|
def _engine_rate_limited():
|
|
"""Apply a small per-source fixed-window limit before accepting LLM work."""
|
|
now = time.monotonic()
|
|
source = request.remote_addr or "unknown"
|
|
cutoff = now - _ENGINE_RATE_WINDOW_SECONDS
|
|
with _engine_rate_lock:
|
|
for key in list(_engine_rate_hits):
|
|
recent = [stamp for stamp in _engine_rate_hits[key] if stamp > cutoff]
|
|
if recent:
|
|
_engine_rate_hits[key] = recent
|
|
else:
|
|
del _engine_rate_hits[key]
|
|
recent = _engine_rate_hits.get(source, [])
|
|
if len(recent) >= _ENGINE_RATE_LIMIT:
|
|
return True
|
|
recent.append(now)
|
|
_engine_rate_hits[source] = recent
|
|
return False
|
|
|
|
|
|
def _run_extraction_job(batch_id, pending, owner):
|
|
"""Keep the lease alive during LLM calls and persist every final outcome."""
|
|
stop = threading.Event()
|
|
|
|
def renew_lease():
|
|
while not stop.wait(15):
|
|
if not extraction_jobs.heartbeat(batch_id, owner):
|
|
return
|
|
|
|
heartbeat = threading.Thread(target=renew_lease, daemon=True)
|
|
heartbeat.start()
|
|
try:
|
|
result = _run_v3_extraction(
|
|
batch_id, pending,
|
|
progress_callback=lambda completed, item: extraction_jobs.progress(batch_id, owner, completed, item),
|
|
)
|
|
extraction_jobs.finish(batch_id, owner, result=result)
|
|
return result
|
|
except Exception as error:
|
|
extraction_jobs.finish(batch_id, owner, error=str(error))
|
|
raise
|
|
finally:
|
|
stop.set()
|
|
heartbeat.join()
|
|
|
|
|
|
def _run_v3_extraction(batch_id, pending, progress_callback=None, llm_client=None):
|
|
"""Собрать extraction и dossier summaries в один ответ API."""
|
|
parsed_documents = [{
|
|
"id": document["id"],
|
|
"filename": document["filename"],
|
|
"elements": stateless_extractor._decode_elements(document.get("elements_json")),
|
|
} for document in pending]
|
|
return analyze_batch(
|
|
batch_id,
|
|
parsed_documents,
|
|
extractor=stateless_extractor.extract_batch,
|
|
dossier_builder=build_dossiers_for_batch,
|
|
llm_client=llm_client,
|
|
progress_callback=progress_callback,
|
|
)
|
|
|
|
|
|
@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"] not in {"failed", "interrupted"},
|
|
"v3": True,
|
|
"finished": job["state"] in {"done", "failed", "interrupted"},
|
|
"state": job["state"],
|
|
"done": job["done"],
|
|
"total": job["total"],
|
|
"failed": job["failed"],
|
|
"classes": job["classes"],
|
|
}
|
|
if job.get("error"):
|
|
response["error"] = job["error"]
|
|
if job.get("result") is not None:
|
|
response["result"] = job["result"]
|
|
return jsonify(response)
|