106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
"""Documents CRUD."""
|
|
from .connection import query, execute, execute_returning
|
|
|
|
|
|
def insert(filename, mime_type, original_bytes, status="uploaded", batch_id=None, zip_source=None):
|
|
"""Insert document, return row dict."""
|
|
return execute_returning(
|
|
"""INSERT INTO documents (filename, mime_type, original_bytes, status, batch_id, zip_source)
|
|
VALUES (%s, %s, %s, %s, %s, %s) RETURNING *""",
|
|
(filename, mime_type, original_bytes, status, batch_id, zip_source),
|
|
)
|
|
|
|
|
|
def get(doc_id):
|
|
rows = query("SELECT * FROM documents WHERE id = %s", (doc_id,))
|
|
return rows[0] if rows else None
|
|
|
|
|
|
def set_parsed(doc_id, elements_json):
|
|
"""Update elements_json + status='parsed'."""
|
|
import json
|
|
return execute(
|
|
"UPDATE documents SET elements_json = %s::jsonb, status = 'parsed' WHERE id = %s",
|
|
(json.dumps(elements_json, ensure_ascii=False), doc_id),
|
|
)
|
|
|
|
|
|
def set_error(doc_id, error_message):
|
|
return execute(
|
|
"UPDATE documents SET status = 'error', error_message = %s WHERE id = %s",
|
|
(error_message, doc_id),
|
|
)
|
|
|
|
|
|
def delete(doc_id):
|
|
return execute("DELETE FROM documents WHERE id = %s", (doc_id,))
|
|
|
|
|
|
def set_classification(doc_id, doc_type, own_number, parent_number, doc_date, counterparty, classify_raw=None, classify_input=None):
|
|
"""Store LLM classification results + raw response + input text."""
|
|
return execute(
|
|
"""UPDATE documents SET doc_type=%s, own_number=%s, parent_number=%s,
|
|
doc_date=%s, counterparty=%s, classify_status='classified', classify_raw=%s, classify_input=%s
|
|
WHERE id=%s""",
|
|
(doc_type, own_number, parent_number, doc_date, counterparty, classify_raw, classify_input, doc_id),
|
|
)
|
|
|
|
|
|
def set_classify_failed(doc_id, error):
|
|
return execute(
|
|
"UPDATE documents SET classify_status='failed', error_message=%s WHERE id=%s",
|
|
(error, doc_id),
|
|
)
|
|
|
|
|
|
def set_classify_garbage(doc_id, reason=""):
|
|
"""Mark document as garbage (Stage 1-2 filter, no LLM call)."""
|
|
return execute(
|
|
"UPDATE documents SET doc_type='garbage', classify_status='garbage', error_message=%s WHERE id=%s",
|
|
(f"garbage: {reason}", doc_id),
|
|
)
|
|
|
|
|
|
def list_pending(batch_id):
|
|
"""Documents waiting for classification."""
|
|
return query(
|
|
"SELECT id, filename, elements_json FROM documents WHERE batch_id=%s AND classify_status='pending'",
|
|
(batch_id,),
|
|
)
|
|
|
|
|
|
def reset_classify_status(batch_id):
|
|
"""Сбросить classify_status на 'pending' для всех документов батча (включая 'processing' — crash recovery)."""
|
|
return execute(
|
|
"UPDATE documents SET classify_status='pending', error_message=NULL WHERE batch_id=%s",
|
|
(batch_id,),
|
|
)
|
|
|
|
|
|
def set_classify_processing(doc_id):
|
|
"""Mark document as being processed (for crash recovery)."""
|
|
return execute(
|
|
"UPDATE documents SET classify_status='processing' WHERE id=%s",
|
|
(doc_id,),
|
|
)
|
|
|
|
|
|
def list_by_batch(batch_id):
|
|
"""All documents in a batch with classification fields."""
|
|
return query(
|
|
"""SELECT id, filename, status, doc_type, own_number, parent_number,
|
|
doc_date, counterparty, classify_status, error_message, classify_raw, classify_input,
|
|
zip_source
|
|
FROM documents WHERE batch_id=%s ORDER BY created_at""",
|
|
(batch_id,),
|
|
)
|
|
|
|
|
|
def count_by_status(batch_id):
|
|
"""Count documents by classify_status."""
|
|
rows = query(
|
|
"SELECT classify_status, count(*) as cnt FROM documents WHERE batch_id=%s GROUP BY classify_status",
|
|
(batch_id,),
|
|
)
|
|
return {r["classify_status"]: r["cnt"] for r in rows}
|