92 lines
3.0 KiB
Python
92 lines
3.0 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.classify import classify_batch
|
|
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] = {}
|
|
|
|
|
|
@pipeline_bp.route("/process-v2", methods=["GET"])
|
|
def process_v2():
|
|
"""SSE-стриминг сравнения договоров."""
|
|
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():
|
|
"""Запустить классификацию. ≤10 файлов — sync, >10 — async (Thread)."""
|
|
body = request.get_json()
|
|
batch_id = body.get("batch_id")
|
|
if not batch_id:
|
|
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_pending(batch_id)
|
|
total = len(pending)
|
|
if total == 0:
|
|
return jsonify(ok=False, error="no pending documents"), 400
|
|
|
|
# Sync для малых батчей
|
|
if total <= 10:
|
|
result = classify_batch(batch_id)
|
|
return jsonify(result)
|
|
|
|
# Async для больших
|
|
def _run():
|
|
try:
|
|
classify_batch(batch_id)
|
|
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), 202
|