contractor/deploy/db/prompts.py

"""Prompts CRUD."""
from .connection import query, execute, execute_returning


def _serialize(row):
    """Convert datetime fields to strings for JSON serialization (non-mutating)."""
    if row and row.get("created_at"):
        row = dict(row)
        row["created_at"] = str(row["created_at"])
    return row


def get_active(role):
    rows = query(
        "SELECT * FROM prompts WHERE role = %s AND is_active = true ORDER BY created_at DESC LIMIT 1",
        (role,),
    )
    return _serialize(rows[0]) if rows else None


def get(prompt_id):
    rows = query("SELECT * FROM prompts WHERE id = %s", (prompt_id,))
    return rows[0] if rows else None


def seed_defaults():
    """Auto-seed default prompts if table is empty."""
    # Check each role separately — don't skip if one is missing
    existing_roles = set(r["role"] for r in query("SELECT DISTINCT role FROM prompts"))
    
    if "extract" not in existing_roles:
        extract_body = (
            "Ты — анализатор договоров облачного провайдера.\n\n"
            "Ниже текст спецификации услуг из ПЕРВОГО документа (базовый договор).\n"
            "Извлеки ВСЕ строки спецификации в JSON-массив.\n\n"
            "Верни СТРОГО JSON без пояснений. Не используй markdown-блоки.\n\n"
            "ФОРМАТ:\n"
            '{\n  "mode": "partial",\n  "ops": [\n'
            '    {"action": "ADD", "new_row": {"name": "полное наименование", "price": число, "qty": число, "sum": число, "date_start": "YYYY-MM-DD"}, "comment": ""}\n'
            "  ]\n}\n\n"
            "ПРАВИЛА:\n"
            "1. Извлеки КАЖДУЮ строку таблицы спецификации как отдельную ADD-операцию.\n"
            '2. Пропускай итоговые строки ("Итого...") и строки с подписями.\n'
            "3. Если ячейка пустая — ставь null (не пиши 0).\n"
            "4. price, qty, sum — ЧИСЛА, не строки.\n\n"
            "ТЕКСТ ДОКУМЕНТА:\n---\n{doc_text}\n---"
        )
        execute(
            "INSERT INTO prompts (role, name, body, is_active, notes) VALUES (%s, %s, %s, %s, %s)",
            ("extract", "default-v1", extract_body, True, "Авто-создан из llm_prompt.py _build_initial"),
        )
    
    if "diff" not in existing_roles:
        diff_body = (
            "Ты — анализатор допсоглашений к договорам облачного провайдера.\n\n"
            "У тебя есть текущая спецификация услуг и текст нового допсоглашения (ДС).\n"
            "Твоя задача — определить, какие изменения вносит ДС в текущую спецификацию.\n\n"
            "Верни СТРОГО JSON без пояснений. Не используй markdown-блоки.\n\n"
            "ФОРМАТ:\n"
            '{\n  "mode": "partial" | "full_replace",\n  "ops": [\n'
            '    {"action": "ADD", "new_row": {"name": "...", "price": число, "qty": число, "sum": число, "date_start": "YYYY-MM-DD"}, "comment": "..."},\n'
            '    {"action": "UPDATE", "target_id": "rN", "new_values": {"price": число}, "comment": "..."},\n'
            '    {"action": "DELETE", "target_id": "rN", "comment": "..."},\n'
            '    {"action": "UNRESOLVED", "new_values": {"name": "...", "price": число, ...}, "reason": "почему не смог сопоставить"}\n'
            "  ]\n}\n\n"
            "ПРАВИЛА:\n"
            '1. mode = "full_replace" — если в тексте есть фразы: «в следующей редакции», «заменить приложение», «излагается в следующей редакции». При full_replace — опиши ВСЕ новые строки как ADD.\n'
            '2. mode = "partial" — если ДС меняет только отдельные строки.\n'
            "3. Для UPDATE/DELETE — укажи target_id (r1, r2...) из списка текущей спецификации. НЕ придумывай новые id.\n"
            "4. new_values в UPDATE — только ИЗМЕНЁННЫЕ поля (не все).\n"
            "5. Если не можешь однозначно сопоставить строку — action: UNRESOLVED с reason.\n"
            "6. Пропускай итоговые строки и подписи.\n"
            "7. price, qty, sum — ЧИСЛА (не строки).\n\n"
            "ТЕКУЩАЯ СПЕЦИФИКАЦИЯ:\n{spec_current}\n\n"
            "ТЕКСТ ДОПСОГЛАШЕНИЯ:\n---\n{doc_text}\n---"
        )
        execute(
            "INSERT INTO prompts (role, name, body, is_active, notes) VALUES (%s, %s, %s, %s, %s)",
            ("diff", "default-v1", diff_body, True, "Авто-создан из llm_prompt.py _build_diff"),
        )
    _ensure_classify_prompt()


def list_by_role(role):
    """List all versions for a role, newest first."""
    rows = query(
        "SELECT id, role, name, is_active, created_by, notes, created_at FROM prompts WHERE role=%s ORDER BY created_at DESC",
        (role,),
    )
    for r in rows:
        if r.get("created_at"):
            r["created_at"] = str(r["created_at"])
    return rows


def save_new_version(role, name, body, notes="", is_active=True):
    """Save new prompt version. Deactivates all others for this role, inserts new one."""
    if is_active:
        execute("UPDATE prompts SET is_active=false WHERE role=%s", (role,))
    return execute_returning(
        """INSERT INTO prompts (role, name, body, is_active, notes)
           VALUES (%s, %s, %s, %s, %s) RETURNING *""",
        (role, name, body, is_active, notes),
    )


def activate(prompt_id):
    """Activate a prompt version (deactivates others for same role)."""
    row = query("SELECT role FROM prompts WHERE id=%s", (prompt_id,))
    if not row:
        return False
    role = row[0]["role"]
    execute("UPDATE prompts SET is_active=false WHERE role=%s", (role,))
    execute("UPDATE prompts SET is_active=true WHERE id=%s", (prompt_id,))
    return True


def delete_prompt(prompt_id):
    """Delete a prompt version (cannot delete active one)."""
    row = query("SELECT is_active FROM prompts WHERE id=%s", (prompt_id,))
    if not row:
        return False
    if row[0]["is_active"]:
        return False  # cannot delete active
    execute("DELETE FROM prompts WHERE id=%s", (prompt_id,))
    return True


def _ensure_classify_prompt():
    """Ensure classify prompt exists (idempotent)."""
    existing = query("SELECT id FROM prompts WHERE role = 'classify' AND is_active = true LIMIT 1")
    if existing:
        return
    body = (
        "Ты — система классификации договорных документов. "
        "Проанализируй текст и верни СТРОГО ВАЛИДНЫЙ JSON ОДНОЙ СТРОКОЙ "
        "(без переносов строк, без markdown, без лишних пробелов в начале/конце).\n\n"
        "Поля:\n"
        '- doc_type: "contract" (договор) / "supplement" (допсоглашение) / "specification" (спецификация/приложение) / "other"\n'
        "- own_number: номер ЭТОГО документа, строка без лишних пробелов (или null)\n"
        '- parent_number: номер родительского договора из фразы «к Договору №...» (или null)\n'
        '- doc_date: дата в YYYY-MM-DD. «27 февраля 2026» → 2026-02-27 (или null)\n'
        "- counterparty: полное название контрагента (Заказчик/Арендатор), без сокращений (или null)\n\n"
        "Пример вывода:\n"
        '{"doc_type":"supplement","own_number":"1","parent_number":"01300_2","doc_date":"2026-02-27","counterparty":"АО XXX003"}\n\n'
        "ДОКУМЕНТ:\n---\n{header_text}\n---"
    )
    execute(
        "INSERT INTO prompts (role, name, body, is_active, notes) VALUES (%s, %s, %s, %s, %s)",
        ("classify", "default-v1", body, True, "Авто-создан: classify prompt v2 — LLM сам предложил"),
    )