Prompt Management System

-- Prompt registry schema
CREATE TYPE prompt_status AS ENUM (
    'draft', 'review', 'approved', 'deployed', 'archived'
);

CREATE TABLE prompts (
    id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    slug        TEXT UNIQUE NOT NULL,  -- e.g., 'customer-support-v2'
    name        TEXT NOT NULL,
    description TEXT,
    created_by  TEXT NOT NULL,
    created_at  TIMESTAMPTZ DEFAULT now()
);

CREATE TABLE prompt_versions (
    id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    prompt_id   UUID REFERENCES prompts(id) ON DELETE CASCADE,
    version     INTEGER NOT NULL,
    template    TEXT NOT NULL,          -- Prompt text with {{variable}} placeholders
    variables   JSONB DEFAULT '[]',    -- Expected variables and their types
    model_hint  TEXT,                   -- Recommended model
    status      prompt_status DEFAULT 'draft',
    change_note TEXT,
    created_by  TEXT NOT NULL,
    reviewed_by TEXT,
    deployed_at TIMESTAMPTZ,
    created_at  TIMESTAMPTZ DEFAULT now(),
    UNIQUE (prompt_id, version)
);

-- Only one deployed version per prompt
CREATE UNIQUE INDEX idx_one_deployed
    ON prompt_versions (prompt_id)
    WHERE status = 'deployed';

CREATE TABLE prompt_ab_tests (
    id            UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    prompt_id     UUID REFERENCES prompts(id),
    control_id    UUID REFERENCES prompt_versions(id),
    variant_id    UUID REFERENCES prompt_versions(id),
    traffic_pct   INTEGER DEFAULT 10,  -- % to variant
    started_at    TIMESTAMPTZ DEFAULT now(),
    ended_at      TIMESTAMPTZ,
    winner_id     UUID REFERENCES prompt_versions(id)
);

"""Prompt management SDK for runtime prompt fetching."""

from __future__ import annotations

import hashlib
import json
import re
from typing import Any

import redis.asyncio as redis


class PromptClient:
    """SDK for fetching and rendering managed prompts."""

    def __init__(self, redis_url: str, db_fallback=None):
        self._redis = redis.from_url(redis_url)
        self._db = db_fallback
        self._local_cache: dict[str, dict] = {}

    async def get_prompt(
        self,
        slug: str,
        variables: dict[str, Any] | None = None,
        ab_identifier: str | None = None,
    ) -> str:
        """Fetch and render a prompt by slug.

        Args:
            slug: The prompt slug (e.g., 'customer-support').
            variables: Template variables to render.
            ab_identifier: User/session ID for A/B test assignment.

        Returns:
            The rendered prompt string.
        """
        prompt_data = await self._fetch(slug)

        # Check for active A/B test
        if ab_identifier and prompt_data.get("ab_test"):
            ab = prompt_data["ab_test"]
            bucket = self._hash_to_bucket(ab_identifier)
            if bucket < ab["traffic_pct"]:
                template = ab["variant_template"]
            else:
                template = prompt_data["template"]
        else:
            template = prompt_data["template"]

        return self._render(template, variables or {})

    async def _fetch(self, slug: str) -> dict:
        """Fetch prompt data from cache with DB fallback."""
        # L1: in-process cache
        if slug in self._local_cache:
            return self._local_cache[slug]

        # L2: Redis cache
        cached = await self._redis.get(f"prompt:{slug}")
        if cached:
            data = json.loads(cached)
            self._local_cache[slug] = data
            return data

        # L3: Database fallback
        if self._db:
            data = await self._db.fetch_deployed_prompt(slug)
            await self._redis.setex(
                f"prompt:{slug}", 300, json.dumps(data)
            )
            self._local_cache[slug] = data
            return data

        raise ValueError(f"Prompt '{slug}' not found")

    @staticmethod
    def _hash_to_bucket(identifier: str) -> int:
        """Deterministic hash to 0-99 for A/B assignment."""
        h = hashlib.sha256(identifier.encode()).hexdigest()
        return int(h[:8], 16) % 100

    @staticmethod
    def _render(template: str, variables: dict) -> str:
        """Render template variables like {{name}}."""
        def replacer(match):
            key = match.group(1).strip()
            return str(variables.get(key, match.group(0)))
        return re.sub(r"\{\{(.+?)\}\}", replacer, template)

# Prompt definition file for CI/CD pipeline
slug: customer-support
name: Customer Support Agent
description: System prompt for the customer-facing support assistant
model_hint: claude-sonnet-4-20250514

template: |
  You are a helpful customer support agent for {{company_name}}.

  Guidelines:
  - Be empathetic and professional
  - Answer questions using only the provided context
  - If you do not know the answer, say so honestly
  - Never make up product features or policies
  - Escalate billing and account deletion requests to a human agent

  Customer tier: {{customer_tier}}
  Context from knowledge base:
  {{retrieved_context}}

variables:
  - name: company_name
    type: string
    required: true
  - name: customer_tier
    type: string
    required: true
    enum: [free, starter, pro, enterprise]
  - name: retrieved_context
    type: string
    required: true