Kairo-Jinkendo/backend/prompt_registry.py

"""Prompt definition registry with versioning and optional steps."""

from __future__ import annotations

import json
from dataclasses import dataclass

from db import get_connection
from feature_registry import VALID_CONTEXT_KINDS, VALID_EXECUTION_MODES

VALID_STEP_TYPES = frozenset(
    {"template", "resolver", "llm_call", "postprocess", "validation"}
)


@dataclass(frozen=True)
class PromptStepRegistration:
    step_key: str
    step_order: int
    step_type: str
    template_body: str
    input_mapping: dict | None = None
    output_mapping: dict | None = None
    enabled: bool = True


@dataclass(frozen=True)
class PromptVersionRegistration:
    version: str
    body: str
    input_schema: dict | None = None
    output_schema: dict | None = None
    model_policy_key: str | None = None
    changelog: str = ""
    steps: tuple[PromptStepRegistration, ...] = ()


@dataclass(frozen=True)
class PromptRegistration:
    prompt_key: str
    name: str
    purpose: str
    context_kind: str
    execution_mode: str
    active_version: str
    versions: tuple[PromptVersionRegistration, ...]


_REGISTRY: dict[str, PromptRegistration] = {}


def register_prompt(registration: PromptRegistration) -> None:
    if not registration.prompt_key or not registration.prompt_key.strip():
        raise ValueError("Prompt key is required")
    if registration.context_kind not in VALID_CONTEXT_KINDS:
        raise ValueError(f"Unknown context_kind: {registration.context_kind}")
    if registration.execution_mode not in VALID_EXECUTION_MODES:
        raise ValueError(f"Unknown execution_mode: {registration.execution_mode}")
    if not registration.versions:
        raise ValueError("At least one prompt version is required")
    version_ids = {v.version for v in registration.versions}
    if registration.active_version not in version_ids:
        raise ValueError(f"active_version not found in versions: {registration.active_version}")
    for step in _iter_steps(registration):
        if step.step_type not in VALID_STEP_TYPES:
            raise ValueError(f"Unknown step_type: {step.step_type}")
    if registration.prompt_key in _REGISTRY:
        raise ValueError(f"Prompt already registered: {registration.prompt_key}")
    _REGISTRY[registration.prompt_key] = registration


def _iter_steps(registration: PromptRegistration):
    for version in registration.versions:
        yield from version.steps


def get_registered_prompts() -> tuple[PromptRegistration, ...]:
    return tuple(_REGISTRY.values())


def clear_prompt_registry_for_tests() -> None:
    _REGISTRY.clear()


def sync_prompts_to_db() -> int:
    if not _REGISTRY:
        print("[prompt_registry] Keine Prompts registriert — Sync übersprungen")
        return 0

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            for reg in _REGISTRY.values():
                cur.execute(
                    """
                    INSERT INTO prompt_definitions (
                        prompt_key, name, purpose, context_kind, execution_mode, status
                    )
                    VALUES (%s, %s, %s, %s, %s, 'active')
                    ON CONFLICT (prompt_key) DO UPDATE SET
                        name = EXCLUDED.name,
                        purpose = EXCLUDED.purpose,
                        context_kind = EXCLUDED.context_kind,
                        execution_mode = EXCLUDED.execution_mode,
                        status = 'active',
                        updated_at = NOW()
                    RETURNING id
                    """,
                    (
                        reg.prompt_key,
                        reg.name,
                        reg.purpose,
                        reg.context_kind,
                        reg.execution_mode,
                    ),
                )
                definition_id = cur.fetchone()[0]

                active_version_id = None
                for version in reg.versions:
                    cur.execute(
                        """
                        INSERT INTO prompt_versions (
                            prompt_definition_id, version, body,
                            input_schema, output_schema, model_policy_key, changelog
                        )
                        VALUES (%s, %s, %s, %s::jsonb, %s::jsonb, %s, %s)
                        ON CONFLICT (prompt_definition_id, version) DO UPDATE SET
                            body = EXCLUDED.body,
                            input_schema = EXCLUDED.input_schema,
                            output_schema = EXCLUDED.output_schema,
                            model_policy_key = EXCLUDED.model_policy_key,
                            changelog = EXCLUDED.changelog
                        RETURNING id
                        """,
                        (
                            definition_id,
                            version.version,
                            version.body,
                            json.dumps(version.input_schema) if version.input_schema else None,
                            json.dumps(version.output_schema) if version.output_schema else None,
                            version.model_policy_key,
                            version.changelog,
                        ),
                    )
                    version_id = cur.fetchone()[0]
                    if version.version == reg.active_version:
                        active_version_id = version_id

                    for step in version.steps:
                        cur.execute(
                            """
                            INSERT INTO prompt_steps (
                                prompt_version_id, step_key, step_order, step_type,
                                template_body, input_mapping, output_mapping, enabled
                            )
                            VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s)
                            ON CONFLICT (prompt_version_id, step_key) DO UPDATE SET
                                step_order = EXCLUDED.step_order,
                                step_type = EXCLUDED.step_type,
                                template_body = EXCLUDED.template_body,
                                input_mapping = EXCLUDED.input_mapping,
                                output_mapping = EXCLUDED.output_mapping,
                                enabled = EXCLUDED.enabled,
                                updated_at = NOW()
                            """,
                            (
                                version_id,
                                step.step_key,
                                step.step_order,
                                step.step_type,
                                step.template_body,
                                json.dumps(step.input_mapping) if step.input_mapping else None,
                                json.dumps(step.output_mapping) if step.output_mapping else None,
                                step.enabled,
                            ),
                        )

                if active_version_id:
                    cur.execute(
                        """
                        UPDATE prompt_definitions
                        SET active_version_id = %s, updated_at = NOW()
                        WHERE id = %s
                        """,
                        (active_version_id, definition_id),
                    )

        conn.commit()
        print(f"[prompt_registry] Sync OK — {len(_REGISTRY)} Prompt(s)")
        return 0
    except Exception as exc:
        conn.rollback()
        print(f"[prompt_registry] Sync FAIL: {exc}")
        return 1
    finally:
        conn.close()


def load_prompt_definition(prompt_key: str) -> dict | None:
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, prompt_key, name, purpose, context_kind, execution_mode,
                       status, active_version_id
                FROM prompt_definitions
                WHERE prompt_key = %s
                """,
                (prompt_key,),
            )
            row = cur.fetchone()
            if not row:
                return None
            return {
                "id": str(row[0]),
                "prompt_key": row[1],
                "name": row[2],
                "purpose": row[3],
                "context_kind": row[4],
                "execution_mode": row[5],
                "status": row[6],
                "active_version_id": str(row[7]) if row[7] else None,
            }
    finally:
        conn.close()


def load_active_prompt_version(definition: dict) -> dict | None:
    version_id = definition.get("active_version_id")
    if not version_id:
        return None
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, version, body, input_schema, output_schema, model_policy_key
                FROM prompt_versions
                WHERE id = %s
                """,
                (version_id,),
            )
            row = cur.fetchone()
            if not row:
                return None
            return {
                "id": str(row[0]),
                "version": row[1],
                "body": row[2],
                "input_schema": row[3],
                "output_schema": row[4],
                "model_policy_key": row[5],
            }
    finally:
        conn.close()


def load_template_steps(version_id: str) -> list[dict]:
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT step_key, step_order, step_type, template_body, enabled
                FROM prompt_steps
                WHERE prompt_version_id = %s AND enabled = TRUE
                ORDER BY step_order, step_key
                """,
                (version_id,),
            )
            return [
                {
                    "step_key": row[0],
                    "step_order": row[1],
                    "step_type": row[2],
                    "template_body": row[3],
                    "enabled": row[4],
                }
                for row in cur.fetchall()
            ]
    finally:
        conn.close()


def list_prompt_definitions() -> list[dict]:
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT prompt_key, name, purpose, context_kind, execution_mode, status
                FROM prompt_definitions
                ORDER BY prompt_key
                """
            )
            return [
                {
                    "prompt_key": row[0],
                    "name": row[1],
                    "purpose": row[2],
                    "context_kind": row[3],
                    "execution_mode": row[4],
                    "status": row[5],
                }
                for row in cur.fetchall()
            ]
    finally:
        conn.close()