Operating AI Systems

# script_id: day_088_operating_ai_systems/first_kill_switch
# A kill switch every replica reads: one shared hash plus a short local cache.
# Install: none (stdlib). In production, replace MemoryHash() with redis.Redis.from_url(url, decode_responses=True).
import time


class MemoryHash(dict):                       # stands in for Redis: the same hset/hgetall calls
    def hset(self, name: str, key: str, value: str) -> None: self.setdefault(name, {})[key] = value
    def hgetall(self, name: str) -> dict[str, str]: return dict(self.get(name, {}))


store, _cache = MemoryHash(), {"at": float("-inf"), "value": {}}


def overrides(max_age_s: float = 5.0) -> dict[str, str]:
    if time.monotonic() - _cache["at"] >= max_age_s:     # at most one store read per replica every 5 s
        _cache.update(at=time.monotonic(), value=store.hgetall("ops:overrides"))
    return _cache["value"]


store.hset("ops:overrides", "degraded", "retrieval_only")
store.hset("ops:overrides", "degraded_reason", "provider 503s")
print(f"degraded={overrides()['degraded']} (set by you: {overrides()['degraded_reason']!r})")

degraded=retrieval_only (set by you: 'provider 503s')

# script_id: day_088_operating_ai_systems/atlas_ops
# atlas/ops.py (blocks 1-4) + scripts/atlas_ops.py (blocks 5-6): every operational lever, audited.
# Install: uv add redis     Env: REDIS_URL (unset = in-process store, for this demo), ATLAS_OPERATOR
# Services: in real use, Day 78's Redis. Every API replica and every operator must use the same one.
from __future__ import annotations

import argparse
import json
import os
import shlex
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Literal

OVERRIDES_KEY, AUDIT_KEY = "ops:overrides", "ops:audit"
Degraded = Literal["off", "retrieval_only", "queue"]
ALLOWED = {"degraded": {"off", "retrieval_only", "queue"}, "force_fallback": {"1"}, "approval_only": {"1"}}


@dataclass(frozen=True)
class Overrides:
    degraded: Degraded = "off"                      # retrieval_only = cited links, no generation (Day 84)
    force_fallback: bool = False                    # skip the primary model even if its breaker is closed
    disabled_routes: frozenset[str] = frozenset()   # e.g. {"investigate"}
    index_version: str | None = None                # None = whatever the index DB marks active
    tenant_caps: dict[str, int] = field(default_factory=dict)   # temporary tokens-per-hour caps
    approval_only: bool = False                     # every tool side effect waits for a human

    @classmethod
    def from_hash(cls, h: dict[str, str]) -> Overrides:
        caps = {k.removeprefix("tenant_cap:"): int(v) for k, v in h.items() if k.startswith("tenant_cap:")}
        degraded = h.get("degraded", "off")
        return cls(degraded=degraded if degraded in ALLOWED["degraded"] else "off",
                   force_fallback=h.get("force_fallback") == "1",
                   disabled_routes=frozenset(r for r in h.get("disabled_routes", "").split(",") if r),
                   index_version=h.get("index_version") or None,
                   tenant_caps=caps, approval_only=h.get("approval_only") == "1")

# script_id: day_088_operating_ai_systems/atlas_ops
class MemoryRedis:
    """In-process stand-in for the six redis-py calls used here (tests, laptops, this demo)."""

    def __init__(self) -> None:
        self.hashes: dict[str, dict[str, str]] = {}
        self.lists: dict[str, list[str]] = {}

    def hset(self, name: str, key: str, value: str) -> int:
        self.hashes.setdefault(name, {})[key] = value
        return 1

    def hdel(self, name: str, *keys: str) -> int:
        return sum(self.hashes.get(name, {}).pop(k, None) is not None for k in keys)

    def hgetall(self, name: str) -> dict[str, str]:
        return dict(self.hashes.get(name, {}))

    def lpush(self, name: str, *values: str) -> int:
        self.lists[name] = [*reversed(values), *self.lists.get(name, [])]
        return len(self.lists[name])

    def lrange(self, name: str, start: int, end: int) -> list[str]:    # end is inclusive; -1 = last
        return self.lists.get(name, [])[start: None if end == -1 else end + 1]

    def ltrim(self, name: str, start: int, end: int) -> bool:
        self.lists[name] = self.lrange(name, start, end)
        return True


_MEMORY = MemoryRedis()


def get_redis(url: str | None = None):
    """REDIS_URL unset → one shared in-process store; otherwise real Redis returning str values."""
    url = url or os.environ.get("REDIS_URL", "")
    if not url:
        return _MEMORY
    from redis import Redis                         # lazy import: the demo needs no server
    return Redis.from_url(url, decode_responses=True)

# script_id: day_088_operating_ai_systems/atlas_ops
_cache: dict[int, tuple[float, Overrides]] = {}


def get_overrides(redis, *, max_age_s: float = 5) -> Overrides:
    """Read shared overrides at most once per max_age_s per process; keep the last value if Redis fails."""
    at, value = _cache.get(id(redis), (float("-inf"), Overrides()))
    if time.monotonic() - at >= max_age_s:
        try:
            value = Overrides.from_hash(redis.hgetall(OVERRIDES_KEY))
        except Exception as exc:                    # an ops-store outage must not become an API outage
            print(f"warning: overrides unavailable ({exc!r}); using last known", file=sys.stderr)
        _cache[id(redis)] = (time.monotonic(), value)
    return value


def set_override(redis, key: str, value: str, *, actor: str, reason: str) -> None:
    """Change one lever and append who/what/when/why to ops:audit. An empty value clears the lever."""
    if not reason.strip():
        raise ValueError("every change needs --reason; it goes into the audit log and the postmortem")
    if value and key in ALLOWED and value not in ALLOWED[key]:
        raise ValueError(f"{key} must be one of {sorted(ALLOWED[key])}, got {value!r}")
    if value and key.startswith("tenant_cap:") and not value.isdigit():
        raise ValueError(f"{key} must be a whole number of tokens per hour, got {value!r}")
    if value:
        redis.hset(OVERRIDES_KEY, key, value)
    else:
        redis.hdel(OVERRIDES_KEY, key)
    entry = {"at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
             "actor": actor, "key": key, "value": value, "reason": reason}
    redis.lpush(AUDIT_KEY, json.dumps(entry))
    redis.ltrim(AUDIT_KEY, 0, 999)                  # keep the last 1,000 here; ship older ones to your logs

# script_id: day_088_operating_ai_systems/atlas_ops
@dataclass(frozen=True)
class RequestPlan:
    status: int                       # 200 serve · 202 queued as a job · 429 capped · 503 route paused
    mode: str                         # "full" | "retrieval_only" | "queued" | "refused"
    force_fallback: bool = False
    index_version: str | None = None
    approval_only: bool = False
    message: str = ""


def plan_request(route: str, tenant: str, ov: Overrides, *, tokens_this_hour: int = 0) -> RequestPlan:
    """The first call in every API handler. Pure, so each kill switch has a unit test."""
    if route in ov.disabled_routes:
        return RequestPlan(503, "refused", message=f"/{route} is paused by the on-call team; retry later")
    cap = ov.tenant_caps.get(tenant)
    if cap is not None and tokens_this_hour >= cap:
        return RequestPlan(429, "refused", message=f"{tenant} reached its temporary cap of {cap:,} tokens/hour")
    if ov.degraded == "queue":
        return RequestPlan(202, "queued", message="accepted; the answer arrives as a background job (Day 78)")
    mode = "retrieval_only" if ov.degraded == "retrieval_only" else "full"
    return RequestPlan(200, mode, ov.force_fallback, ov.index_version, ov.approval_only)

# Wiring into Day 77's FastAPI handler, as the first lines of POST /v1/ask:
#   plan = plan_request("ask", principal.tenant, get_overrides(redis), tokens_this_hour=used)
#   if plan.status >= 400: raise HTTPException(plan.status, plan.message)
#   then pass plan.mode / plan.force_fallback / plan.index_version down to ask()

# script_id: day_088_operating_ai_systems/atlas_ops
def build_parser() -> argparse.ArgumentParser:
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--reason", default="", help="why; required for every change")
    p = argparse.ArgumentParser(prog="atlas_ops", description="Atlas kill switches (audited)")
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("status")
    sub.add_parser("degrade", parents=[common]).add_argument("mode", choices=["on", "queue", "off"])
    sub.add_parser("fallback", parents=[common]).add_argument("mode", choices=["force", "auto"])
    sub.add_parser("approvals", parents=[common]).add_argument("mode", choices=["on", "off"])
    r = sub.add_parser("route", parents=[common])
    r.add_argument("action", choices=["disable", "enable"])
    r.add_argument("name")
    t = sub.add_parser("tenant", parents=[common])
    t.add_argument("action", choices=["cap"])
    t.add_argument("tenant")
    t.add_argument("tokens", help="tokens per hour, or 'clear'")
    i = sub.add_parser("index", parents=[common])
    i.add_argument("action", choices=["use"])
    i.add_argument("name", help="v1, v2 … or 'auto' to follow the index DB")
    return p

# script_id: day_088_operating_ai_systems/atlas_ops
def main(argv: list[str] | None = None) -> None:
    a = build_parser().parse_args(argv)
    redis, actor = get_redis(), os.environ.get("ATLAS_OPERATOR") or os.environ.get("USER", "unknown")
    ov = get_overrides(redis, max_age_s=0)                   # operators always want the live value
    if a.cmd == "status":
        print(json.dumps({**ov.__dict__, "disabled_routes": sorted(ov.disabled_routes)}))
        return
    if a.cmd == "degrade":
        key, value = "degraded", {"on": "retrieval_only", "queue": "queue", "off": ""}[a.mode]
    elif a.cmd in ("fallback", "approvals"):
        key = "force_fallback" if a.cmd == "fallback" else "approval_only"
        value = "1" if a.mode in ("force", "on") else ""
    elif a.cmd == "route":
        on = set(ov.disabled_routes)
        routes = on | {a.name} if a.action == "disable" else on - {a.name}
        key, value = "disabled_routes", ",".join(sorted(routes))
    elif a.cmd == "tenant":
        key, value = f"tenant_cap:{a.tenant}", "" if a.tokens == "clear" else a.tokens
    else:
        key, value = "index_version", "" if a.name == "auto" else a.name
    set_override(redis, key, value, actor=actor, reason=a.reason)
    print(f"OK {key}={value or '(cleared)'} by {actor}: {a.reason}")

# script_id: day_088_operating_ai_systems/atlas_ops
DEMO = ["degrade on --reason 'provider 503s, breaker open (INC-42)'",
        "route disable investigate --reason 'cut load while degraded'",
        "tenant cap acme 50000 --reason 'acme investigate loop burning tokens'",
        "degrade off", "status"]

if __name__ == "__main__":
    if len(sys.argv) > 1:
        main()
    else:                                                    # demo transcript on the in-process store
        os.environ.setdefault("ATLAS_OPERATOR", "oncall-you")
        for cmd in DEMO:
            print(f"$ atlas_ops {cmd}")
            try:
                main(shlex.split(cmd))
            except ValueError as exc:
                print(f"refused: {exc}")
        ov = get_overrides(get_redis())                      # what any replica sees within 5 s
        print("replica:", plan_request("investigate", "acme", ov).message)
        print("replica:", plan_request("ask", "acme", ov, tokens_this_hour=51_000).message)

$ atlas_ops degrade on --reason 'provider 503s, breaker open (INC-42)'
OK degraded=retrieval_only by oncall-you: provider 503s, breaker open (INC-42)
$ atlas_ops route disable investigate --reason 'cut load while degraded'
OK disabled_routes=investigate by oncall-you: cut load while degraded
$ atlas_ops tenant cap acme 50000 --reason 'acme investigate loop burning tokens'
OK tenant_cap:acme=50000 by oncall-you: acme investigate loop burning tokens
$ atlas_ops degrade off
refused: every change needs --reason; it goes into the audit log and the postmortem
$ atlas_ops status
{"degraded": "retrieval_only", "force_fallback": false, "disabled_routes": ["investigate"], "index_version": null, "tenant_caps": {"acme": 50000}, "approval_only": false}
replica: /investigate is paused by the on-call team; retry later
replica: acme reached its temporary cap of 50,000 tokens/hour

<!-- file: ops/runbooks/provider_outage.md -->
# Runbook: LLM provider outage

## Summary
The primary model provider returns 5xx errors or times out. Atlas's circuit breaker (Day 84) opens and
traffic moves to the fallback model, but retries and half-open probes still fail a share of requests.

## Severity
- SEV1: error rate above 5% for 5 minutes on all tenants, or the fallback model is failing too.
- SEV2: fallback share above 20% with error rate below 5%.

## Signals
- Alert `error_rate_high` (this runbook) and the "calls by model and status" panel from Day 86.
- The provider's status page. Check it, but don't wait for it to confirm.

## First 5 minutes
1. Acknowledge the page. Post "Investigating Atlas errors" in #atlas-incidents.
2. Run `atlas_ops status` and note any override that is already set.
3. Open the "calls by model and status" panel: errors only on the primary model = provider outage.

## Mitigate
1. `atlas_ops fallback force --reason "primary 5xx, INC-<n>"`
2. Fallback failing as well: `atlas_ops degrade on --reason "both models failing, INC-<n>"`
3. Queue depth or latency still climbing: `atlas_ops route disable investigate --reason "shed load, INC-<n>"`

## Verify
Within 60 s: `error_rate_high` clears, fallback share is near 100%, and 3 manual `/v1/ask` calls succeed.

## Escalate
SEV1, or no improvement 15 minutes after step 1: page the incident commander on the rota.

## Status message
"Atlas is answering with a backup model while our AI provider has an outage. Answers may be slower
or shorter. Next update in 30 minutes."

## Follow-ups
- Revert with `atlas_ops fallback auto --reason "provider green 30 min"` once the provider has been healthy for 30 minutes.
- Postmortem within 2 working days. Fix any step in this runbook that slowed you down.

<!-- file: ops/runbooks/quality_regression.md -->
# Runbook: quality regression

## Summary
Requests succeed but answers got worse: wrong version, missing citations, refusals. Usual causes: an index
migration, a prompt or model flag (Day 87), a docs change, or a vector store returning nothing.

## Severity
- SEV1: online pass rate down 25+ points, or wrong answers on security or data-loss topics.
- SEV2: online pass rate down 10+ points on n >= 100 sampled requests.

## Signals
- Alert `pass_rate_low`. The "zero-chunk requests" panel: a jump means retrieval is silently empty.
- Thumbs-down rate (Day 86 feedback) and the failing samples linked from the alert.

## First 5 minutes
1. Acknowledge, post in #atlas-incidents.
2. `atlas_ops status`, then list the last day's changes: `git log --since=24.hours -- release.toml prompts/`.
3. Read the index events: `sqlite3 atlas_index.db "SELECT * FROM index_events ORDER BY at DESC LIMIT 5"`.

## Mitigate
1. Index changed recently: `atlas_ops index use v1 --reason "pass rate -15 after cutover, INC-<n>"`
2. Prompt or model flag changed: roll that flag back in `release.toml` (Day 87), then redeploy the config.

## Verify
Online pass rate within 5 points of baseline over the next 100 sampled requests; zero-chunk share back to normal.

## Escalate
No known change in the last 24 h, or no recovery after both levers: page the RAG owner and the IC.

## Status message
"Some Atlas answers were less accurate than usual since <time>. We have rolled back a recent change
and are verifying. Next update in 30 minutes."

## Follow-ups
- Add the failing questions to the eval set (Day 86 `feedback_to_cases`).
- Make the missing check a gate, so the same change can't ship again unnoticed.

# script_id: day_088_operating_ai_systems/runbook_lint
# scripts/runbook_lint.py: required sections present, no vague steps. Run it in CI on ops/runbooks/.
# Install: none (stdlib). Usage: python runbook_lint.py ops/runbooks/*.md   (no args = demo)
import re
import sys
from pathlib import Path

REQUIRED = ["## Summary", "## Severity", "## Signals", "## First 5 minutes", "## Mitigate",
            "## Verify", "## Escalate", "## Status message", "## Follow-ups"]
# a vague verb + object, unless the same line names the exact command or query in backticks
VAGUE = re.compile(r"\b(check|look at|investigate|monitor) (the )?(logs|dashboards?|metrics)\b(?!.*`)", re.I)


def lint(text: str, *, sections: bool = True) -> list[str]:
    problems = [f"missing section {s!r}" for s in REQUIRED if sections and s not in text]
    for n, line in enumerate(text.splitlines(), 1):
        if VAGUE.search(line):
            problems.append(f"line {n}: vague step {line.strip()!r}: name the command, query or panel")
    return problems


if __name__ == "__main__":
    if sys.argv[1:]:
        found = {p: lint(Path(p).read_text(encoding="utf-8")) for p in sys.argv[1:]}
        for path, problems in found.items():
            print("\n".join(f"{path}: {x}" for x in problems) or f"{path}: OK")
        sys.exit(1 if any(found.values()) else 0)
    for step in ["3. Check the logs for errors.",
                 "3. Check the logs: `atlas_ops status`, then the 'calls by model and status' panel."]:
        print(lint(step, sections=False) or "OK", "<-", step[:45])

["line 1: vague step '3. Check the logs for errors.': name the command, query or panel"] <- 3. Check the logs for errors.
OK <- 3. Check the logs: `atlas_ops status`, then t

# script_id: day_088_operating_ai_systems/game_day
# scripts/game_day.py: inject a fault, replay traffic, time detection and recovery, draft the postmortem.
# Install: none (stdlib). Simulated clock + in-process Atlas stand-in, so a 10-minute drill takes ~1 s.
from __future__ import annotations

import argparse
import random
from dataclasses import dataclass, field
from typing import Callable


@dataclass
class Event:                       # compact Day 86 RequestEvent: one replayed request
    t: float
    ok: bool                       # HTTP success
    fallback: bool                 # answered by the fallback model
    passed: bool                   # the sampled online check passed (Day 86)
    tokens: int = 1_200
    flagged: bool = False          # a guardrail fired (Day 70)


@dataclass
class DrillReport:
    scenario: str
    mttd_s: float | None = None    # fault start → first alert
    mttr_s: float | None = None    # fault start → alerts clear after the mitigation
    alerts: list[str] = field(default_factory=list)
    timeline: list[tuple[float, str]] = field(default_factory=list)


def rate(window: list[Event], attr: str) -> float:
    return sum(getattr(e, attr) for e in window) / max(len(window), 1)


FOR_TICKS = 3                      # a rule must hold for 3 evaluations (30 s) before it pages
ALERTS: dict[str, Callable[[list[Event]], bool]] = {    # compact Day 86 AlertRules, last 60 s
    "error_rate_high": lambda w: 1 - rate(w, "ok") > 0.05,
    "pass_rate_low": lambda w: rate(w, "passed") < 0.75,
}
RUNBOOK_FIRST_LEVER = {"provider_outage": "fallback force", "quality_regression": "index use v1"}

# script_id: day_088_operating_ai_systems/game_day
class SimAtlas:
    """In-process stand-in for the API: only the fault and the operator's lever change its behaviour."""

    def __init__(self, scenario: str, seed: int = 7) -> None:
        self.scenario, self.rng = scenario, random.Random(seed)
        self.fault_on = False
        self.lever: str | None = None              # what the responder ran, e.g. "fallback force"

    def send(self, t: float) -> Event:
        r = self.rng.random()
        if self.fault_on and self.scenario == "provider_outage":
            if self.lever == "fallback force":     # skip the dead primary straight away
                return Event(t, ok=True, fallback=True, passed=r < 0.88)
            return Event(t, ok=r > 0.3, fallback=r > 0.3, passed=0.3 < r < 0.85)   # timeouts, then fallback
        if self.fault_on and self.scenario == "quality_regression" and self.lever != "index use v1":
            return Event(t, ok=True, fallback=False, passed=r < 0.6)   # a third of docs missing: wrong, not down
        return Event(t, ok=True, fallback=False, passed=r < 0.9)


def scripted_responder(scenario: str, *, delay_s: float) -> Callable[[float, list[str]], str | None]:
    """Stands in for you: runs the runbook's first lever delay_s after the first alert."""
    first: list[float] = []

    def respond(t: float, firing: list[str]) -> str | None:
        first[:] = first or [t]
        return RUNBOOK_FIRST_LEVER.get(scenario) if t - first[0] >= delay_s else None
    return respond

# script_id: day_088_operating_ai_systems/game_day
def run_game_day(scenario: str, *, minutes: int = 10, fault_at_s: int = 120, rps: int = 3,
                 atlas: SimAtlas | None = None,
                 responder: Callable[[float, list[str]], str | None] | None = None) -> DrillReport:
    atlas = atlas or SimAtlas(scenario)
    responder = responder or scripted_responder(scenario, delay_s=120)
    rep, events, streak = DrillReport(scenario), [], {}
    for t in range(0, minutes * 60, 10):                       # one tick = 10 simulated seconds
        if t == fault_at_s:
            atlas.fault_on = True
            rep.timeline.append((t, f"fault injected: {scenario}"))
        events += [atlas.send(t) for _ in range(rps * 10)]
        window = [e for e in events if e.t > t - 60]
        for name, rule in ALERTS.items():
            streak[name] = streak.get(name, 0) + 1 if rule(window) else 0
        firing = [name for name, n in streak.items() if n >= FOR_TICKS]
        for name in firing:
            if name not in rep.alerts:
                rep.alerts.append(name)
                rep.timeline.append((t, f"ALERT {name}"))
                rep.mttd_s = t - fault_at_s if rep.mttd_s is None else rep.mttd_s
        if firing and atlas.lever is None and (lever := responder(t, firing)):
            atlas.lever = lever
            rep.timeline.append((t, f"mitigation: atlas_ops {lever}"))
        if atlas.lever and not firing and rep.mttr_s is None:
            rep.mttr_s = t - fault_at_s
            rep.timeline.append((t, "recovered: all alerts clear"))
    return rep

# script_id: day_088_operating_ai_systems/game_day
def minutes(s: float | None) -> str:
    return "not reached" if s is None else f"{s / 60:.1f} min"


def postmortem(rep: DrillReport) -> str:
    lines = [f"# Postmortem (blameless): game day `{rep.scenario}`", "",
             f"- Time to detect: {minutes(rep.mttd_s)} · time to recover: {minutes(rep.mttr_s)}",
             f"- Alerts: {', '.join(rep.alerts) or 'none: a detection gap, write that up first'}", "",
             "## Timeline (seconds from drill start)"]
    lines += [f"- t={t:>4}s  {what}" for t, what in rep.timeline]
    lines += ["", "## Impact", "- Who was affected, for how long, how badly (from the dashboard)",
              "## Contributing factors (systems and gaps, never people)", "- ",
              "## Action items (owner · due date · ticket)", "- [ ] Fix the runbook step that slowed us down"]
    return "\n".join(lines)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Run a game-day drill on a simulated clock")
    ap.add_argument("--scenario", action="append", choices=sorted(RUNBOOK_FIRST_LEVER))
    ap.add_argument("--responder-delay", type=float, default=120, help="seconds you take to act")
    args = ap.parse_args()
    for scenario in args.scenario or sorted(RUNBOOK_FIRST_LEVER):
        rep = run_game_day(scenario, responder=scripted_responder(scenario, delay_s=args.responder_delay))
        print(f"{scenario:<19} detect={minutes(rep.mttd_s):<8} recover={minutes(rep.mttr_s):<8} {rep.alerts}")
    print("\n" + postmortem(rep))

provider_outage     detect=0.3 min  recover=3.2 min  ['error_rate_high', 'pass_rate_low']
quality_regression  detect=0.8 min  recover=3.3 min  ['pass_rate_low']

# Postmortem (blameless): game day `quality_regression`

- Time to detect: 0.8 min · time to recover: 3.3 min
- Alerts: pass_rate_low

## Timeline (seconds from drill start)
- t= 120s  fault injected: quality_regression
- t= 170s  ALERT pass_rate_low
- t= 290s  mitigation: atlas_ops index use v1
- t= 320s  recovered: all alerts clear

## Impact
- Who was affected, for how long, how badly (from the dashboard)
## Contributing factors (systems and gaps, never people)
-
## Action items (owner · due date · ticket)
- [ ] Fix the runbook step that slowed us down

# script_id: day_088_operating_ai_systems/mixed_vectors
# Break it: compare one model's vector with another model's. No error, only a meaningless score.
# Install: uv add numpy. Two seeded "models" stand in for text-embedding-3-small and -large (offline).
import hashlib

import numpy as np


def fake_embed(text: str, model: str, dims: int = 256) -> np.ndarray:
    """Hashed bag of words. A different model name sends each word to a different axis, like a real model."""
    v = np.zeros(dims, dtype=np.float32)
    for word in text.lower().split():
        h = int.from_bytes(hashlib.sha256(f"{model}:{word}".encode()).digest()[:8], "little")
        v[h % dims] += 1.0 if h >> 63 else -1.0
    return v / (np.linalg.norm(v) or 1.0)


doc = "workers renew their lease with a heartbeat every heartbeat_seconds"
query = "how often do workers send a lease heartbeat"
print(f"same model : cosine={fake_embed(doc, 'small') @ fake_embed(query, 'small'):.2f}")
print(f"mixed      : cosine={fake_embed(doc, 'small') @ fake_embed(query, 'large'):.2f}  <- no error, no meaning")

same model : cosine=0.47
mixed      : cosine=0.12  <- no error, no meaning

# script_id: day_088_operating_ai_systems/index_migration
# atlas/index_migration.py: re-embed Atlas's index with zero downtime (expand → migrate → contract).
# Install: uv add numpy     Data: ORBIT_DATA/docs/*.md, ORBIT_DATA/eval/retrieval_queries.jsonl
# Runs offline on SQLite + NumPy with a fake embedder. Production: the same steps on Postgres + pgvector
# (Go Deeper) with openai_embed() from the next script; the function signatures don't change.
from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable

import numpy as np

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))


@dataclass(frozen=True)
class IndexVersion:
    name: str
    embed_model: str
    dims: int
    table: str


INDEX_VERSIONS = {
    "v1": IndexVersion("v1", "text-embedding-3-small", 1536, "chunks_v1"),
    "v2": IndexVersion("v2", "text-embedding-3-large", 1024, "chunks_v2"),   # dimensions=1024: HNSW-friendly
}
Embedder = Callable[[list[str], IndexVersion], np.ndarray]

# script_id: day_088_operating_ai_systems/index_migration
def fake_embed(texts: list[str], v: IndexVersion) -> np.ndarray:
    """Offline stand-in: hashed bag of words seeded by the model name. Tests mechanics, not quality."""
    out = np.zeros((len(texts), v.dims), dtype=np.float32)
    for i, text in enumerate(texts):
        words = text.lower().split() if v.name == "v1" else re.findall(r"[a-z0-9_]+", text.lower())
        for w in words:
            h = int.from_bytes(hashlib.sha256(f"{v.embed_model}:{w}".encode()).digest()[:8], "little")
            out[i, h % v.dims] += 1.0 if h >> 63 else -1.0
    return out / np.maximum(np.linalg.norm(out, axis=1, keepdims=True), 1e-9)

# script_id: day_088_operating_ai_systems/index_migration
def chunk(doc_id: str, text: str, max_chars: int = 1200) -> list[dict]:
    """Compact Day 25 chunker: paragraphs packed up to max_chars; ids stay the same in every version."""
    parts, buf = [], ""
    for para in text.split("\n\n"):
        if buf and len(buf) + len(para) > max_chars:
            parts, buf = [*parts, buf], ""
        buf = f"{buf}\n\n{para}" if buf else para
    parts += [buf] if buf else []
    return [{"id": f"{doc_id}#{i}", "doc_id": doc_id, "body": p,
             "hash": hashlib.sha256(p.encode()).hexdigest()[:16]} for i, p in enumerate(parts)]


def open_db(path: str | Path) -> sqlite3.Connection:
    conn = sqlite3.connect(path, timeout=5, isolation_level=None, check_same_thread=False)
    conn.execute("PRAGMA journal_mode=WAL")               # readers keep reading while cutover writes
    conn.execute("CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)")
    conn.execute("CREATE TABLE IF NOT EXISTS index_events (at TEXT, actor TEXT, event TEXT)")
    return conn


def create_version(conn: sqlite3.Connection, v: IndexVersion) -> None:
    """Expand: a new table beside the old one. Nothing reads it yet. (Postgres DDL: see Go Deeper.)"""
    conn.execute(f"CREATE TABLE IF NOT EXISTS {v.table} (id TEXT PRIMARY KEY, doc_id TEXT NOT NULL,"
                 " body TEXT NOT NULL, content_hash TEXT NOT NULL, embedding BLOB NOT NULL)")


def upsert(conn: sqlite3.Connection, v: IndexVersion, chunks: list[dict], embed: Embedder) -> None:
    vecs = embed([c["body"] for c in chunks], v)
    conn.executemany(f"INSERT INTO {v.table} VALUES (?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET"
                     " body=excluded.body, content_hash=excluded.content_hash, embedding=excluded.embedding",
                     [(c["id"], c["doc_id"], c["body"], c["hash"], vec.astype(np.float32).tobytes())
                      for c, vec in zip(chunks, vecs)])


def dual_write(conn: sqlite3.Connection, chunk: dict, versions: list[IndexVersion], embed: Embedder) -> None:
    """Ingest writes every live version during a migration, so the backfill never falls behind."""
    for v in versions:
        upsert(conn, v, [chunk], embed)

# script_id: day_088_operating_ai_systems/index_migration
@dataclass
class BackfillReport:
    copied: int
    skipped: int
    seconds: float


def backfill(conn: sqlite3.Connection, src: IndexVersion, dst: IndexVersion, embed: Embedder, *,
             batch: int = 100, pause_s: float = 0.0) -> BackfillReport:
    """Re-embed every src row that dst lacks or holds stale. Idempotent: rerun it after a crash."""
    start = time.monotonic()
    have = dict(conn.execute(f"SELECT id, content_hash FROM {dst.table}").fetchall())
    rows = conn.execute(f"SELECT id, doc_id, body, content_hash FROM {src.table} ORDER BY id").fetchall()
    todo = [r for r in rows if have.get(r[0]) != r[3]]
    for i in range(0, len(todo), batch):
        part = [{"id": r[0], "doc_id": r[1], "body": r[2], "hash": r[3]} for r in todo[i:i + batch]]
        upsert(conn, dst, part, embed)
        print(f"  backfill {dst.name}: {len(rows) - len(todo) + i + len(part)}/{len(rows)} rows")
        time.sleep(pause_s)                     # stay under the embeddings TPM (tokens-per-minute) limit
    return BackfillReport(len(todo), len(rows) - len(todo), round(time.monotonic() - start, 2))


def parity(conn: sqlite3.Connection, a: IndexVersion, b: IndexVersion) -> dict[str, int]:
    """Row diff by id + content hash. Cutover needs missing == stale == extra == 0."""
    ra = dict(conn.execute(f"SELECT id, content_hash FROM {a.table}").fetchall())
    rb = dict(conn.execute(f"SELECT id, content_hash FROM {b.table}").fetchall())
    return {f"rows_{a.name}": len(ra), f"rows_{b.name}": len(rb),
            "missing": len(ra.keys() - rb.keys()), "extra": len(rb.keys() - ra.keys()),
            "stale": sum(ra[k] != rb[k] for k in ra.keys() & rb.keys())}

# script_id: day_088_operating_ai_systems/index_migration
def search(conn: sqlite3.Connection, v: IndexVersion, query: str, embed: Embedder, k: int = 5) -> list[str]:
    """Top-k doc ids by cosine. Brute force suits ~500 rows; pgvector uses its HNSW index instead."""
    rows = conn.execute(f"SELECT doc_id, embedding FROM {v.table}").fetchall()
    if not rows:
        return []                              # the silent-empty case: alert on it (zero-chunk share)
    m = np.frombuffer(b"".join(r[1] for r in rows), dtype=np.float32).reshape(len(rows), v.dims)
    out: list[str] = []
    for i in np.argsort(-(m @ embed([query], v)[0])):
        if rows[i][0] not in out:
            out.append(rows[i][0])
        if len(out) == k:
            break
    return out


def compare_versions(conn: sqlite3.Connection, a: IndexVersion, b: IndexVersion, queries: list[dict],
                     embed: Embedder) -> dict[str, dict[str, float]]:
    """Shadow read: the same queries against both versions, Day 29 metrics. Users still see only `a`."""
    report = {}
    for v in (a, b):
        recall, rr, ms = [], [], []
        for q in queries:
            t0 = time.perf_counter()
            got, rel = search(conn, v, q["query"], embed), set(q["relevant_doc_ids"])
            ms.append((time.perf_counter() - t0) * 1000)
            recall.append(len(rel & set(got)) / len(rel))
            rr.append(next((1 / (i + 1) for i, d in enumerate(got) if d in rel), 0.0))
        report[v.name] = {"recall@5": round(float(np.mean(recall)), 3), "mrr": round(float(np.mean(rr)), 3),
                          "p95_ms": round(float(np.percentile(ms, 95)), 1)}
    return report

# script_id: day_088_operating_ai_systems/index_migration
class CutoverRefused(RuntimeError):
    pass


def active_version(conn: sqlite3.Connection) -> str:
    row = conn.execute("SELECT value FROM index_meta WHERE key = 'active'").fetchone()
    return row[0] if row else "v1"


def resolve_version(conn: sqlite3.Connection, override: str | None) -> IndexVersion:
    """Per request: the ops override (Overrides.index_version) wins, then the DB pointer."""
    return INDEX_VERSIONS[override or active_version(conn)]


def cutover(conn: sqlite3.Connection, name: str, *, actor: str,
            gate: Callable[[], str | None] | None = None) -> None:
    """Atomically repoint reads. Refuses unless the target matches the current version row for row."""
    cur = INDEX_VERSIONS[active_version(conn)]
    p = parity(conn, cur, INDEX_VERSIONS[name])
    if p["missing"] or p["stale"] or p["extra"]:
        raise CutoverRefused(f"{name} is not in parity with {cur.name}: {p}")
    if gate and (why := gate()):                # e.g. an eval gate (Exercise 3)
        raise CutoverRefused(why)
    upsert_meta = "INSERT INTO index_meta VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value"
    conn.execute("BEGIN IMMEDIATE")
    conn.execute(upsert_meta, ("previous", cur.name))
    conn.execute(upsert_meta, ("active", name))
    conn.execute("INSERT INTO index_events VALUES (?, ?, ?)",
                 (datetime.now(timezone.utc).isoformat(timespec="seconds"), actor, f"cutover {cur.name} -> {name}"))
    conn.execute("COMMIT")
    print(f"cutover {cur.name} -> {name} by {actor}; rollback = cutover(conn, {cur.name!r}, actor=...)")


def contract(conn: sqlite3.Connection, old: IndexVersion, *, actor: str) -> None:
    """After the bake window: stop dual-writing `old` in ingest first, then drop it."""
    assert active_version(conn) != old.name, "never drop the version that is serving"
    conn.execute(f"DROP TABLE {old.table}")
    conn.execute("INSERT INTO index_events VALUES (datetime('now'), ?, ?)", (actor, f"contract: dropped {old.name}"))

# script_id: day_088_operating_ai_systems/index_migration
def load_docs() -> dict[str, str]:
    return {p.stem: p.read_text(encoding="utf-8") for p in sorted((DATA_DIR / "docs").glob("*.md"))}

# script_id: day_088_operating_ai_systems/index_migration
if __name__ == "__main__":
    import tempfile
    v1, v2 = INDEX_VERSIONS["v1"], INDEX_VERSIONS["v2"]
    conn = open_db(Path(tempfile.mkdtemp()) / "atlas_index.db")
    docs = load_docs()
    create_version(conn, v1)
    upsert(conn, v1, [c for d, text in docs.items() for c in chunk(d, text)], fake_embed)   # today's index
    create_version(conn, v2)                                                    # 1 expand
    try:
        cutover(conn, "v2", actor="oncall")                                     # the 02:10 mistake
    except CutoverRefused as exc:
        print("refused:", str(exc)[:96])
    edited = chunk("faq", docs["faq"] + "\n\nQ: Does 3.2.1 change leases? A: No.")[-1]
    dual_write(conn, edited, [v1, v2], fake_embed)                              # 2 a doc edit lands in both
    print(backfill(conn, v1, v2, fake_embed, batch=200))                        # 3 backfill
    print(parity(conn, v1, v2))                                                 # 4a parity
    queries = [json.loads(line) for line in (DATA_DIR / "eval" / "retrieval_queries.jsonl").open()]
    for name, metrics in compare_versions(conn, v1, v2, queries, fake_embed).items():   # 4b shadow read
        print(f"  {name}: {metrics}")
    cutover(conn, "v2", actor="oncall")                                         # 5 cutover; 6 bake begins
    print("serving:", resolve_version(conn, None).name, "| with override:", resolve_version(conn, "v1").name)

# your timings will differ
refused: v2 is not in parity with v1: {'rows_v1': 257, 'rows_v2': 0, 'missing': 257, 'extra': 0, 'stale':
  backfill v2: 201/257 rows
  backfill v2: 257/257 rows
BackfillReport(copied=256, skipped=1, seconds=0.04)
{'rows_v1': 257, 'rows_v2': 257, 'missing': 0, 'extra': 0, 'stale': 0}
  v1: {'recall@5': 0.542, 'mrr': 0.548, 'p95_ms': 0.4}
  v2: {'recall@5': 0.658, 'mrr': 0.722, 'p95_ms': 0.3}
cutover v1 -> v2 by oncall; rollback = cutover(conn, 'v1', actor=...)
serving: v2 | with override: v1

# script_id: day_088_operating_ai_systems/openai_embed
# The real embedder for index_migration (same signature as fake_embed). Env: OPENAI_API_KEY
# Install: uv add "openai>=3.46,<4" numpy
import numpy as np
from openai import OpenAI

client = OpenAI()                                     # the SDK retries transient errors twice by default


def openai_embed(texts: list[str], v) -> np.ndarray:  # v: IndexVersion (uses .embed_model and .dims)
    """Passes v.dims as `dimensions=`, so a 3-large index stays under pgvector's 2,000-dim HNSW limit."""
    vectors: list[list[float]] = []
    for i in range(0, len(texts), 256):               # several small requests; each input ≤ 8,192 tokens
        resp = client.embeddings.create(model=v.embed_model, input=texts[i:i + 256], dimensions=v.dims)
        vectors += [d.embedding for d in resp.data]
    return np.asarray(vectors, dtype=np.float32)

# script_id: day_088_operating_ai_systems/check_deprecations
# scripts/check_deprecations.py: fail CI when code or config names a model that shuts down soon.
# Install: none (stdlib). Usage: python check_deprecations.py atlas/ scripts/ release.toml [--days 60]
# Dates/replacements: https://developers.openai.com/api/docs/deprecations (as of 2026-11-06, verify).
from __future__ import annotations

import argparse
import os
import re
import sys
import tempfile
from dataclasses import dataclass
from datetime import date
from pathlib import Path


@dataclass(frozen=True)
class Deprecation:
    model_id: str
    shutdown: date
    replacement: str


DEPRECATIONS = [
    *(Deprecation(m, date(2026, 12, 11), "gpt-5.6-sol/-terra/-luna")
      for m in ("gpt-5", "gpt-5-mini", "gpt-5-nano", "gpt-5-pro", "o3", "o3-pro")),
    *(Deprecation(m, date(2027, 1, 6), "gpt-realtime-2.1-mini via Realtime")
      for m in ("tts-1", "tts-1-hd", "gpt-4o-mini-tts")),
    *(Deprecation(m, date(2027, 1, 14), "gpt-6-luna") for m in ("gpt-4o", "gpt-4o-mini", "gpt-4.1", "gpt-4.1-mini")),
    Deprecation("gpt-realtime", date(2027, 1, 20), "gpt-realtime-2.1"),
    Deprecation("gpt-realtime-mini", date(2027, 1, 20), "gpt-realtime-2.1-mini"),
    *(Deprecation(m, date(2027, 2, 26), "gpt-transcribe")
      for m in ("whisper-1", "gpt-4o-transcribe", "gpt-4o-mini-transcribe", "gpt-4o-transcribe-diarize")),
]
# whole ID only, plus an optional dated snapshot: "gpt-4o" matches gpt-4o-2024-08-06, never gpt-4o-mini
PATTERNS = [(d, re.compile(rf"(?<![\w.-]){re.escape(d.model_id)}(?:-\d{{4}}-\d{{2}}-\d{{2}})?(?![\w.-])"))
            for d in DEPRECATIONS]

# script_id: day_088_operating_ai_systems/check_deprecations
@dataclass(frozen=True)
class Finding:
    path: str
    line: int
    dep: Deprecation
    days_left: int


def scan(paths: list[str], *, today: date | None = None) -> list[Finding]:
    """Every deprecated model ID in code and config files (not docs: history is allowed there)."""
    today, found = today or date.today(), []
    files = [f for p in map(Path, paths) for f in ([p] if p.is_file() else sorted(p.rglob("*")))
             if f.is_file() and (f.suffix in {".py", ".toml", ".yaml", ".yml", ".json"} or f.name.startswith(".env"))]
    for f in files:
        for n, text in enumerate(f.read_text(encoding="utf-8", errors="replace").splitlines(), 1):
            found += [Finding(str(f), n, d, (d.shutdown - today).days) for d, rx in PATTERNS if rx.search(text)]
    return sorted(found, key=lambda x: x.days_left)

# script_id: day_088_operating_ai_systems/check_deprecations
def demo_repo() -> str:
    d = Path(tempfile.mkdtemp())
    (d / "judges.py").write_text('JUDGE_MODEL = "gpt-4o-mini"   # picked in 2025 as the cheap judge\n')
    (d / "voice.py").write_text('TTS_MODEL = "tts-1"\n')
    (d / "release.toml").write_text('[ask]\nmodel = "gpt-6-luna"\n[voice]\nstt = "whisper-1"\n')
    return str(d)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Fail CI on soon-to-be-retired model IDs")
    ap.add_argument("paths", nargs="*", help="files or directories (none = demo repo)")
    ap.add_argument("--days", type=int, default=60, help="FAIL when a shutdown is this close")
    ap.add_argument("--today", type=date.fromisoformat, default=None, help="YYYY-MM-DD, for tests")
    a = ap.parse_args(argv)
    if not a.paths:
        os.chdir(demo_repo())
    findings = scan(a.paths or ["."], today=a.today)
    for x in findings:
        level = "FAIL" if x.days_left <= a.days else "WARN"
        print(f"{level} {x.dep.model_id} in {os.path.relpath(x.path)}:{x.line}: shuts down "
              f"{x.dep.shutdown} ({x.days_left} days) → {x.dep.replacement}")
    failed = any(x.days_left <= a.days for x in findings)
    return 1 if failed and a.paths else 0     # the demo always exits 0


if __name__ == "__main__":
    sys.exit(main())

# days are counted from today; these are for 2026-11-07
FAIL tts-1 in voice.py:1: shuts down 2027-01-06 (60 days) → gpt-realtime-2.1-mini via Realtime
WARN gpt-4o-mini in judges.py:1: shuts down 2027-01-14 (68 days) → gpt-6-luna
WARN whisper-1 in release.toml:4: shuts down 2027-02-26 (111 days) → gpt-transcribe