contracts-app/site/routes/llm.py

"""routes/llm.py — LLM-извлечение + сравнение + чат."""

import json as _json
import time as _time

from flask import Response, request, jsonify

import db
import extractor
import differ
import llm_client


def process(cid):
    """
    GET /llm/process/<cid> — SSE-поток полного цикла:

      1. Для каждого файла договора: parsed_text → LLM → spec_rows
      2. Сравнение допников через differ.diff()
      3. Результат → SSE-сообщения в браузер
    """
    def generate():
        start_time = _time.time()

        # ── Получить все файлы договора ─────────────────────
        supps, _ = db.query(
            "SELECT s.id as supp_id, s.type, d.id as doc_id, d.filename, d.parsed_text "
            "FROM supplements s JOIN documents d ON s.document_id = d.id "
            "WHERE s.contract_id = %s AND d.parsed_text IS NOT NULL "
            "ORDER BY s.created_at",
            (cid,),
        )
        if not supps:
            yield f"data: {_json.dumps({'type': 'error', 'message': 'Нет распарсенных файлов. Сначала нажмите Парсинг.'})}\n\n"
            return

        supp_rows = [dict(zip(supps["columns"], r)) for r in supps["rows"]]
        extracted = {}  # supp_id → [rows]

        # ── Шаг 1: LLM-извлечение для каждого файла ─────────
        for s in supp_rows:
            # Пропустить уже извлечённые
            existing, _ = db.query(
                "SELECT 1 FROM spec_rows WHERE supplement_id = %s LIMIT 1",
                (s["supp_id"],),
            )
            if existing and existing["rows"]:
                # Загрузить существующие
                rows_res, _ = db.query(
                    "SELECT row_num, name, price, qty, sum, date_start "
                    "FROM spec_rows WHERE supplement_id = %s ORDER BY row_num",
                    (s["supp_id"],),
                )
                rows = [dict(zip(rows_res["columns"], r)) for r in rows_res["rows"]] if rows_res else []
                extracted[s["supp_id"]] = rows
                yield f"data: {_json.dumps({'type': 'extract_skip', 'name': s['filename'], 'reason': 'уже извлечено', 'count': len(rows)})}\n\n"
                continue

            yield f"data: {_json.dumps({'type': 'extract_start', 'name': s['filename']})}\n\n"

            t0 = _time.time()
            result = extractor.extract(s["parsed_text"])

            if "error" in result:
                yield f"data: {_json.dumps({'type': 'extract_error', 'name': s['filename'], 'error': result['error']})}\n\n"
                continue

            rows = result.get("rows", [])
            unresolved = result.get("unresolved", [])
            elapsed = round(_time.time() - t0, 2)

            # Сохранить в БД
            db.execute("DELETE FROM spec_rows WHERE supplement_id = %s", (s["supp_id"],))
            for row in rows:
                db.execute(
                    "INSERT INTO spec_rows (supplement_id, row_num, name, price, qty, sum, date_start) "
                    "VALUES (%s, %s, %s, %s, %s, %s, %s)",
                    (
                        s["supp_id"],
                        row.get("row_num"),
                        row.get("name"),
                        row.get("price"),
                        row.get("qty"),
                        row.get("sum"),
                        row.get("date_start"),
                    ),
                )

            extracted[s["supp_id"]] = rows
            yield f"data: {_json.dumps({'type': 'extract_done', 'name': s['filename'], 'time_s': elapsed, 'count': len(rows), 'unresolved': unresolved})}\n\n"

        # ── Шаг 2: Сравнение ────────────────────────────────
        yield f"data: {_json.dumps({'type': 'diff_start', 'files': len(supp_rows)})}\n\n"

        # Группируем: ищем amendments, если нет — сравниваем первый с остальными
        initial_rows = None
        amendments = []

        for s in supp_rows:
            if s["type"] != "initial":
                amendments.append({
                    "supp_id": s["supp_id"],
                    "filename": s["filename"],
                    "type": s["type"],
                    "rows": extracted.get(s["supp_id"], []),
                })
            elif initial_rows is None:
                # Первый initial — базовый
                initial_rows = extracted.get(s["supp_id"], [])
            else:
                # Последующие initial — как допники (fallback)
                amendments.append({
                    "supp_id": s["supp_id"],
                    "filename": s["filename"],
                    "type": "initial",
                    "rows": extracted.get(s["supp_id"], []),
                })

        all_changes = []

        if initial_rows is not None:
            for am in amendments:
                if not am["rows"]:
                    continue
                d = differ.diff(initial_rows, am["rows"])
                changes = d.get("changes", [])
                summary = d.get("summary", {})
                all_changes.append({
                    "filename": am["filename"],
                    "type": am["type"],
                    "changes": changes,
                    "summary": summary,
                })
                yield f"data: {_json.dumps({'type': 'diff_file', 'name': am['filename'], 'changes': len(changes), 'summary': summary})}\n\n"

        # ── Итог ────────────────────────────────────────────
        total_time = round(_time.time() - start_time, 2)
        yield f"data: {_json.dumps({'type': 'done', 'total_time_s': total_time, 'all_changes': all_changes})}\n\n"

    return Response(generate(), mimetype="text/event-stream")


def chat(cid):
    """
    POST /llm/chat/<cid> — задать вопрос по договору.

    Тело: {"question": "..."}
    Контекст: все строки spec_rows этого договора.
    Ответ: {"answer": "..."}
    """
    data = request.get_json(silent=True) or {}
    question = data.get("question", "").strip()
    if not question:
        return jsonify({"error": "Введите вопрос"}), 400

    # ── Собрать контекст: все строки спецификации ──────────
    rows_res, _ = db.query(
        "SELECT s.type, s.number, d.filename, sr.row_num, sr.name, sr.price, sr.qty, sr.sum, sr.date_start "
        "FROM spec_rows sr "
        "JOIN supplements s ON sr.supplement_id = s.id "
        "JOIN documents d ON s.document_id = d.id "
        "WHERE s.contract_id = %s ORDER BY s.created_at, sr.row_num",
        (cid,),
    )
    if not rows_res or not rows_res["rows"]:
        return jsonify({"answer": "Нет извлечённых данных. Сначала запустите сравнение."})

    # Форматируем контекст
    lines = []
    for r in rows_res["rows"]:
        row = dict(zip(rows_res["columns"], r))
        lines.append(
            f"[{row['filename']}] строка {row['row_num']}: {row['name']} | "
            f"цена={row['price']} | объём={row['qty']} | сумма={row['sum']} | начало={row['date_start']}"
        )
    context = "\n".join(lines)

    # ── Промпт ──────────────────────────────────────────────
    prompt = f"""Ты — анализатор договоров облачного провайдера. У тебя есть извлечённые строки спецификации услуг.

ДАННЫЕ:
{context}

ВОПРОС ПОЛЬЗОВАТЕЛЯ:
{question}

Ответь кратко и по делу, опираясь ТОЛЬКО на данные выше. Если данных недостаточно — скажи об этом."""

    # ── Вызов LLM ───────────────────────────────────────────
    result = llm_client.ask(prompt, max_tokens=1000)
    if "error" in result:
        return jsonify({"answer": f"Ошибка LLM: {result['error']}"})

    return jsonify({"answer": result.get("text", "Нет ответа")})