217 lines
9.4 KiB
Python
217 lines
9.4 KiB
Python
"""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()
|
|
custom_prompt = _prompts.get(cid) or None
|
|
result = extractor.extract(s["parsed_text"], custom_prompt=custom_prompt)
|
|
|
|
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")
|
|
|
|
|
|
# ── Хранилище промптов (в памяти) ──────────────────────────────
|
|
|
|
_prompts = {} # cid → prompt
|
|
|
|
|
|
def save_prompt(cid):
|
|
"""POST /llm/prompt/<cid> — сохранить промпт. Тело: {"prompt": "..."}"""
|
|
data = request.get_json(silent=True) or {}
|
|
prompt = data.get("prompt", "").strip()
|
|
if prompt:
|
|
_prompts[cid] = prompt
|
|
return jsonify({"ok": True})
|
|
return jsonify({"error": "пустой промпт"}), 400
|
|
|
|
|
|
def get_prompt(cid):
|
|
"""GET /llm/prompt/<cid> — получить сохранённый промпт (или DEFAULT_PROMPT)"""
|
|
import extractor
|
|
return jsonify({"prompt": _prompts.get(cid, "")})
|
|
|
|
|
|
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", "Нет ответа")})
|