125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
"""Classify service — LLM-based document classification."""
|
|
import json, re, os
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
import httpx
|
|
from db import documents as db_docs
|
|
from llm_prompt import build_classify_prompt
|
|
|
|
log = __import__("logging").getLogger(__name__)
|
|
|
|
MAX_WORKERS = 4
|
|
LLM_URL = "https://api.aillm.ru/v1/chat/completions"
|
|
LLM_KEY = os.environ.get("LLM_KEY") or os.environ.get("LLM_API_KEY", "")
|
|
LLM_MODEL = "gpt-oss-120b"
|
|
|
|
|
|
def _call_llm_classify(header_text):
|
|
"""Call LLM for classification. Returns parsed JSON dict."""
|
|
prompt, _ = build_classify_prompt(header_text)
|
|
payload = {
|
|
"model": LLM_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 500,
|
|
"temperature": 0.1,
|
|
}
|
|
with httpx.Client(http2=True, timeout=60, verify=False) as client:
|
|
resp = client.post(
|
|
LLM_URL, json=payload,
|
|
headers={"Authorization": f"Bearer {LLM_KEY}", "Content-Type": "application/json"},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
raw = data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
# Extract JSON from possible markdown wrapping
|
|
json_text = raw
|
|
if "```json" in json_text:
|
|
json_text = json_text.split("```json")[1].split("```")[0]
|
|
elif "```" in json_text:
|
|
json_text = json_text.split("```")[1].split("```")[0]
|
|
return json.loads(json_text.strip())
|
|
|
|
|
|
def classify_batch(batch_id):
|
|
"""Classify all pending documents in a batch. Returns summary dict."""
|
|
pending = db_docs.list_pending(batch_id)
|
|
if not pending:
|
|
return {"ok": False, "error": "no pending documents"}
|
|
|
|
total = len(pending)
|
|
done = 0
|
|
failed = 0
|
|
|
|
def _classify_one(doc):
|
|
try:
|
|
text = _smart_extract(doc["elements_json"])
|
|
result = _call_llm_classify(text)
|
|
db_docs.set_classification(
|
|
doc["id"],
|
|
result.get("doc_type", "other"),
|
|
result.get("own_number"),
|
|
result.get("parent_number"),
|
|
result.get("doc_date"),
|
|
result.get("counterparty"),
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
db_docs.set_classify_failed(doc["id"], str(e))
|
|
return False
|
|
|
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
|
|
futures = {pool.submit(_classify_one, d): d for d in pending}
|
|
for f in as_completed(futures):
|
|
if f.result():
|
|
done += 1
|
|
else:
|
|
failed += 1
|
|
|
|
return {"ok": True, "total": total, "done": done, "failed": failed}
|
|
|
|
|
|
def _smart_extract(elements_json):
|
|
"""Extract smart header: first ~1500 chars + marker lines from full doc."""
|
|
if not elements_json:
|
|
return ""
|
|
|
|
if isinstance(elements_json, str):
|
|
try:
|
|
elements = json.loads(elements_json)
|
|
except json.JSONDecodeError:
|
|
return elements_json[:2000]
|
|
elif isinstance(elements_json, list):
|
|
elements = elements_json
|
|
else:
|
|
return str(elements_json)[:2000]
|
|
|
|
# Build full text
|
|
lines = []
|
|
for el in elements:
|
|
if isinstance(el, dict):
|
|
t = el.get("type") or el.get("TYPE", "")
|
|
if t == "paragraph":
|
|
txt = el.get("text") or el.get("TEXT", "")
|
|
if txt:
|
|
lines.append(txt)
|
|
elif t == "table":
|
|
rows = el.get("rows") or el.get("ROWS", [])
|
|
for row in rows:
|
|
lines.append(" | ".join(str(c) for c in row))
|
|
|
|
full_text = "\n".join(lines)
|
|
|
|
# Header: first ~1500 chars
|
|
header = full_text[:1500]
|
|
|
|
# Marker lines: grep for key patterns
|
|
markers = re.findall(
|
|
r'.{0,200}(?:договор|№|соглашен|приложен|специф|контрагент|заказчик|арендатор).{0,200}',
|
|
full_text, re.IGNORECASE,
|
|
)
|
|
unique_markers = list(dict.fromkeys(markers))[:10]
|
|
|
|
combined = header + "\n---\n" + "\n".join(unique_markers)
|
|
return combined[:3000]
|