app/llm/prompt_cache.py

"""OpenAI prompt-cache helpers: stable system prefix, cache keys, usage metrics."""

from __future__ import annotations

import hashlib
from typing import Any

import structlog
import tiktoken

from app.config import settings

_log = structlog.get_logger(__name__)
_ENCODING = tiktoken.get_encoding("cl100k_base")

# Static block appended when the system prompt is below OpenAI automatic cache threshold (~1024 tokens).
_CACHE_PAD_HEADER = (
    "\n\n--- RICS REPORT SCHEMA REFERENCE (STATIC — DO NOT EDIT) ---\n"
    "This block exists only to stabilise the cached system prefix for inference. "
    "Section codes follow RICS Home Survey Standard layout (A–L, D1–D9, E1–E9, etc.). "
    "Condition ratings: NI, 1, 2, 3. Survey levels: 1 Condition Report, 2 Home Survey, 3 Building Survey. "
    "British English only. Non-invention: property-specific facts originate from inspector notes only. "
)


def prompt_caching_active() -> bool:
    return bool(settings.enable_async_pipeline and settings.enable_prompt_caching)


def count_prompt_tokens(text: str) -> int:
    return len(_ENCODING.encode(text or ""))


def ensure_cacheable_system_prefix(system: str, *, min_tokens: int | None = None) -> str:
    """Pad the system prompt so identical prefixes exceed OpenAI automatic cache minimum."""
    floor = int(min_tokens or settings.prompt_cache_min_system_tokens)
    body = (system or "").strip()
    if count_prompt_tokens(body) >= floor:
        return body
    pad = _CACHE_PAD_HEADER
    # Grow pad deterministically until token floor is met.
    while count_prompt_tokens(body + pad) < floor:
        pad += (
            " Standard paragraphs map notes to HB-BS clauses. "
            "Retrieval tiers: document, section, paragraph. "
            "Evidence must be traceable to notes or approved standard wording. "
        )
    return body + pad


def normalize_messages_for_caching(messages: list[Any]) -> list[Any]:
    """Pad the first system message so multi-turn and raw message lists cache reliably."""
    if not prompt_caching_active() or not messages:
        return messages
    out: list[Any] = []
    padded_system = False
    for msg in messages:
        if not isinstance(msg, dict):
            out.append(msg)
            continue
        role = msg.get("role")
        content = msg.get("content")
        if (
            not padded_system
            and role == "system"
            and isinstance(content, str)
        ):
            out.append({**msg, "content": ensure_cacheable_system_prefix(content)})
            padded_system = True
        else:
            out.append(dict(msg))
    return out


def build_chat_messages(*, system: str, user: str) -> list[dict[str, str]]:
    """System (static) first, user (dynamic) second — maximises prefix cache hits."""
    sys = ensure_cacheable_system_prefix(system) if prompt_caching_active() else system
    return [
        {"role": "system", "content": sys},
        {"role": "user", "content": user},
    ]


def prompt_cache_key(
    *,
    phase: str,
    model: str,
    survey_level: int | None = None,
    interference_level: str | None = None,
    tenant_id: str | None = None,
) -> str:
    il = (interference_level or "").strip().lower() or "legacy"
    lvl = int(survey_level or 0)
    tenant = (tenant_id or "").strip() or "_"
    raw = f"rics|{tenant}|{model}|{phase}|L{lvl}|{il}"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]


def openai_extra_kwargs(
    *,
    phase: str,
    model: str,
    survey_level: int | None = None,
    interference_level: str | None = None,
    tenant_id: str | None = None,
) -> dict[str, Any]:
    if not prompt_caching_active():
        return {}
    return {"prompt_cache_key": prompt_cache_key(
        phase=phase,
        model=model,
        survey_level=survey_level,
        interference_level=interference_level,
        tenant_id=tenant_id,
    )}


def log_openai_cache_usage(
    response: Any,
    *,
    phase: str,
    section_id: str | None,
) -> bool | None:
    """Log cache_read_input_tokens / input_tokens; return True if cache was used."""
    usage = getattr(response, "usage", None)
    if usage is None:
        return None
    input_tokens = int(getattr(usage, "prompt_tokens", 0) or getattr(usage, "input_tokens", 0) or 0)
    cached = int(getattr(usage, "prompt_tokens_details", None) and getattr(
        usage.prompt_tokens_details, "cached_tokens", 0
    ) or getattr(usage, "cache_read_input_tokens", 0) or 0)
    # OpenAI SDK: cached_tokens may live on prompt_tokens_details
    details = getattr(usage, "prompt_tokens_details", None)
    if details is not None:
        cached = int(getattr(details, "cached_tokens", 0) or 0)
    rate = (cached / input_tokens) if input_tokens > 0 else 0.0
    cache_hit = cached > 0
    _log.info(
        event="openai_prompt_cache",
        phase=phase,
        section_id=section_id,
        cache_hit=cache_hit,
        cache_read_input_tokens=cached,
        input_tokens=input_tokens,
        cache_hit_rate=round(rate, 4),
    )
    return cache_hit