111 lines
4.1 KiB
Python
111 lines
4.1 KiB
Python
"""CRUD для результатов stateless-извлечения отдельных документов.
|
|
|
|
Исходный документ и результат LLM разделены намеренно: повторный запуск
|
|
извлечения может обновить паспорт, не затрагивая байты файла и его структуру.
|
|
Составные поля сохраняются как JSON-текст, поскольку SQLite не требует
|
|
отдельного JSON-типа и существующий слой БД уже использует такой подход.
|
|
"""
|
|
import json
|
|
import uuid
|
|
|
|
from db.connection import execute, query
|
|
|
|
|
|
_JSON_FIELDS = (
|
|
"referenced_numbers_json",
|
|
"sections_json",
|
|
"math_validation_json",
|
|
)
|
|
_FIELDS = (
|
|
"doc_class",
|
|
"intent_mode",
|
|
"own_number",
|
|
"own_date",
|
|
"counterparty_name",
|
|
"counterparty_inn",
|
|
*_JSON_FIELDS,
|
|
"status",
|
|
"prompt_version",
|
|
"raw_llm_response",
|
|
)
|
|
|
|
|
|
def _encode_json(value):
|
|
"""Сериализовать значения в JSON, сохраняя Decimal как точную строку."""
|
|
return json.dumps(value, ensure_ascii=False, default=str)
|
|
|
|
|
|
def _decode_row(row):
|
|
"""Вернуть составные поля паспорта как Python-объекты для сервисов."""
|
|
if not row:
|
|
return None
|
|
result = dict(row)
|
|
for field in _JSON_FIELDS:
|
|
raw_value = result.get(field)
|
|
try:
|
|
result[field.removesuffix("_json")] = json.loads(raw_value or "null")
|
|
except (TypeError, json.JSONDecodeError):
|
|
result[field.removesuffix("_json")] = None
|
|
return result
|
|
|
|
|
|
def upsert_extraction(document_id, data):
|
|
"""Создать или заменить паспорт документа, не меняя его стабильный ID.
|
|
|
|
Входные JSON-поля принимают как Python-списки/словари, так и уже
|
|
сериализованные строки. Поля, которых нет в `data`, получают безопасные
|
|
значения по умолчанию при вставке и остаются без изменений при обновлении.
|
|
"""
|
|
existing = query(
|
|
"SELECT id FROM document_extractions WHERE document_id = %s",
|
|
(document_id,),
|
|
)
|
|
values = {}
|
|
for field in _FIELDS:
|
|
if field not in data:
|
|
continue
|
|
value = data[field]
|
|
if field in _JSON_FIELDS and not isinstance(value, str):
|
|
value = _encode_json(value)
|
|
values[field] = value
|
|
|
|
if existing:
|
|
if values:
|
|
assignments = ", ".join(f"{field} = %s" for field in values)
|
|
params = [*values.values(), document_id]
|
|
execute(
|
|
f"UPDATE document_extractions SET {assignments}, "
|
|
"updated_at = datetime('now') WHERE document_id = %s",
|
|
params,
|
|
)
|
|
return get_by_document(document_id)
|
|
|
|
extraction_id = str(uuid.uuid4())
|
|
columns = ["id", "document_id", *values.keys()]
|
|
placeholders = ", ".join(["%s"] * len(columns))
|
|
execute(
|
|
f"INSERT INTO document_extractions ({', '.join(columns)}) "
|
|
f"VALUES ({placeholders})",
|
|
[extraction_id, document_id, *values.values()],
|
|
)
|
|
return get_by_document(document_id)
|
|
|
|
|
|
def get_by_document(document_id):
|
|
"""Получить паспорт документа или `None`, если он ещё не извлечён."""
|
|
rows = query(
|
|
"SELECT * FROM document_extractions WHERE document_id = %s",
|
|
(document_id,),
|
|
)
|
|
return _decode_row(rows[0]) if rows else None
|
|
|
|
|
|
def list_by_batch(batch_id):
|
|
"""Получить паспорта всех документов batch в порядке загрузки."""
|
|
rows = query(
|
|
"""SELECT e.* FROM document_extractions e
|
|
JOIN documents d ON d.id = e.document_id
|
|
WHERE d.batch_id = %s ORDER BY d.created_at, d.id""",
|
|
(batch_id,),
|
|
)
|
|
return [_decode_row(row) for row in rows] |