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