531 lines
24 KiB
Python
531 lines
24 KiB
Python
"""app.py — Точка входа и сборка слоёв."""
|
||
|
||
from dotenv import load_dotenv
|
||
from flask import Flask, render_template, request, redirect, url_for, Response, jsonify
|
||
import json
|
||
|
||
import sys as _sys, os as _os
|
||
_sys.path.insert(0, _os.path.join(_os.path.dirname(__file__), '..'))
|
||
|
||
import schema
|
||
import db
|
||
import parser as parser_mod
|
||
import textify
|
||
import extractor
|
||
import differ
|
||
import mimeutil
|
||
from test_routes import test_bp
|
||
from upload import upload_bp
|
||
from api import api_bp
|
||
from v2.chunk_api import bp as v2_bp
|
||
import redis_client
|
||
import redis_worker
|
||
|
||
load_dotenv()
|
||
|
||
import time as _time
|
||
import re as _re
|
||
|
||
# Логи в памяти (быстро, не тормозит ответ)
|
||
_log_memory = []
|
||
# Хранилище чанков в памяти
|
||
_chunks_mem = {}
|
||
|
||
def _log(step, detail=""):
|
||
"""Писать лог в память (мгновенно)."""
|
||
_log_memory.append({"step": step, "detail": str(detail)[:500], "time": _time.time()})
|
||
if len(_log_memory) > 500:
|
||
_log_memory.pop(0)
|
||
|
||
|
||
def _save_file_to_db(filename, file_bytes, mime, contract_id, conn=None):
|
||
"""Сохранить файл в БД. Если conn передан — использовать его, иначе свои подключения."""
|
||
if conn:
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"INSERT INTO documents (filename, mime_type, original_bytes, status) VALUES (%s,%s,%s,'uploaded')",
|
||
(filename, mime, file_bytes),
|
||
)
|
||
cur.execute(
|
||
"SELECT id FROM documents WHERE filename=%s AND status='uploaded' ORDER BY created_at DESC LIMIT 1",
|
||
(filename,),
|
||
)
|
||
row = cur.fetchone()
|
||
doc_id = row[0] if row else None
|
||
if doc_id:
|
||
cur.execute(
|
||
"INSERT INTO supplements (contract_id, document_id, type) VALUES (%s,%s,'initial')",
|
||
(str(contract_id), str(doc_id)),
|
||
)
|
||
conn.commit()
|
||
cur.close()
|
||
return doc_id
|
||
|
||
# Без соединения — каждое действие со своим (старый режим)
|
||
db.execute(
|
||
"INSERT INTO documents (filename, mime_type, original_bytes, status) VALUES (%s,%s,%s,'uploaded')",
|
||
(filename, mime, file_bytes),
|
||
)
|
||
doc, _ = db.query_one(
|
||
"SELECT id FROM documents WHERE filename=%s AND status='uploaded' ORDER BY created_at DESC LIMIT 1",
|
||
(filename,),
|
||
)
|
||
doc_id = doc["id"] if doc else None
|
||
if doc_id:
|
||
db.execute(
|
||
"INSERT INTO supplements (contract_id, document_id, type) VALUES (%s,%s,'initial')",
|
||
(str(contract_id), str(doc_id)),
|
||
)
|
||
return doc_id
|
||
|
||
|
||
class ContractsApp:
|
||
def __init__(self):
|
||
self.app = Flask(__name__)
|
||
schema.ensure_schema()
|
||
redis_worker.start_worker()
|
||
self._add_routes()
|
||
|
||
def _add_routes(self):
|
||
self.app.add_url_rule("/", "index", self._index, methods=["GET", "POST"])
|
||
self.app.add_url_rule("/parse/<cid>", "parse", self._parse, methods=["GET"])
|
||
self.app.add_url_rule("/chunk_json", "chunk_json", self._chunk_json, methods=["POST"])
|
||
self.app.add_url_rule("/upload", "upload_json", self._upload_json, methods=["POST"])
|
||
self.app.add_url_rule("/health", "health", self._health)
|
||
self.app.add_url_rule("/logs", "logs", self._logs)
|
||
self.app.register_blueprint(test_bp)
|
||
self.app.register_blueprint(upload_bp)
|
||
self.app.register_blueprint(api_bp)
|
||
self.app.register_blueprint(v2_bp)
|
||
|
||
def _health(self):
|
||
return "OK", 200, {"Content-Type": "text/plain"}
|
||
|
||
def _logs(self):
|
||
return jsonify(_log_memory[-50:]) # последние 50 записей
|
||
|
||
# ── Шаг 1: загрузка файлов ──────────────────────────────────
|
||
|
||
def _index(self):
|
||
if request.method == "POST":
|
||
return self._upload_files()
|
||
cid = request.args.get("id")
|
||
if cid:
|
||
return self._show_contract(cid)
|
||
return render_template("upload.html")
|
||
|
||
def _upload_files(self):
|
||
"""Сохранить файлы. ?cid=X — добавить к существующему договору."""
|
||
files = request.files.getlist("files")
|
||
if not files or not any(f.filename for f in files):
|
||
return jsonify({"error": "Нет файлов"}), 400
|
||
|
||
cid = request.args.get("cid")
|
||
|
||
if cid:
|
||
contract_id = cid
|
||
else:
|
||
conn, err = db.connect()
|
||
if err:
|
||
return jsonify({"error": f"БД: {err}"}), 500
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"INSERT INTO contracts (number, client) VALUES (%s, %s) RETURNING id",
|
||
("б/н " + __import__("datetime").datetime.now().strftime("%Y%m%d-%H%M"), ""),
|
||
)
|
||
contract_id = cur.fetchone()[0]
|
||
conn.commit()
|
||
cur.close()
|
||
db.put_conn(conn)
|
||
|
||
for f in files:
|
||
if not f.filename:
|
||
continue
|
||
filename = f.filename
|
||
mime = mimeutil.guess_mime(filename) or f.content_type or "application/octet-stream"
|
||
_save_file_to_db(filename, f.read(), mime, str(contract_id))
|
||
|
||
return jsonify({"contract_id": str(contract_id)})
|
||
|
||
# ── Загрузка base64 (JSON) ──────────────────────────────────
|
||
|
||
def _upload_json(self):
|
||
"""Принять файл как base64. Без JSON-парсинга — сырой разбор."""
|
||
_log("upload_entry", "start")
|
||
try:
|
||
raw = request.get_data(as_text=True)
|
||
import re
|
||
m = re.search(r'"data"\s*:\s*"([^"]*)"', raw)
|
||
if not m:
|
||
return jsonify({"error": "Нет data"}), 400
|
||
b64 = m.group(1)
|
||
|
||
m = re.search(r'"filename"\s*:\s*"([^"]*)"', raw)
|
||
filename = m.group(1) if m else ""
|
||
|
||
m = re.search(r'"cid"\s*:\s*"([^"]*)"', raw)
|
||
cid = m.group(1) if m else None
|
||
|
||
if not filename or not b64:
|
||
return jsonify({"error": "Нет filename или data"}), 400
|
||
|
||
mime = mimeutil.guess_mime(filename) or "application/octet-stream"
|
||
|
||
# Контракт — синхронно (быстро, нужно для следующих загрузок)
|
||
if not cid:
|
||
conn, err = db.connect()
|
||
if err:
|
||
return jsonify({"error": f"БД: {err}"}), 500
|
||
try:
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"INSERT INTO contracts (number, client) VALUES (%s, %s) RETURNING id",
|
||
("б/н " + __import__("datetime").datetime.now().strftime("%Y%m%d-%H%M"), ""),
|
||
)
|
||
cid = cur.fetchone()[0]
|
||
conn.commit()
|
||
cur.close()
|
||
finally:
|
||
db.put_conn(conn)
|
||
|
||
# Файл — в Redis (в фоне, не блокируем ответ)
|
||
import threading
|
||
task = {"filename": filename, "b64": b64, "mime": mime, "cid": str(cid)}
|
||
threading.Thread(target=redis_client.push, args=(task,), daemon=True).start()
|
||
|
||
return jsonify({"contract_id": str(cid)})
|
||
except Exception as e:
|
||
_log("upload_outer_error", str(e))
|
||
return jsonify({"error": f"Крах: {e}"}), 500
|
||
|
||
# ── Чанковая загрузка base64 (30KB) ───────────────────────
|
||
|
||
def _chunk_json(self):
|
||
"""Принять чанк base64 (в памяти). На последнем — в Redis."""
|
||
data = request.get_json(silent=True) or {}
|
||
upload_id = data.get("upload_id", "")
|
||
filename = data.get("filename", "")
|
||
chunk_index = data.get("chunk_index", 0)
|
||
total_chunks = data.get("total_chunks", 1)
|
||
b64_piece = data.get("data", "")
|
||
cid = data.get("cid")
|
||
|
||
if not upload_id or not b64_piece:
|
||
return jsonify({"error": "Нет upload_id или data"}), 400
|
||
|
||
if upload_id not in _chunks_mem:
|
||
_chunks_mem[upload_id] = {"chunks": [None]*total_chunks, "filename": filename, "cid": cid, "received": 0}
|
||
|
||
store = _chunks_mem[upload_id]
|
||
if store["chunks"][chunk_index] is None:
|
||
store["chunks"][chunk_index] = b64_piece
|
||
store["received"] += 1
|
||
|
||
if store["received"] == total_chunks:
|
||
full_b64 = "".join(store["chunks"])
|
||
fname = store["filename"]
|
||
fcid = store["cid"]
|
||
mime = mimeutil.guess_mime(fname) or "application/octet-stream"
|
||
import threading
|
||
task = {"filename": fname, "b64": full_b64, "mime": mime, "cid": fcid or ""}
|
||
threading.Thread(target=redis_client.push, args=(task,), daemon=True).start()
|
||
del _chunks_mem[upload_id]
|
||
return jsonify({"contract_id": fcid or "pending"})
|
||
|
||
return jsonify({"chunk": chunk_index, "received": store["received"], "total": total_chunks})
|
||
|
||
# ── Старая чанковая загрузка ───────────────────────────────
|
||
|
||
def _chunk_upload(self):
|
||
"""Принять чанк. Хранить в БД. На последнем — собрать и сохранить."""
|
||
upload_id = request.form.get("upload_id", "")
|
||
filename = request.form.get("filename", "")
|
||
chunk_index = int(request.form.get("chunk_index", 0))
|
||
total_chunks = int(request.form.get("total_chunks", 1))
|
||
cid = request.form.get("cid") or None
|
||
chunk_file = request.files.get("chunk")
|
||
if not chunk_file:
|
||
return jsonify({"error": "Нет чанка"}), 400
|
||
|
||
chunk_data = chunk_file.read()
|
||
mime = mimeutil.guess_mime(filename) or "application/octet-stream"
|
||
|
||
_log("chunk_start", f"{filename} idx={chunk_index}/{total_chunks} size={len(chunk_data)} cid={cid}")
|
||
|
||
# Сохранить чанк в БД (ON CONFLICT — идемпотентно)
|
||
rc, err = db.execute(
|
||
"INSERT INTO chunks (upload_id, chunk_index, chunk_data, filename, mime, total_chunks, cid) "
|
||
"VALUES (%s,%s,%s,%s,%s,%s,%s) ON CONFLICT (upload_id, chunk_index) DO NOTHING",
|
||
(upload_id, chunk_index, chunk_data, filename, mime, total_chunks, cid),
|
||
)
|
||
_log("chunk_insert", f"idx={chunk_index} rc={rc} err={err}")
|
||
|
||
# Сколько уже получено
|
||
count_res, cerr = db.query(
|
||
"SELECT COUNT(*) FROM chunks WHERE upload_id=%s", (upload_id,)
|
||
)
|
||
received = count_res["rows"][0][0] if count_res else 0
|
||
_log("chunk_count", f"received={received}/{total_chunks} err={cerr}")
|
||
|
||
if received == total_chunks:
|
||
_log("assembly_start", f"{filename} {total_chunks} chunks, {sum(len(c) if c else 0 for c in [chunk_data])}b this chunk")
|
||
try:
|
||
conn, err = db.connect()
|
||
_log("assembly_connect", f"err={err}")
|
||
if err:
|
||
return jsonify({"error": f"БД: {err}"}), 500
|
||
cur = conn.cursor()
|
||
|
||
cur.execute(
|
||
"SELECT chunk_data FROM chunks WHERE upload_id=%s ORDER BY chunk_index",
|
||
(upload_id,),
|
||
)
|
||
all_chunks = cur.fetchall()
|
||
_log("assembly_select", f"{len(all_chunks)} rows")
|
||
|
||
file_bytes = b"".join(r[0] for r in all_chunks)
|
||
_log("assembly_joined", f"{len(file_bytes)} bytes")
|
||
|
||
cur.execute(
|
||
"SELECT filename, mime, cid FROM chunks WHERE upload_id=%s LIMIT 1",
|
||
(upload_id,),
|
||
)
|
||
meta = cur.fetchone()
|
||
fname = meta[0] if meta else filename
|
||
fmime = meta[1] if meta else mime
|
||
fcid = meta[2] if meta else cid
|
||
_log("assembly_meta", f"fname={fname} mime={fmime} cid={fcid}")
|
||
|
||
if fcid:
|
||
contract_id = fcid
|
||
else:
|
||
cur.execute(
|
||
"INSERT INTO contracts (number, client) VALUES (%s, %s) RETURNING id",
|
||
("б/н " + __import__("datetime").datetime.now().strftime("%Y%m%d-%H%M"), ""),
|
||
)
|
||
contract_id = cur.fetchone()[0]
|
||
_log("assembly_contract", f"new cid={contract_id}")
|
||
|
||
conn.commit()
|
||
_log("assembly_commit", "ok")
|
||
|
||
doc_id = _save_file_to_db(fname, file_bytes, fmime, str(contract_id), conn=conn)
|
||
_log("save_file", f"doc_id={doc_id}")
|
||
if not doc_id:
|
||
conn.rollback()
|
||
cur.close()
|
||
db.put_conn(conn)
|
||
return jsonify({"error": "Не удалось сохранить файл в БД"}), 500
|
||
|
||
cur.close()
|
||
db.put_conn(conn)
|
||
_log("assembly_done", f"cid={contract_id}")
|
||
except Exception as e:
|
||
_log("assembly_error", str(e))
|
||
try: conn.rollback()
|
||
except: pass
|
||
try: cur.close()
|
||
except: pass
|
||
try: db.put_conn(conn)
|
||
except: pass
|
||
return jsonify({"error": f"Сборка: {e}"}), 500
|
||
finally:
|
||
db.execute("DELETE FROM chunks WHERE upload_id=%s", (upload_id,))
|
||
_log("chunks_cleanup", f"upload_id={upload_id}")
|
||
|
||
return jsonify({"contract_id": str(contract_id)})
|
||
|
||
_log("chunk_ok", f"idx={chunk_index} received={received}/{total_chunks}")
|
||
return jsonify({"chunk": chunk_index, "received": received, "total": total_chunks})
|
||
|
||
def _show_contract(self, cid):
|
||
"""Показать договор с кнопкой «Обработать»."""
|
||
contract, _ = db.query_one("SELECT id,number,created_at FROM contracts WHERE id=%s", (cid,))
|
||
if not contract:
|
||
return render_template("upload.html", error="Договор не найден")
|
||
|
||
supps, _ = db.query(
|
||
"SELECT s.id, s.created_at, d.filename, d.status, d.parsed_text IS NOT NULL as has_text "
|
||
"FROM supplements s JOIN documents d ON s.document_id=d.id "
|
||
"WHERE s.contract_id=%s ORDER BY s.created_at",
|
||
(cid,),
|
||
)
|
||
supp_list = [dict(zip(supps["columns"], r)) for r in supps["rows"]] if supps else []
|
||
|
||
# Есть ли уже результаты?
|
||
rows_res, _ = db.query(
|
||
"SELECT 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 "
|
||
"WHERE s.contract_id=%s ORDER BY sr.row_num",
|
||
(cid,),
|
||
)
|
||
all_rows = [dict(zip(rows_res["columns"], r)) for r in rows_res["rows"]] if rows_res else []
|
||
|
||
# Есть необработанные допники?
|
||
unprocessed = [s for s in supp_list if s["status"] in ("uploaded", "parsed")]
|
||
|
||
return render_template("upload.html",
|
||
contract=contract, supplements=supp_list,
|
||
all_rows=all_rows, unprocessed=unprocessed)
|
||
|
||
# ── Шаг 2: Парсинг (SSE) ────────────────────────────────────
|
||
|
||
def _parse(self, cid):
|
||
"""SSE-поток: парсинг файлов договора (без LLM)."""
|
||
import time as time_mod
|
||
|
||
def generate():
|
||
start_time = time_mod.time()
|
||
total_bytes = 0
|
||
files_processed = 0
|
||
|
||
supps, _ = db.query(
|
||
"SELECT s.id, d.id as doc_id, d.filename, d.mime_type, "
|
||
"LENGTH(d.original_bytes) as file_size "
|
||
"FROM supplements s JOIN documents d ON s.document_id=d.id "
|
||
"WHERE s.contract_id=%s",
|
||
(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"]]
|
||
|
||
def _file_done(name, elapsed, elements, errors, text):
|
||
paragraphs = sum(1 for e in elements if e.get("type") == "paragraph")
|
||
tables = sum(1 for e in elements if e.get("type") == "table")
|
||
table_rows = sum(len(e.get("rows", [])) for e in elements if e.get("type") == "table")
|
||
table_headers = [
|
||
e["rows"][0] if e.get("rows") else []
|
||
for e in elements if e.get("type") == "table"
|
||
]
|
||
return {
|
||
"type": "file_done",
|
||
"name": name,
|
||
"time_s": elapsed,
|
||
"elements": len(elements),
|
||
"paragraphs": paragraphs,
|
||
"tables": tables,
|
||
"table_rows": table_rows,
|
||
"table_headers": table_headers,
|
||
"errors": errors,
|
||
"text_len": len(text) if text else 0,
|
||
"text_preview": text[:200] if text else "",
|
||
}
|
||
|
||
for s in supp_rows:
|
||
# Пропустить уже распарсенные
|
||
existing, _ = db.query(
|
||
"SELECT 1 FROM documents WHERE id=%s AND status IN ('parsed','expanded')", (s["doc_id"],)
|
||
)
|
||
if existing and existing["rows"]:
|
||
continue
|
||
|
||
file_start = time_mod.time()
|
||
file_bytes = s.get("file_size", 0) or 0
|
||
total_bytes += file_bytes
|
||
|
||
yield f"data: {json.dumps({'type': 'file_start', 'name': s['filename'], 'bytes': file_bytes})}\n\n"
|
||
|
||
# Получить байты: сначала original_bytes, иначе original_b64
|
||
doc, _ = db.query_one("SELECT original_bytes, original_b64 FROM documents WHERE id=%s", (s["doc_id"],))
|
||
raw = doc.get("original_bytes") if doc else None
|
||
b64 = doc.get("original_b64") if doc else None
|
||
if raw:
|
||
file_data = bytes(raw) if isinstance(raw, memoryview) else raw
|
||
elif b64:
|
||
import base64 as b64mod
|
||
file_data = b64mod.b64decode(b64)
|
||
else:
|
||
yield f"data: {json.dumps({'type': 'file_error', 'name': s['filename'], 'error': 'Нет данных'})}\n\n"
|
||
continue
|
||
|
||
# Парсинг: ZIP — распаковать и парсить каждый файл отдельно
|
||
if s["mime_type"] == "application/zip":
|
||
import io as io_mod, zipfile as zf_mod
|
||
zip_names = []
|
||
try:
|
||
with zf_mod.ZipFile(io_mod.BytesIO(file_data)) as zf:
|
||
for zname in zf.namelist():
|
||
if zname.endswith("/"):
|
||
continue
|
||
zdata = zf.read(zname)
|
||
zmime = mimeutil.guess_mime(zname) or "application/octet-stream"
|
||
|
||
# Только PDF и Word
|
||
if zmime not in (
|
||
"application/pdf",
|
||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
"application/msword",
|
||
):
|
||
continue
|
||
|
||
zip_names.append(zname)
|
||
|
||
# Сохранить как отдельный документ
|
||
db.execute(
|
||
"INSERT INTO documents (filename, mime_type, original_bytes, status) VALUES (%s,%s,%s,'uploaded')",
|
||
(zname, zmime, zdata),
|
||
)
|
||
zdoc, _ = db.query_one(
|
||
"SELECT id FROM documents WHERE filename=%s ORDER BY created_at DESC LIMIT 1",
|
||
(zname,),
|
||
)
|
||
zdoc_id = zdoc["id"] if zdoc else None
|
||
if zdoc_id:
|
||
db.execute(
|
||
"INSERT INTO supplements (contract_id, document_id, type) VALUES (%s,%s,'initial')",
|
||
(cid, str(zdoc_id)),
|
||
)
|
||
|
||
# Парсинг внутреннего файла
|
||
zf_start = time_mod.time()
|
||
total_bytes += len(zdata)
|
||
yield f"data: {json.dumps({'type': 'file_start', 'name': zname, 'bytes': len(zdata)})}\n\n"
|
||
|
||
try:
|
||
zpr = parser_mod.parse(zdata, zmime)
|
||
zelements = zpr.get("elements", [])
|
||
ztext = textify.to_text(zelements) if zelements else ""
|
||
db.execute(
|
||
"UPDATE documents SET parsed_text=%s, elements_json=%s, status='parsed' WHERE id=%s",
|
||
(ztext, json.dumps(zelements, ensure_ascii=False), str(zdoc_id)),
|
||
)
|
||
zelapsed = round(time_mod.time() - zf_start, 2)
|
||
files_processed += 1
|
||
yield f"data: {json.dumps(_file_done(zname, zelapsed, zelements, zpr.get('errors',[]), ztext))}\n\n"
|
||
except Exception as e:
|
||
yield f"data: {json.dumps({'type': 'file_error', 'name': zname, 'error': str(e)})}\n\n"
|
||
|
||
except Exception as e:
|
||
yield f"data: {json.dumps({'type': 'file_error', 'name': s['filename'], 'error': 'ZIP: ' + str(e)})}\n\n"
|
||
|
||
# ZIP сам — expanded
|
||
db.execute("UPDATE documents SET status='expanded' WHERE id=%s", (s["doc_id"],))
|
||
yield f"data: {json.dumps({'type': 'zip_expanded', 'name': s['filename'], 'count': len(zip_names)})}\n\n"
|
||
|
||
else:
|
||
pr = parser_mod.parse(file_data, s["mime_type"])
|
||
elements = pr.get("elements", [])
|
||
text = textify.to_text(elements) if elements else ""
|
||
db.execute(
|
||
"UPDATE documents SET parsed_text=%s, elements_json=%s, status='parsed' WHERE id=%s",
|
||
(text, json.dumps(elements, ensure_ascii=False), str(s["doc_id"])),
|
||
)
|
||
elapsed = round(time_mod.time() - file_start, 2)
|
||
files_processed += 1
|
||
yield f"data: {json.dumps(_file_done(s['filename'], elapsed, elements, pr.get('errors',[]), text))}\n\n"
|
||
|
||
total_time = round(time_mod.time() - start_time, 2)
|
||
yield f"data: {json.dumps({'type': 'summary', 'total_time_s': total_time, 'total_bytes': total_bytes, 'files_processed': files_processed})}\n\n"
|
||
|
||
return Response(generate(), mimetype="text/event-stream")
|
||
|
||
def run(self):
|
||
self.app.run(host="0.0.0.0", port=5000)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
ContractsApp().run()
|