fix(ревью#2): db_docs→_db + classify_raw в MemRepo + try/except в build_classify_prompt

This commit is contained in:
2026-06-28 09:40:00 +04:00
parent 8a804da494
commit d8e8b35092
4 changed files with 89 additions and 23 deletions
+8 -6
View File
@@ -226,12 +226,14 @@ def build_prompt(current_spec: list, doc_text: str) -> tuple:
def build_classify_prompt(header_text): def build_classify_prompt(header_text):
"""Build classify prompt. Returns (prompt, prompt_id).""" """Build classify prompt. Returns (prompt, prompt_id)."""
from db import prompts as db_prompts try:
prompt = db_prompts.get_active("classify") from db import prompts as db_prompts
if prompt: prompt = db_prompts.get_active("classify")
body = prompt["body"].replace("{header_text}", header_text) if prompt:
return body, prompt.get("id", "") body = prompt["body"].replace("{header_text}", header_text)
# Fallback return body, prompt.get("id", "")
except Exception:
pass # DB unavailable — use fallback
body = """Ты — классификатор договорных документов облачного провайдера НУБЕС. body = """Ты — классификатор договорных документов облачного провайдера НУБЕС.
Ниже фрагмент текста документа. Определи: Ниже фрагмент текста документа. Определи:
+12 -10
View File
@@ -31,7 +31,8 @@ class Repository(Protocol):
def set_classification(self, doc_id: str, doc_type: str, own_number: str = None, def set_classification(self, doc_id: str, doc_type: str, own_number: str = None,
parent_number: str = None, doc_date: str = None, parent_number: str = None, doc_date: str = None,
counterparty: str = None) -> None: counterparty: str = None,
classify_raw: str = None, classify_input: str = None) -> None:
"""Сохранить результат классификации.""" """Сохранить результат классификации."""
... ...
@@ -107,10 +108,12 @@ class PgRepository:
documents.set_error(doc_id, error) documents.set_error(doc_id, error)
def set_classification(self, doc_id, doc_type, own_number=None, parent_number=None, def set_classification(self, doc_id, doc_type, own_number=None, parent_number=None,
doc_date=None, counterparty=None): doc_date=None, counterparty=None,
classify_raw=None, classify_input=None):
from db import documents from db import documents
documents.set_classification(doc_id, doc_type, own_number, parent_number, documents.set_classification(doc_id, doc_type, own_number, parent_number,
doc_date, counterparty) doc_date, counterparty,
classify_raw=classify_raw, classify_input=classify_input)
def set_classify_garbage(self, doc_id, reason=""): def set_classify_garbage(self, doc_id, reason=""):
from db import documents from db import documents
@@ -199,15 +202,14 @@ class MemRepository:
self.documents[doc_id]["error_message"] = error self.documents[doc_id]["error_message"] = error
def set_classification(self, doc_id, doc_type, own_number=None, parent_number=None, def set_classification(self, doc_id, doc_type, own_number=None, parent_number=None,
doc_date=None, counterparty=None): doc_date=None, counterparty=None,
classify_raw=None, classify_input=None):
if doc_id in self.documents: if doc_id in self.documents:
d = self.documents[doc_id] d = self.documents[doc_id]
d["doc_type"] = doc_type d.update({"doc_type": doc_type, "own_number": own_number,
d["own_number"] = own_number "parent_number": parent_number, "doc_date": doc_date,
d["parent_number"] = parent_number "counterparty": counterparty, "classify_raw": classify_raw,
d["doc_date"] = doc_date "classify_input": classify_input, "classify_status": "classified"})
d["counterparty"] = counterparty
d["classify_status"] = "classified"
def set_classify_garbage(self, doc_id, reason=""): def set_classify_garbage(self, doc_id, reason=""):
if doc_id in self.documents: if doc_id in self.documents:
+7 -7
View File
@@ -92,8 +92,8 @@ def classify_batch(batch_id, llm_client=None, repo=None):
_llm = llm_client if llm_client else _get_classify_client() _llm = llm_client if llm_client else _get_classify_client()
# Сбросить статус — allow re-classify after file changes # Сбросить статус — allow re-classify after file changes
db_docs.reset_classify_status(batch_id) _db.reset_classify_status(batch_id)
pending = db_docs.list_pending(batch_id) pending = _db.list_pending(batch_id)
if not pending: if not pending:
return {"ok": False, "error": "no pending documents"} return {"ok": False, "error": "no pending documents"}
@@ -110,24 +110,24 @@ def classify_batch(batch_id, llm_client=None, repo=None):
try: try:
# Stage 1: garbage by filename (0 tokens) # Stage 1: garbage by filename (0 tokens)
if _is_garbage_by_filename(doc["filename"]): if _is_garbage_by_filename(doc["filename"]):
db_docs.set_classify_garbage(doc["id"], "filename_regex") _db.set_classify_garbage(doc["id"], "filename_regex")
garbage += 1 garbage += 1
return True return True
# Stage 2: garbage by header keywords (0 tokens) # Stage 2: garbage by header keywords (0 tokens)
text = _smart_extract(doc["elements_json"]) text = _smart_extract(doc["elements_json"])
if _is_garbage_by_header(text): if _is_garbage_by_header(text):
db_docs.set_classify_garbage(doc["id"], "header_keywords") _db.set_classify_garbage(doc["id"], "header_keywords")
garbage += 1 garbage += 1
return True return True
# Stage 3: LLM classification (only for remaining) # Stage 3: LLM classification (only for remaining)
db_docs.set_classify_processing(doc["id"]) # crash recovery marker _db.set_classify_processing(doc["id"]) # crash recovery marker
result, raw, needed_fix = _call_llm_classify(text, _llm) result, raw, needed_fix = _call_llm_classify(text, _llm)
json_total += 1 json_total += 1
if needed_fix: if needed_fix:
json_fixes += 1 json_fixes += 1
db_docs.set_classification( _db.set_classification(
doc["id"], doc["id"],
result.get("doc_type", "other"), result.get("doc_type", "other"),
result.get("own_number"), result.get("own_number"),
@@ -139,7 +139,7 @@ def classify_batch(batch_id, llm_client=None, repo=None):
) )
return True return True
except Exception as e: except Exception as e:
db_docs.set_classify_failed(doc["id"], str(e)) _db.set_classify_failed(doc["id"], str(e))
return False return False
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool: with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+62
View File
@@ -0,0 +1,62 @@
"""Test classify_batch with MemRepository + FakeLLM."""
from services.classify import classify_batch
from services.llm_client import FakeLLMClient
from repository import MemRepository
class TestClassifyBatch:
def test_garbage_by_filename(self):
"""Счёт-фактура отфильтровывается без LLM."""
repo = MemRepository()
llm = FakeLLMClient({"default": '{"doc_type":"contract"}'})
repo.insert_document("счет-фактура №123.docx", "mime", "data", batch_id="b1")
result = classify_batch("b1", llm_client=llm, repo=repo)
assert result["ok"] is True
assert result["garbage"] == 1
assert result["done"] == 1
assert len(llm.calls) == 0 # LLM не вызывался
def test_garbage_by_header(self):
"""Акт сверки в тексте отфильтровывается."""
repo = MemRepository()
llm = FakeLLMClient({"default": '{"doc_type":"contract"}'})
doc = repo.insert_document("документ.docx", "mime", "data", batch_id="b1")
repo.set_document_parsed(doc["id"], [
{"type": "paragraph", "text": "АКТ СВЕРКИ взаимных расчётов", "style": ""}
])
result = classify_batch("b1", llm_client=llm, repo=repo)
assert result["garbage"] == 1
assert len(llm.calls) == 0
def test_llm_classify(self):
"""Договор классифицируется через LLM."""
repo = MemRepository()
llm = FakeLLMClient({
"default": '{"doc_type":"contract","own_number":"03700_1","doc_date":"2026-02-01","counterparty":"ЗАО XXX001"}'
})
doc = repo.insert_document("договор-XXX001.docx", "mime", "data", batch_id="b1")
repo.set_document_parsed(doc["id"], [
{"type": "paragraph", "text": "ДОГОВОР № 03700_1", "style": ""}
])
result = classify_batch("b1", llm_client=llm, repo=repo)
assert result["done"] == 1, f"done={result['done']}, failed={result['failed']}, doc={repo.get_document(doc['id'])}"
assert result["garbage"] == 0
assert len(llm.calls) == 1
updated = repo.get_document(doc["id"])
assert updated["doc_type"] == "contract"
assert updated["own_number"] == "03700_1"
assert updated["counterparty"] == "ЗАО XXX001"
def test_no_pending(self):
"""Пустой батч."""
repo = MemRepository()
result = classify_batch("empty", llm_client=FakeLLMClient(), repo=repo)
assert result["ok"] is False
assert "no pending" in result["error"]