Spaces:

defenestrate
/

asa

Sleeping

asa / scripts /lib /chat /session_core.py

ASA Space: squashed history (instructions and settings live in the private pack)

526950e 24 days ago

10.6 kB

"""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)

	"""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)