Blame for scripts/lib/chat/session_core.py

"""Shared core curriculum: a byte-stable, nightly-built prompt block of
the highest-savings docs, sized dynamically by a savings floor.

The block sits in Anthropic's `system` field with extended-cache markers
so its first read in any 1-hour window is the only write.
"""
from __future__ import annotations
import json
import logging
import os
import statistics
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path

from scripts.lib.chat.pricing import (
    PRICE_SONNET_IN,
    PRICE_SONNET_CACHE_READ,
    PRICE_SONNET_CACHE_WRITE,
    PRICE_SONNET_CACHE_WRITE_1H,
)
from scripts.lib.chat_eval.turn_log import iter_turns
from scripts.lib.runtime_config import DERIVED_STUDENT_CHAT_DIR


logger = logging.getLogger(__name__)

_DERIVED_DIR = DERIVED_STUDENT_CHAT_DIR
_DEFAULT_JSON_PATH = _DERIVED_DIR / "session_core.json"
_MTIME_CACHE_SECONDS = 60.0


@dataclass
class SessionCore:
    doc_ids: list[str]
    frozen_prompt_text: str
    build_timestamp: str
    concurrent_sessions_estimate: int
    projected_savings_usd: float
    floor_usd: float


# path -> (loaded_at_monotonic, file_mtime, SessionCore)
_CACHE: dict[Path, tuple[float, float, SessionCore]] = {}


def load_session_core(path: Path = _DEFAULT_JSON_PATH) -> SessionCore | None:
    """Load core from disk. Returns None if the file is missing.

    Caches in-process for 60s OR until the json file's mtime changes,
    whichever comes first. The 60s ceiling means orchestrators pick up
    a nightly rebuild within a minute without explicit invalidation.
    """
    import time
    if not Path(path).is_file():
        return None
    stat = Path(path).stat()
    cached = _CACHE.get(Path(path))
    now = time.monotonic()
    if cached is not None:
        loaded_at, cached_mtime, core = cached
        if cached_mtime == stat.st_mtime and (now - loaded_at) < _MTIME_CACHE_SECONDS:
            return core

    try:
        data = json.loads(Path(path).read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError) as e:
        logger.warning("session_core.load: failed to read %s (%s)", path, e)
        return None
    prompt_path = Path(path).parent / data.get("prompt_path", "session_core_prompt.txt")
    try:
        prompt_text = prompt_path.read_text(encoding="utf-8")
    except OSError as e:
        logger.warning("session_core.load: failed to read prompt %s (%s)", prompt_path, e)
        return None

    core = SessionCore(
        doc_ids=list(data.get("doc_ids", [])),
        frozen_prompt_text=prompt_text,
        build_timestamp=data.get("build_timestamp", ""),
        concurrent_sessions_estimate=int(data.get("concurrent_sessions_estimate", 1)),
        projected_savings_usd=float(data.get("projected_savings_usd", 0.0)),
        floor_usd=float(data.get("floor_usd", 0.01)),
    )
    _CACHE[Path(path)] = (now, stat.st_mtime, core)
    return core


def _doc_savings(
    *, body_tokens: int, retrieval_count: int,
    concurrent_sessions: int, ttl_hours: float,
    write_rate: float,
) -> float:
    """Projected savings (USD) of pinning one doc into the cached prefix.

    Each TTL window: write the doc once, read it on every retrieval that
    hits during the window. Aggregate retrieval rate scales with
    concurrent_sessions because every active session may read the same
    cached prefix.

    expected_reads_per_TTL = retrieval_rate × concurrent_sessions × ttl_hours
    savings_per_TTL = body_tokens × (
        expected_reads_per_TTL × (in_rate − read_rate)
      − (write_rate − in_rate)
    )

    `retrieval_rate` is the per-session-hour retrieval count. For a 7-day
    log we approximate it as retrieval_count / (7 days × 24 hours).
    """
    in_rate = PRICE_SONNET_IN
    read_rate = PRICE_SONNET_CACHE_READ
    retrieval_rate_per_hour = retrieval_count / (7.0 * 24.0)
    expected_reads_per_TTL = (
        retrieval_rate_per_hour * concurrent_sessions * ttl_hours
    )
    return body_tokens * (
        expected_reads_per_TTL * (in_rate - read_rate)
        - (write_rate - in_rate)
    )


def _render_doc_block(rec: dict) -> str:
    """Render one doc as a deterministic XML envelope.

    Identical doc_id + kind + topic + concept_profile → identical bytes.
    Kept simpler than prompt_builder._format_doc_envelope on purpose:
    the core is *static reference material*, not session-specific
    retrieval, so we don't inline cited turns or other dynamic content.
    """
    doc_id = rec.get("doc_id", "")
    kind = rec.get("kind", "")
    topic = rec.get("topic", "")
    body = rec.get("concept_profile", "") or ""
    return (
        f'<doc id="{doc_id}" kind="{kind}" topic="{topic}">\n'
        f"  <body>{body}</body>\n"
        f"</doc>"
    )


def build_session_core(
    *,
    log_path: Path,
    output_dir: Path = _DERIVED_DIR,
    concurrent_sessions: int,
    savings_floor_usd: float = 0.01,
    ttl_hours: float = 1.0,
    catalog_records_by_id: dict,
    since: str | None = None,
) -> SessionCore:
    """Build the shared core from a turn log + catalog snapshot.

    Walks the log to count retrievals + capture body_chars per doc_id,
    projects savings under the concurrency model, then includes every
    doc whose projection clears `savings_floor_usd`. Output ordering is
    by savings descending, with doc_id tiebreaker for determinism.
    """
    write_rate = (
        PRICE_SONNET_CACHE_WRITE_1H if ttl_hours >= 1.0
        else PRICE_SONNET_CACHE_WRITE
    )

    by_doc: dict[str, dict] = {}
    for turn in iter_turns(log_path, since=since):
        for r in turn.get("retrieved", []):
            doc_id = r.get("doc_id", "")
            if not doc_id:
                continue
            entry = by_doc.setdefault(doc_id, {
                "retrieval_count": 0, "body_chars": 0,
            })
            entry["retrieval_count"] += 1
            if not entry["body_chars"]:
                entry["body_chars"] = r.get("body_chars", 0) or 0

    candidates: list[tuple[float, str, dict]] = []
    total_savings = 0.0
    for doc_id, info in by_doc.items():
        rec = catalog_records_by_id.get(doc_id)
        if rec is None:
            continue  # stale catalog: doc was removed since log was written
        body_chars = info["body_chars"] or len(rec.get("concept_profile", "") or "")
        if body_chars <= 0 or info["retrieval_count"] < 2:
            continue
        body_tokens = body_chars // 4
        savings = _doc_savings(
            body_tokens=body_tokens,
            retrieval_count=info["retrieval_count"],
            concurrent_sessions=concurrent_sessions,
            ttl_hours=ttl_hours,
            write_rate=write_rate,
        )
        if savings <= 0 or savings < savings_floor_usd:
            continue
        candidates.append((savings, doc_id, rec))
        total_savings += savings

    # Deterministic order: highest savings first, doc_id ascending tiebreaker.
    candidates.sort(key=lambda x: (-x[0], x[1]))

    # Always-pin: course_meta records (the syllabus, today). These bypass
    # the savings filter — pinning the syllabus is a policy choice, not an
    # emergent metric. Without this, the syllabus could fall out of the core
    # whenever its retrieval count drops (which is exactly what happens once
    # it's pinned, since students no longer need to ask for it explicitly).
    existing_ids = {doc_id for _, doc_id, _ in candidates}
    must_pin: list[tuple[float, str, dict]] = []
    for doc_id, rec in catalog_records_by_id.items():
        if rec.get("kind") != "course_meta":
            continue
        if doc_id in existing_ids:
            continue
        must_pin.append((0.0, doc_id, rec))
    must_pin.sort(key=lambda x: x[1])
    candidates = must_pin + candidates

    doc_ids = [doc_id for _, doc_id, _ in candidates]
    rendered_blocks = [_render_doc_block(rec) for _, _, rec in candidates]
    prompt_text = "\n\n".join(rendered_blocks)

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    prompt_path = output_dir / "session_core_prompt.txt"
    prompt_staging = prompt_path.with_name(prompt_path.name + ".tmp")
    prompt_staging.write_text(prompt_text, encoding="utf-8")

    from datetime import datetime, timezone
    metadata = {
        "doc_ids": doc_ids,
        "build_timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "concurrent_sessions_estimate": concurrent_sessions,
        "projected_savings_usd": round(total_savings, 4),
        "floor_usd": savings_floor_usd,
        "prompt_path": prompt_path.name,
    }
    json_path = output_dir / "session_core.json"
    json_staging = json_path.with_name(json_path.name + ".tmp")
    json_staging.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

    # Both staging files written; rename them as close together as possible.
    # If the process is killed between the two renames, the loader will read
    # the old json (still pointing at the old prompt path) — the new prompt
    # is left at session_core_prompt.txt, but the old json keeps consumers on
    # the previous build until next nightly run.
    os.replace(prompt_staging, prompt_path)
    os.replace(json_staging, json_path)

    return SessionCore(
        doc_ids=doc_ids,
        frozen_prompt_text=prompt_text,
        build_timestamp=metadata["build_timestamp"],
        concurrent_sessions_estimate=concurrent_sessions,
        projected_savings_usd=metadata["projected_savings_usd"],
        floor_usd=savings_floor_usd,
    )


def estimate_concurrent_sessions(
    log_path: Path,
    *,
    since: str | None = None,
) -> int:
    """Estimate steady-state concurrent-session count via p50 of distinct
    sessions per active hour.

    'Active hour' = any UTC hour with at least one logged turn. Hours
    with zero turns don't count (they would drag the median to 1
    artificially during off-peak).

    Returns max(1, p50). Used by the builder when concurrent_sessions
    isn't explicitly passed.
    """
    by_hour: dict[str, set[str]] = defaultdict(set)
    for turn in iter_turns(log_path, since=since):
        ts = turn.get("ts", "")
        sid = turn.get("session_id", "")
        if not ts or not sid:
            continue
        # ts format: 2026-05-01T10:05:00Z → hour key: 2026-05-01T10
        hour_key = ts[:13]
        by_hour[hour_key].add(sid)

    if not by_hour:
        return 1
    counts = sorted(len(s) for s in by_hour.values())
    p50 = int(statistics.median(counts))
    return max(1, p50)