v1.0.162: auto-classify backend — classify/grouping services, 7 new columns, /classify-batch, /api/groups, /apply-groups

2026-06-24 08:22:21 +04:00
parent 740f75a791
commit 7a96f5b95b
8 changed files with 372 additions and 6 deletions
@@ -0,0 +1,124 @@
+"""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]