"""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)