diff --git a/.gitignore b/.gitignore index 365770d..11fa760 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,14 @@ go.work.sum /trace.jsonl /report.md /report.json +/model.json +/predicted-trace.jsonl +/cloudoracle-trace.jsonl +/cloudoracle-model.json + +# Python sidecar artifacts +__pycache__/ +*.py[cod] +.pytest_cache/ +*.egg-info/ +.venv/ diff --git a/README.md b/README.md index 984e825..9fdaff5 100644 --- a/README.md +++ b/README.md @@ -182,6 +182,34 @@ Retries and fan-out scale the whole call (more calls); context growth inflates only the prompt side, not the completion — so the model reflects *which* driver moved, not a flat fudge factor. +### Predict cost without running (Python sidecar) + +The what-if knobs above re-cost a recorded trace, but they hold output length +fixed. The one thing you genuinely can't guess for an agent is **completion +length** — and it's where the cost surprise lives. The optional +[`sidecar/`](sidecar/) (Python) learns output length from a recorded trace, then +projects cost for inputs you never ran: + +```sh +# Learn output_tokens ≈ a + b·input_tokens per model, from one recorded run. +python -m augur_predict fit --trace trace.jsonl --out model.json + +# "What if prompts grow 1.5× as conversations lengthen?" — predict, don't run. +python -m augur_predict emit-trace --model model.json --out grown.jsonl --input-scale 1.5 + +# The predicted trace flows through the normal gate — no tokens spent. +augur gate --trace grown.jsonl --traffic traffic.yaml --budget budget.yaml +``` + +Coupling is the trace file and nothing else — no RPC, no shared library. Unlike +`--context-growth`, the sidecar feeds the larger prompt *back through the output +model*, so a bigger ask predicts a longer answer, not just a costlier prompt. +The default fit is an honest linear baseline (it reports R² and falls back to the +mean when the signal is weak); `--dist quantile` fits conditional quantiles +directly to model the right-skewed tail the gate's p95 lives in, and +`--run-correlation` widens the per-run spread to match correlated runs. See +[`sidecar/README.md`](sidecar/README.md). + ### Self-hosted models (TCO) For a model you run yourself there's no per-token API price — you pay for an @@ -252,7 +280,8 @@ example. truthfully. Classifying those calls into retries vs sub-agent fan-out needs labeling the trace does not yet carry. - **Running the agent in CI spends real tokens.** Keep the scenario set small and - `runs` modest; a record-once/replay mode is on the roadmap. + `runs` modest, or record once and replay (`--record`/`--replay`) so CI pushes + spend nothing. ## Status @@ -266,13 +295,42 @@ dependency is `gopkg.in/yaml.v3`): | Hito 2 | scenario runner + per-scenario aggregation | | Hito 3 | projection engine with bootstrap confidence intervals | | Hito 4 | budget gate + Markdown/JSON report + CI exit codes | -| Hito 5 | record/replay cassette (`--record`/`--replay`), what-if knobs (`--retry-rate`/`--fanout`/`--context-growth`), self-hosted TCO mode (`augur tco`, `--tco`), GitHub Action (`action.yml`) | +| Hito 5 | record/replay cassette (`--record`/`--replay`), what-if knobs (`--retry-rate`/`--fanout`/`--context-growth`), self-hosted TCO mode (`augur tco`, `--tco`), GitHub Action (`action.yml`), Python output-length prediction sidecar ([`sidecar/`](sidecar/)) | -**Roadmap (stretch):** a Python output-length prediction sidecar (the analytical -piece that would earn a second language). +Every SPEC milestone and stretch is implemented. The Go core is pure Go (only +external dependency `gopkg.in/yaml.v3`); the optional prediction sidecar is +Python (numpy + scipy), coupled to the core through the trace file alone. See [`SPEC.md`](SPEC.md) for the full design. +## Dogfooded on a real agent + +Augur has been run against a real, non-trivial agent — the Insights Agent from +its sibling project [**CloudOracle**](https://github.com/Cro22/CloudOracle) (a +LangGraph supervisor multi-agent with tool calls and guardrails), not just +synthetic traces. CloudOracle is FinOps for *cloud* (runtime); Augur is FinOps +for *AI agents* (pre-prod) — dogfooding one on the other closes the loop. + +Because that agent talks to its model natively through LangChain (no OpenAI +`base_url`), the capture used the SPEC's documented proxy fallback (ADR D1): a +LangChain **callback shim** that writes Augur's trace schema from the usage every +call reports — *without editing CloudOracle*. Full walkthrough and harness: +[`examples/cloudoracle/`](examples/cloudoracle/). + +**What it found (Claude Haiku 4.5, 20 runs): gate PASS** at `$/request p95 +$0.0198` (budget $0.02). The headline is the `find-savings` scenario — its **p95 +cost is 2.3× its median**, driven by a call-count tail (5 → 13 calls/run when the +savings specialist's tool loop keeps going). That agentic cost driver, observed on +a real agent, is exactly what Augur exists to catch; gating on the mean would hide +it. + +Dogfooding also surfaced real findings *about the agent* — it isn't +provider-portable (its LLM judge sends a system-only message Anthropic rejects but +Gemini tolerates), Claude Haiku hallucinates dollar figures that trip the agent's +grounding check, and the sidecar's input→output signal is weak there (output +length is driven by call *role*, not prompt size). Details in the +[example README](examples/cloudoracle/README.md). + ## Naming **Augur** — a seer who reads omens to foretell what's coming; here, an agent's diff --git a/SPEC.md b/SPEC.md index 8a6912f..0febf22 100644 --- a/SPEC.md +++ b/SPEC.md @@ -177,6 +177,7 @@ budget.yaml ──┐││ LLM calls (base_url ## 10. Open questions for Jesús (decide before/at Hito 0) 1. **First dogfood target.** Strong candidate: **CloudOracle's own Insights Agent** (LangGraph multi-agent + RAG = real, messy cost drivers, and it ties the two repos together in the story). Alternatives: a toy LangGraph agent, or Despachito if it gains an LLM feature. → *Recommend CloudOracle's Insights Agent.* + - **✅ Resolved (done).** Dogfooded on [CloudOracle's Insights Agent](https://github.com/Cro22/CloudOracle) via a LangChain callback shim (ADR D1's proxy fallback — the agent speaks its model natively), CloudOracle source untouched. Gate PASS; the `find-savings` scenario showed a real call-count tail (p95 2.3× median). Harness + findings: [`examples/cloudoracle/`](examples/cloudoracle/). 2. **Pricing data:** ship a dated snapshot only (v1), or attempt a live fetch? 3. **CI token-cost tolerance:** how much real spend per CI run is acceptable? This decides whether record/replay is v1 or stretch. 4. **Proxy vs harness for the runner:** does Jesús want the tool to *drive* the agent (needs an entrypoint contract), or just provide the proxy + headers and let the user run their own harness? (Lighter v1 = the latter.) diff --git a/examples/cloudoracle/README.md b/examples/cloudoracle/README.md new file mode 100644 index 0000000..02ce234 --- /dev/null +++ b/examples/cloudoracle/README.md @@ -0,0 +1,141 @@ +# Dogfooding Augur on CloudOracle's Insights Agent + +> Augur measuring a *real* agent, not a synthetic trace: the Insights Agent from +> [**CloudOracle**](https://github.com/Cro22/CloudOracle) — a LangGraph supervisor +> multi-agent with tool calls and guardrails. + +CloudOracle is Augur's sibling project: FinOps for *cloud* (runtime, post-hoc), +where Augur is FinOps for *AI agents* (pre-prod, predictive). Dogfooding Augur on +CloudOracle's own agent closes the loop between the two — and is the first test of +Augur against an agent it didn't author. + +This is the first dogfood: point Augur at an agent it didn't author and see +whether the cost picture holds up. It does — and it surfaced real findings, +which is the point of dogfooding. + +## The integration: a callback shim, not the proxy + +Augur's recording proxy is OpenAI-compatible, but the Insights Agent talks to its +model **natively** through LangChain (no OpenAI `base_url` to redirect). So we use +the capture path the [SPEC](../../SPEC.md) documents as the fallback "for +frameworks where the proxy is awkward" (ADR D1): a **LangChain callback** that +writes Augur's exact `trace.jsonl` schema from the token usage every call reports. + +The agent reuses one chat model across the supervisor, every specialist's ReAct +loop, and the synthesizer, so a callback on that model sees the **entire call +graph** of a run — the real agentic fan-out. Coupling is one direction only: +[`augur_dogfood.py`](augur_dogfood.py) *imports* the agent's public building +blocks (`build_supervisor_graph`, `build_tools`, `run_guarded`) and builds them +with our model; **CloudOracle's source is never edited.** Delete this folder and +both repos are untouched. + +## Files + +| file | what | +|---|---| +| [`augur_dogfood.py`](augur_dogfood.py) | the harness: builds the agent graph with our model + trace callback, runs the scenarios | +| [`scenarios.yaml`](scenarios.yaml) | representative FinOps questions, spread across the agent's three specialists | +| [`traffic.yaml`](traffic.yaml) | production volume assumptions for the projection | +| [`budget.yaml`](budget.yaml) | the thresholds the gate enforces | +| [`pricing-gemini.yaml`](pricing-gemini.yaml) | a Gemini price snapshot (for the `--provider gemini` path; the Claude run uses Augur's default `pricing.yaml`, which already lists Claude) | + +## Running it + +Prereqs — bring up CloudOracle's stack so the agent's tools have data to fetch: + +```sh +cd /path/to/CloudOracle +docker compose up -d postgres # pgvector +go build -o oracle.exe ./cmd/oracle +# DB defaults (oracle/oracle_dev/cloudoracle) match docker-compose; seed + serve: +CLOUDORACLE_PROVIDER=synthetic LOG_LEVEL=info ./oracle.exe seed --count 120 +CLOUDORACLE_PROVIDER=synthetic LOG_LEVEL=info ./oracle.exe serve --port 8080 & +``` + +Then run the agent under the Augur callback (uses CloudOracle's venv + `.env`): + +```sh +set -a; source insights-agent/.env; set +a # GEMINI/ANTHROPIC + CLOUDORACLE keys +unset DATABASE_URL # RAG off for a clean first run (see caveats) + +insights-agent/.venv/Scripts/python.exe \ + /path/to/Augur/examples/cloudoracle/augur_dogfood.py \ + --scenarios /path/to/Augur/examples/cloudoracle/scenarios.yaml \ + --out /path/to/Augur/cloudoracle-trace.jsonl \ + --provider anthropic --runs 5 --no-judge +``` + +Pipe the real trace through Augur (Claude prices ship in the default snapshot): + +```sh +cd /path/to/Augur +./augur gate --trace cloudoracle-trace.jsonl \ + --traffic examples/cloudoracle/traffic.yaml \ + --budget examples/cloudoracle/budget.yaml \ + --pricing pricing.yaml + +# And learn the output-length model from the real run: +cd sidecar && python -m augur_predict fit --trace ../cloudoracle-trace.jsonl \ + --out ../cloudoracle-model.json --dist quantile +python -m augur_predict report --model ../cloudoracle-model.json +``` + +`--provider gemini` runs the same harness on Gemini (add `--rps 0.15` to stay under +the free-tier rate limit); `--provider anthropic` (default) uses Claude, whose +higher limits make a clean multi-call run practical. + +## What Augur measured (Claude Haiku 4.5, 20 runs) + +``` +scenario "find-savings" — 5 run(s) + metric mean p50 p95 stdev min max + $/run 0.022428 0.016974 0.038833 0.012192 0.016830 0.044236 + calls/run 6.60 5.00 11.40 3.58 5 13 <-- the tail +``` + +| scenario | calls/run (p50→p95) | $/run p95 | note | +|---|---|---|---| +| cost-breakdown | 4 → 4 | $0.0069 | input-heavy, cheap output | +| concept-rightsizing | 4 → 4 | $0.0105 | stable | +| full-review | 5 → 5 | $0.0197 | the deliberate multi-specialist path | +| **find-savings** | **5 → 11.4 (max 13)** | **$0.0388** | **the savings specialist sometimes loops** | + +**Gate verdict: ✅ PASS** — projected `$/request p95 = $0.0198` (budget $0.02 — it +*just* fits), `$39.72/tenant/month`, `$794/month` at the assumed volume. + +The headline is `find-savings`: its **p95 cost is 2.3× its median**, driven by a +call-count tail (5 → 13 calls/run) when the savings specialist's ReAct loop keeps +going. That is precisely the agentic cost driver Augur exists to catch — observed +on a real agent, not assumed. Gating on the mean would have hidden it; gating on +p95 (and the wide p95 CI `[$0.019, $0.044]`) surfaces it. + +## Findings the dogfood surfaced + +Dogfooding earns its keep by what breaks: + +- **The agent isn't provider-portable.** Built for Gemini, its LLM-judge layer + sends a *system-only* message — which Gemini tolerates but Anthropic rejects + (`messages: at least one message is required`). `--no-judge` works around it for + the Claude run; the fix belongs in CloudOracle, and Augur found it. (This is a + bug in the agent, not in Augur.) +- **Claude Haiku hallucinates figures on this agent.** Most `find-savings` / + `full-review` runs hit the agent's *deterministic grounding* fallback ("answer + states $320.00 … not found in any tool result"). A quality signal, orthogonal + to cost — Augur still measured the real token spend of those runs. +- **The sidecar's output model is weak here (R¹ ≈ 0.03).** Output length isn't + driven by input length for this agent; it's driven by the call's *role* (a + router turn is terse, a synthesis turn is long). The sidecar reports the weak + fit honestly rather than faking a trend — and it motivates the documented next + step: segment the output model by call role / seq, not just `input_tokens`. + +## Caveats / not captured + +- **RAG embeddings.** The callback rides the *chat* model; the agent's embeddings + go through a separate Gemini embeddings client, so RAG retrieval calls aren't in + this trace. `DATABASE_URL` is unset here to keep the run clean. Capturing them + would mean instrumenting `GeminiEmbeddingsProvider` too. +- **Synthetic backend.** CloudOracle serves synthetic resources (`seed`), so the + *tool* outputs are representative-shaped, not a real cloud bill. The *agent's* + token usage — what Augur measures — is real. +- **Small N.** 20 runs is enough to see the find-savings tail but thin for a + precise p95; the gate's wide CI says so. diff --git a/examples/cloudoracle/augur_dogfood.py b/examples/cloudoracle/augur_dogfood.py new file mode 100644 index 0000000..7115a30 --- /dev/null +++ b/examples/cloudoracle/augur_dogfood.py @@ -0,0 +1,365 @@ +"""Dogfood harness: run CloudOracle's Insights Agent under Augur. + +This is the **callback-shim** capture path the SPEC documents as the fallback +to the proxy "for frameworks where the proxy is awkward" (ADR D1). The Insights +Agent talks to Gemini natively through `langchain_google_genai` (no OpenAI +`base_url` to point at Augur's proxy), so instead of a proxy we attach a +LangChain callback to the one chat model the whole graph shares and write the +exact Augur cost-trace schema from the usage every call reports. + +Why this is loose coupling, not a CloudOracle change: we import the agent's +public runtime, build it as the CLI does, then set `.callbacks` on its shared +`_chat` model before running. Nothing in CloudOracle is edited. Delete this file +and both repos are untouched. The output is an ordinary `trace.jsonl` that +`augur aggregate | project | gate` and the sidecar consume with no idea a +callback (not the proxy) produced it. + +Run it with CloudOracle's venv, from the CloudOracle repo so its `.env` and +`src/` are found: + + cd F:/JetBrains/GoProjects/CloudOracle + insights-agent/.venv/Scripts/python.exe \ + F:/JetBrains/GoProjects/Augur/examples/cloudoracle/augur_dogfood.py \ + --scenarios F:/JetBrains/GoProjects/Augur/examples/cloudoracle/scenarios.yaml \ + --out F:/JetBrains/GoProjects/Augur/trace.jsonl --runs 3 +""" + +from __future__ import annotations + +import argparse +import asyncio +import datetime as _dt +import json +import os +import sys +import threading +from pathlib import Path +from typing import Any + +# Make the agent importable when run from the CloudOracle repo root. +_AGENT_SRC = Path.cwd() / "insights-agent" / "src" +if _AGENT_SRC.is_dir(): + sys.path.insert(0, str(_AGENT_SRC)) + +try: + from langchain_core.callbacks import BaseCallbackHandler + from langchain_core.outputs import LLMResult +except Exception as e: # pragma: no cover - environment guard + sys.exit(f"augur_dogfood: LangChain not importable ({e}); run with the " + f"insights-agent venv from the CloudOracle repo root.") + + +class AugurTraceCallback(BaseCallbackHandler): + """Writes one Augur trace row per LLM call from the usage LangChain reports. + + The agent reuses a single chat model across the supervisor, every + specialist's ReAct loop, the synthesizer, and the LLM judge, so a callback + on that model sees the *entire* call graph of a run — the real agentic + fan-out Augur exists to measure. `on_llm_end` is the hook: chat models + populate `usage_metadata` on the returned message (and `llm_output` as a + fallback), which mirrors the provider's own token report. + + The (scenario_id, run_id, seq) tagging matches what the proxy would stamp + from request headers; here we set them directly per run. Seq is the call's + order within the run, preserving the call-graph ordering downstream tools + mine for multipliers. + """ + + def __init__(self, writer: "TraceWriter", model_fallback: str) -> None: + self._writer = writer + self._model_fallback = model_fallback + self._lock = threading.Lock() + self.scenario_id = "" + self.run_id = "" + self._seq = 0 + + def begin_run(self, scenario_id: str, run_id: str) -> None: + with self._lock: + self.scenario_id = scenario_id + self.run_id = run_id + self._seq = 0 + + def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + usage = _extract_usage(response) + if usage is None: + # No usage reported (rare): still record the call with zero tokens + # so the call graph's shape (fan-out count) stays truthful. + usage = {"input": 0, "output": 0, "cached": 0} + model = _extract_model(response) or self._model_fallback + with self._lock: + seq = self._seq + self._seq += 1 + scenario_id, run_id = self.scenario_id, self.run_id + self._writer.write({ + "ts": _dt.datetime.now(_dt.timezone.utc).isoformat(), + "scenario_id": scenario_id, + "run_id": run_id, + "seq": seq, + "model": model, + "input_tokens": usage["input"], + "output_tokens": usage["output"], + "cached_tokens": usage["cached"], + "latency_ms": 0, + "endpoint": "/v1/chat/completions", + "status": 200, + }) + + +def _extract_usage(response: LLMResult) -> dict[str, int] | None: + """Pull (input, output, cached) tokens from an LLMResult, provider-agnostic. + + Prefers the per-message `usage_metadata` LangChain normalises across + providers; falls back to the raw `llm_output` usage block. Cached prompt + tokens, when reported, live in `input_token_details.cache_read` and are a + SUBSET of input (Augur's convention), so we clamp them to input. + """ + for gens in response.generations: + for gen in gens: + msg = getattr(gen, "message", None) + um = getattr(msg, "usage_metadata", None) if msg is not None else None + if um: + inp = int(um.get("input_tokens", 0) or 0) + out = int(um.get("output_tokens", 0) or 0) + details = um.get("input_token_details") or {} + cached = int(details.get("cache_read", 0) or 0) + return {"input": inp, "output": out, "cached": min(cached, inp)} + + out_meta = response.llm_output or {} + usage = out_meta.get("usage_metadata") or out_meta.get("token_usage") or {} + if usage: + inp = int(usage.get("input_tokens", usage.get("prompt_token_count", 0)) or 0) + out = int(usage.get("output_tokens", usage.get("candidates_token_count", 0)) or 0) + cached = int(usage.get("cached_content_token_count", 0) or 0) + return {"input": inp, "output": out, "cached": min(cached, inp)} + return None + + +def _normalize_model(name: str) -> str: + """Match Augur's pricing keys to what providers echo back. + + Gemini returns 'models/gemini-2.5-flash'; Anthropic returns a dated id like + 'claude-haiku-4-5-20251001'. Augur's pricing.yaml keys on the bare family + name, so strip the 'models/' prefix and a trailing -YYYYMMDD date. + """ + import re + name = name.strip() + if name.startswith("models/"): + name = name[len("models/"):] + name = re.sub(r"-\d{8}$", "", name) + return name + + +def _extract_model(response: LLMResult) -> str | None: + meta = response.llm_output or {} + name = meta.get("model_name") or meta.get("model") + if name: + return _normalize_model(str(name)) + for gens in response.generations: + for gen in gens: + rmeta = getattr(getattr(gen, "message", None), "response_metadata", {}) or {} + if rmeta.get("model_name"): + return _normalize_model(str(rmeta["model_name"])) + return None + + +class TraceWriter: + """Append-only JSONL writer, thread-safe for concurrent LLM callbacks.""" + + def __init__(self, path: str) -> None: + self._path = path + self._lock = threading.Lock() + self._fh = open(path, "a", encoding="utf-8") + + def write(self, row: dict[str, Any]) -> None: + line = json.dumps(row, ensure_ascii=False) + with self._lock: + self._fh.write(line + "\n") + self._fh.flush() + + def close(self) -> None: + self._fh.close() + + +def _load_scenarios(path: str) -> tuple[int, str, list[dict[str, str]]]: + """Minimal YAML reader for the dogfood scenarios file. + + We avoid a yaml dependency: the file is a tiny, fixed shape (runs, an + optional default model, and a list of {id, input}). Falls back to PyYAML if + the simple parse misses anything. + """ + try: + import yaml # type: ignore + data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + runs = int(data.get("runs", 3)) + model = str(data.get("model", "gemini-2.5-flash")) + scenarios = [{"id": s["id"], "input": s["input"]} for s in data["scenarios"]] + return runs, model, scenarios + except ImportError: + pass + + runs, model, scenarios = 3, "gemini-2.5-flash", [] + cur: dict[str, str] = {} + for raw in Path(path).read_text(encoding="utf-8").splitlines(): + line = raw.rstrip() + if not line or line.lstrip().startswith("#"): + continue + if line.startswith("runs:"): + runs = int(line.split(":", 1)[1].strip()) + elif line.startswith("model:"): + model = line.split(":", 1)[1].strip().strip('"') + elif line.lstrip().startswith("- id:"): + if cur: + scenarios.append(cur) + cur = {"id": line.split("id:", 1)[1].strip().strip('"')} + elif line.lstrip().startswith("input:"): + cur["input"] = line.split("input:", 1)[1].strip().strip('"') + if cur: + scenarios.append(cur) + return runs, model, scenarios + + +# Provider -> default model when --model is not given. Anthropic has far higher +# rate limits than Gemini's free tier, so it is the default for a clean run. +_DEFAULT_MODELS = { + "anthropic": "claude-haiku-4-5-20251001", + "gemini": "gemini-2.5-flash", +} + + +def _build_model(provider: str, model_name: str, callback: Any, rps: float) -> Any: + """Build the chat model the whole graph runs on, with the callback attached. + + The callback rides on the model, so every graph call (supervisor, each + specialist's ReAct loop, synthesizer, LLM judge) writes a trace row. A rate + limiter is optional — needed for Gemini's free tier, rarely for Anthropic. + """ + kwargs: dict[str, Any] = {"temperature": 0.2, "callbacks": [callback], "max_retries": 6} + if rps > 0: + try: + from langchain_core.rate_limiters import InMemoryRateLimiter + kwargs["rate_limiter"] = InMemoryRateLimiter( + requests_per_second=rps, check_every_n_seconds=0.1, max_bucket_size=2) + except Exception as e: # pragma: no cover + print(f"augur_dogfood: rate limiter unavailable ({e})", file=sys.stderr) + + if provider == "anthropic": + from langchain_anthropic import ChatAnthropic + key = os.environ.get("ANTHROPIC_API_KEY") + if not key: + sys.exit("augur_dogfood: ANTHROPIC_API_KEY not set (load CloudOracle's .env).") + # max_tokens bounds the completion; generous enough for routing + synthesis. + return ChatAnthropic(model=model_name, api_key=key, max_tokens=2048, **kwargs) + + from langchain_google_genai import ChatGoogleGenerativeAI + key = os.environ.get("GEMINI_API_KEY") + if not key: + sys.exit("augur_dogfood: GEMINI_API_KEY not set (load CloudOracle's .env).") + return ChatGoogleGenerativeAI(model=model_name, google_api_key=key, **kwargs) + + +async def _run(args: argparse.Namespace) -> int: + runs, file_model, scenarios = _load_scenarios(args.scenarios) + if args.runs is not None: + runs = args.runs + model_name = args.model or _DEFAULT_MODELS.get(args.provider, file_model) + + # Build the agent's graph from its own public pieces, but with OUR model so + # we can pick the provider. CloudOracle is imported, never edited. + from insights_agent.config import Settings + from insights_agent.graph.supervisor import build_supervisor_graph + from insights_agent.guardrails.runner import run_guarded + from insights_agent.logging import get_logger, setup + from insights_agent.tools.cloudoracle import CloudOracleClient, build_tools + + # The agent's Settings mandates GEMINI_API_KEY even when we drive it with + # Anthropic (GeminiProvider is simply never constructed here). Satisfy the + # required field with a placeholder so a Gemini key isn't needed for a + # Claude run. + if args.provider != "gemini": + os.environ.setdefault("GEMINI_API_KEY", "unused-for-anthropic-run") + + settings = Settings() # type: ignore[call-arg] # env-populated + setup(level="WARNING", fmt="text") # quiet: the agent's own logs to stderr + get_logger("augur_dogfood") + + writer = TraceWriter(args.out) + callback = AugurTraceCallback(writer, _normalize_model(model_name)) + model = _build_model(args.provider, model_name, callback, args.rps) + + client = CloudOracleClient( + base_url=settings.cloudoracle_base_url, + api_key=settings.cloudoracle_api_key, + timeout_seconds=settings.http_timeout_seconds, + ) + tools = list(build_tools(client)) # RAG knowledge tool omitted (DATABASE_URL ignored) + graph = build_supervisor_graph(model, tools, settings.run_limits) + # The agent's LLM judge sends a system-only message, which Anthropic rejects + # ("at least one message is required") though Gemini tolerates it — a + # provider-portability bug the dogfood surfaced in the agent. --no-judge + # skips that layer so a Claude run completes; deterministic grounding still + # runs. (On Gemini the judge works and can stay on.) + judge = model if (settings.enable_llm_judge and not args.no_judge) else None + + total = 0 + print(f"augur_dogfood: provider={args.provider} model={model_name}; " + f"{len(scenarios)} scenario(s) x {runs} run(s) -> {args.out}", file=sys.stderr) + try: + for scenario in scenarios: + sid = scenario["id"] + for i in range(runs): + callback.begin_run(sid, f"{sid}-{i:03d}") + try: + result = await run_guarded( + graph, scenario["input"], + validate=settings.enable_answer_validation, judge_model=judge) + if result.fallback_used: + why = result.error or ( + result.validation.reason if result.validation else "?") + status = f"fallback ({why})" + else: + status = "ok" + except Exception as e: # keep going; a failed run is data too + status = f"error:{e}" + calls = callback._seq + print(f" [{sid} {i+1}/{runs}] {calls} LLM call(s) ({status})", + file=sys.stderr) + total += calls + if args.delay > 0: + await asyncio.sleep(args.delay) + finally: + await client.aclose() + writer.close() + + print(f"augur_dogfood: wrote {total} trace row(s) -> {args.out}", file=sys.stderr) + return 0 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser( + prog="augur_dogfood", + description="Run CloudOracle's Insights Agent under an Augur trace callback.", + ) + p.add_argument("--scenarios", required=True, help="dogfood scenarios.yaml") + p.add_argument("--out", default="trace.jsonl", help="trace file to append to") + p.add_argument("--provider", choices=["anthropic", "gemini"], default="anthropic", + help="which LLM the agent's graph runs on (default: anthropic, " + "for its higher rate limits)") + p.add_argument("--model", default=None, + help="model id (default: provider's default; trace is tagged " + "with the family name to match Augur's pricing.yaml)") + p.add_argument("--runs", type=int, default=None, + help="override repetitions per scenario (default: from the file)") + p.add_argument("--no-judge", action="store_true", + help="skip the LLM judge layer (needed on Anthropic: the " + "agent's judge sends a system-only message Claude rejects)") + p.add_argument("--delay", type=float, default=0.0, + help="seconds to sleep between runs (ease provider rate limits)") + p.add_argument("--rps", type=float, default=0.0, + help="requests/sec cap on the shared chat model (e.g. 0.15 " + "for a 10 RPM free tier). 0 disables the limiter.") + args = p.parse_args(argv) + return asyncio.run(_run(args)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/cloudoracle/budget.yaml b/examples/cloudoracle/budget.yaml new file mode 100644 index 0000000..e6c4987 --- /dev/null +++ b/examples/cloudoracle/budget.yaml @@ -0,0 +1,17 @@ +# Augur — budget for the CloudOracle Insights Agent dogfood +# +# Thresholds the projected unit economics must stay within or `augur gate` fails +# the build. The per-request check uses p95 (the tail is where the cost surprise +# lives). Tuned around a Gemini-2.5-Flash multi-specialist agent; adjust to taste +# after seeing the first projection. + +version: 1 + +# Cap the p95 cost of a single agent question. +max_request_p95_usd: 0.02 + +# Cap the projected cost per tenant per month. +max_tenant_per_month_usd: 200 + +# Cap the projected total monthly bill. +max_monthly_bill_usd: 3000 diff --git a/examples/cloudoracle/pricing-gemini.yaml b/examples/cloudoracle/pricing-gemini.yaml new file mode 100644 index 0000000..1f32e83 --- /dev/null +++ b/examples/cloudoracle/pricing-gemini.yaml @@ -0,0 +1,30 @@ +# Augur — pricing snapshot for the CloudOracle dogfood (Gemini) +# +# Prices are USD per 1,000,000 tokens (per Mtok). DATED SNAPSHOT, copied by hand +# from Google's public Gemini API pricing — APPROXIMATE, verify before quoting. +# Provider prices drift; this exists so `augur gate` can price the agent's trace. +# +# input USD/Mtok for non-cached prompt tokens +# output USD/Mtok for completion tokens (Gemini bills "thinking" here) +# cached_input USD/Mtok for context-cached prompt tokens (a subset of input) + +version: 1 +snapshot_date: "2026-06-22" +currency: USD +unit: per_mtok + +models: + # Chat model the Insights Agent runs on (supervisor, specialists, synthesizer, + # LLM judge). Gemini 2.5 Flash standard tier (<=200k context). + gemini-2.5-flash: + input: 0.30 + output: 2.50 + cached_input: 0.075 + + # RAG embeddings (only billed when DATABASE_URL is set and the run uses the + # knowledge tool). Free tier at the snapshot date; listed so an embeddings call + # is priceable rather than an unknown-model hard error. + text-embedding-004: + input: 0.0 + output: 0.0 + cached_input: 0.0 diff --git a/examples/cloudoracle/scenarios.yaml b/examples/cloudoracle/scenarios.yaml new file mode 100644 index 0000000..0240d5e --- /dev/null +++ b/examples/cloudoracle/scenarios.yaml @@ -0,0 +1,28 @@ +# Augur dogfood scenarios — CloudOracle's Insights Agent +# +# Representative FinOps questions a user asks the agent. They deliberately spread +# across the three specialists (cost_analyst / savings_advisor / concept_expert) +# and include a multi-step one, so the captured trace shows the real supervisor +# fan-out — not a single-call happy path. +# +# Consumed by augur_dogfood.py, which runs each `input` against the agent `runs` +# times and writes a row per LLM call (supervisor + specialists + synthesizer + +# LLM judge) to trace.jsonl. + +# Repetitions per scenario. The agent is non-deterministic (routing, tool loops, +# output length vary), so we want the cost DISTRIBUTION, not one sample. Kept low +# because each run spends real Gemini tokens across several calls. +runs: 3 + +# Default model to tag calls with when the provider doesn't echo one back. +model: gemini-2.5-flash + +scenarios: + - id: cost-breakdown + input: "What are my total cloud costs this month, broken down by service?" + - id: find-savings + input: "Where can I save money? Show me the top optimization recommendations and why." + - id: concept-rightsizing + input: "What is rightsizing, and when should I buy reserved instances instead?" + - id: full-review + input: "Give me a full cost review: current spend, the biggest source of waste, and how to fix it." diff --git a/examples/cloudoracle/traffic.yaml b/examples/cloudoracle/traffic.yaml new file mode 100644 index 0000000..58e3b95 --- /dev/null +++ b/examples/cloudoracle/traffic.yaml @@ -0,0 +1,24 @@ +# Augur — traffic profile for the CloudOracle Insights Agent dogfood +# +# Production assumptions for projecting the observed per-run cost to scale. These +# are a FORECAST (yours), not measured by Augur. Modelled as an internal FinOps +# assistant: an ops/eng org where engineers ask it cost questions through the day. + +version: 1 + +# Projected population: ~500 engineers, a handful of agent questions each per day. +users: 500 +requests_per_user_per_day: 4 + +# Multi-tenant SaaS framing: offered to 20 customer orgs. +tenants: 20 + +days_per_month: 30 + +# Mix of question types in production. Concept questions are cheap and common; +# the full review is the heavy multi-specialist path. Weights need not sum to 1. +scenario_mix: + cost-breakdown: 0.30 + find-savings: 0.25 + concept-rightsizing: 0.30 + full-review: 0.15 diff --git a/sidecar/README.md b/sidecar/README.md new file mode 100644 index 0000000..ea45453 --- /dev/null +++ b/sidecar/README.md @@ -0,0 +1,222 @@ +# augur-predict — the predictive output-length sidecar + +> Estimate an agent's cost for inputs you **haven't run**, by learning completion +> length from a trace Augur already recorded. + +This is the one piece of Augur the [SPEC](../SPEC.md) deliberately writes in +Python instead of Go (Hito 5): a small **predictive output-length model**. It is +optional — delete this directory and Augur's pure-Go v1 is untouched. + +## Why it exists + +Augur's core measures *real* cost by running your agent through a recording +proxy. But running spends tokens, and you can't run every hypothetical. The one +quantity you genuinely cannot guess for an agent is **completion length** — and, +per Augur's thesis, it's where the cost surprise lives (long outputs in the +tail). So we learn output length from a recorded trace, then use it to project +cost for inputs we never executed: a bigger prompt, a context-growth scenario, +more runs. + +This is the PreflightLLMCost direction — an honest linear baseline, not a deep +model — and it's where a second language earns its place: the analytical fit +(numpy OLS, residual spread, prediction bands) is more natural in Python, while +the systems core stays in Go. + +## Loose coupling: the trace file is the only contract + +The sidecar never calls the Go binary and the Go binary never calls the sidecar. +They share exactly one thing: the JSONL cost-trace schema. `fit` reads a recorded +trace; `emit-trace` writes a synthetic one that `augur aggregate | project | +gate` price with no idea a model — not a proxy — produced it. No RPC, no shared +library, no import in either direction. + +``` +trace.jsonl ──► augur-predict fit ──► model.json + │ + augur-predict emit-trace --input-scale 1.5 + │ + ▼ + predicted-trace.jsonl ──► augur aggregate | project | gate +``` + +## Install + +Needs Python ≥ 3.10 and numpy. No install required to run: + +```sh +cd sidecar +python -m augur_predict --help # run in place +# or: +pip install -e . # exposes the `augur-predict` command +pip install -e '.[dev]' && pytest # with the test suite +``` + +## Commands + +### `fit` — learn the model from a recorded trace + +```sh +python -m augur_predict fit --trace trace.jsonl --out model.json +``` + +Fits, **per billed model**, `output_tokens ≈ intercept + slope · input_tokens`, +and captures each observed `(scenario, run)` as a *run template* (its call-graph: +which models, what input sizes, in what order). Only successful calls feed the +fit; a failed-but-billed call (a 429 that burned input and produced nothing) +would distort an output-length fit, so it's excluded — but it still contributes +its structure to the templates, because it's part of the call graph `emit-trace` +replays. + +`--dist` chooses how the output spread is modelled: + +| `--dist` | what it fits | use it when | +|---|---|---| +| `gaussian` *(default)* | OLS mean line + a **symmetric** residual band | the honest baseline; output is roughly symmetric around the trend | +| `quantile` | a grid of **conditional quantiles** by quantile regression | you care about the **tail** — output is right-skewed, and the p95 the gate uses is what matters | + +See [Tail accuracy](#tail-accuracy-quantile-mode--run-correlation) below for why +the default isn't always enough. + +### `report` — is the fit worth trusting? + +```sh +python -m augur_predict report --model model.json +``` + +``` +Output-length model (source: trace.jsonl) + 100 calls, 50 run templates, 2 model(s) + + model n method out~in R2 ±resid + ---------------------------------------------------------- + gpt-4o 50 ols 65+0.299·in 0.97 30 + gpt-4o-mini 50 ols 22+0.017·in 0.60 5 +``` + +Honesty is a first-class output. With too few points or no spread in the inputs, +a slope is noise: the model degrades to predicting the **mean** output +(`method=mean`) and says so. A low R² is flagged. This mirrors the Go side +reporting confidence intervals instead of bare point estimates. + +A quantile model reports the **median and p95 lines** side by side (the two the +gate cares about) plus the Koenker–Machado pseudo-R¹, and flags any model that +fell back to the empirical marginal quantiles (too few points to regress). + +### `predict` — one input size + +```sh +python -m augur_predict predict --model model.json \ + --model-name gpt-4o --input-tokens 3000 --price-out 10.0 +``` + +``` +model=gpt-4o input_tokens=3000.0 + predicted output tokens: 962 (~95% band 903-1021, method=ols) + est. output cost @ $10.0/Mtok: $0.009622 (p95 $0.010214) +``` + +A gaussian model prints a point prediction and a symmetric `~95%` band from the +residual spread; a quantile model prints the **p50 and p95** directly. `--price-out` +is optional and prices only the *completion* side — the unknown the sidecar +models; the prompt cost is already exact from the input you supplied. Full +per-call pricing lives in the Go `aggregate` stage, which `emit-trace` feeds. + +### `emit-trace` — a predicted trace for the Go gate + +```sh +python -m augur_predict emit-trace --model model.json \ + --out predicted-trace.jsonl --runs 50 --input-scale 1.5 --seed 7 +``` + +Resamples the run templates, scales every prompt by `--input-scale`, and draws +each call's output from the fit — sampling the *spread*, not just the mean, so +the predicted distribution (and the p95 the gate cares about) stays honest +instead of collapsing to the average. A gaussian model samples its symmetric +band; a quantile model inverts its conditional CDF, preserving the skew. The +result is an ordinary trace: + +```sh +augur aggregate --trace predicted-trace.jsonl +augur gate --trace predicted-trace.jsonl --traffic traffic.yaml --budget budget.yaml +``` + +`--input-scale` is the predictive analogue of the Go `--context-growth` knob, but +it does more: it feeds the larger prompt *back through the output model*, so a +bigger ask predicts a longer answer — not just a costlier prompt. +`--run-correlation ρ` (0–1, default 0) shares a *verbosity* draw across a run's +calls (see below). Emission is **deterministic per `--seed`** (same seed → same +trace), carrying the Go side's record/replay determinism into the predictive +path. + +## Worked example + +```sh +augur run --scenarios scenarios.yaml --upstream https://api.openai.com # record once +python -m augur_predict fit --trace trace.jsonl --out model.json +# "What if prompts grow 1.5× as conversations lengthen?" — no agent re-run: +python -m augur_predict emit-trace --model model.json --out grown.jsonl --input-scale 1.5 +augur gate --trace grown.jsonl --traffic traffic.yaml --budget budget.yaml +``` + +In Augur's own test of this chain, the recorded trace passed the budget +(~$15.3k/month) while the 1.5× predicted trace failed it (~$22.4k/month, over a +$20k cap) — the cost regression caught *before* anyone ran the bigger workload. + +## Tail accuracy: quantile mode & run correlation + +The whole tool gates on **p95**, because the cost surprise lives in the tail. The +gaussian default has two assumptions that *understate* that tail, and `quantile` +mode plus `--run-correlation` are the opt-in fixes: + +1. **Skew.** A symmetric Gaussian band is thin-tailed; real completion lengths + are right-skewed (rambles, near-`max_tokens` runs), so the gaussian p95 sits + too low. **`--dist quantile`** fits the conditional quantiles directly — it + targets the p95 the gate uses and represents the asymmetry instead of a + symmetric ± band. Quantile regression is solved exactly as a linear program + (`scipy.optimize.linprog`, HiGHS); crossing quantiles are repaired by + monotone rearrangement at prediction time. + +2. **Run-level correlation.** Sampling each call independently understates the + variance of the *per-run* total — but the gate aggregates to per-run cost and + then takes its p95, so that spread is exactly what's gated. Real runs + correlate (a verbose run is verbose throughout). **`--run-correlation ρ`** + models it with a Gaussian copula: each run draws one latent `z_run`, each call + blends it as `z = √ρ·z_run + √(1-ρ)·z_idio`. ρ=0 is independent (the default); + higher ρ widens the per-run p95. + +In a heavy-tailed test against the recorded ground truth (p95 ≈ \$0.0130/req), +the **quantile** projection (\$0.0126, 95% CI [\$0.0117, \$0.0138]) covered the +true value, while the **gaussian** projection (\$0.0121, CI [\$0.0115, \$0.0124]) +sat lower with a falsely tight interval that *missed* it — the precise failure +mode the modes exist to fix. + +Both stay opt-in: `gaussian`/`ρ=0` is the conservative default so you never get +unmodelled assumptions you didn't ask for. + +## Honest limitations + +- **Still a one-feature model.** Both modes use only `input_tokens` as the + predictor. When prompt length doesn't drive output length (low R²/R¹), the fit + falls back — to the mean (gaussian) or the marginal quantiles (quantile) — and + the report flags it. Don't read more into a prediction than its goodness-of-fit + supports. +- **It can't invent structure it never saw.** `emit-trace` resamples observed run + templates; it cannot predict a call path the agent never took in the recorded + trace. Representativeness of the original run still bounds everything. +- **The tail is data-bound.** A p95 from a few dozen runs is intrinsically noisy + no matter the model, which is why the grid stops at 0.95 (a p99 would be false + precision) and why fit quality is always reported. Quantile mode sharpens the + *shape*; it does not manufacture tail data you didn't record. + +## Tests + +```sh +cd sidecar && pytest -q +``` + +57 cases covering the trace round-trip and Go-schema contract; OLS recovery of a +known slope/intercept and the mean fallback; quantile-mode ordering (p05 ≤ p50 ≤ +p95), asymmetric bands on skewed data, quantile p95 > gaussian p95, and the +empirical-quantile fallback; that `--run-correlation` widens the per-run total +variance; and `emit-trace` determinism + validity (cached ≤ input, non-negative +tokens) so emitted rows never fail the Go cost validator. diff --git a/sidecar/augur_predict/__init__.py b/sidecar/augur_predict/__init__.py new file mode 100644 index 0000000..d1b2127 --- /dev/null +++ b/sidecar/augur_predict/__init__.py @@ -0,0 +1,24 @@ +"""augur_predict — the optional Python sidecar for Augur. + +Augur's v1 is pure Go: it measures an agent's real token usage by running it +through a recording proxy, then projects the bill at production scale. This +sidecar is the one analytical piece the SPEC deliberately defers to a second +language (Hito 5): a *predictive output-length model*. + +The premise is PreflightLLMCost's: completion length — not prompt length — is +what you cannot guess for an agent, and it is the dominant cost driver in the +tail. So we learn it from a trace Augur already recorded, then use that model to +estimate cost for inputs we have NOT run, without spending tokens. + +Coupling to the Go core is deliberately loose: this package reads the same +JSONL cost-trace the proxy writes and (for ``emit-trace``) writes one back. No +RPC, no shared library, no import in either direction — just the trace file as +the contract. You can delete this directory and Augur's v1 is untouched. +""" + +__version__ = "0.1.0" + +from .model import Model, ModelFit, fit +from .trace import Record, load_trace + +__all__ = ["Model", "ModelFit", "fit", "Record", "load_trace", "__version__"] diff --git a/sidecar/augur_predict/__main__.py b/sidecar/augur_predict/__main__.py new file mode 100644 index 0000000..7b08493 --- /dev/null +++ b/sidecar/augur_predict/__main__.py @@ -0,0 +1,8 @@ +"""Enable ``python -m augur_predict ...`` as the no-install entry point.""" + +import sys + +from .cli import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sidecar/augur_predict/cli.py b/sidecar/augur_predict/cli.py new file mode 100644 index 0000000..2459376 --- /dev/null +++ b/sidecar/augur_predict/cli.py @@ -0,0 +1,238 @@ +"""Command-line interface for the sidecar. + +Four subcommands, each a thin shell over the library so the logic stays testable +without argv: + + fit trace.jsonl -> model.json + report model.json -> human-readable fit quality + predict model.json + inputs -> predicted output length (+ optional cost) + emit-trace model.json -> predicted trace.jsonl (for the Go gate) + +The contract with Augur's Go core is files only: this never calls the Go binary +and the Go binary never calls this. ``emit-trace`` writes a trace; everything +else reads one. +""" + +from __future__ import annotations + +import argparse +import sys +from typing import List, Optional + +from . import __version__ +from .emit import emit +from .model import Model, fit +from .trace import load_trace, write_trace + + +def _cmd_fit(args: argparse.Namespace) -> int: + records = load_trace(args.trace) + if not records: + print(f"fit: trace {args.trace!r} is empty - nothing to learn", file=sys.stderr) + return 1 + model = fit(records, dist=args.dist) + model.source = args.trace + model.save(args.out) + print(f"fit: learned {len(model.fits)} model(s) [{args.dist}] from " + f"{model.n_records} call(s) across {len(model.templates)} run(s) -> {args.out}") + return 0 + + +def _cmd_report(args: argparse.Namespace) -> int: + model = Model.load(args.model) + print(_render_report(model)) + return 0 + + +def _render_report(model: Model) -> str: + if model.dist == "quantile": + return _render_report_quantile(model) + lines: List[str] = [] + lines.append(f"Output-length model [gaussian] (source: {model.source or 'unknown'})") + lines.append(f" {model.n_records} calls, {len(model.templates)} run templates, " + f"{len(model.fits)} model(s)") + lines.append("") + header = f" {'model':<28} {'n':>5} {'method':>7} {'out~in':>16} {'R2':>6} {'+-resid':>8}" + lines.append(header) + lines.append(" " + "-" * (len(header) - 2)) + for name in sorted(model.fits): + f = model.fits[name] + if f.method == "ols": + rel = f"{f.intercept:.0f}+{f.slope:.3f}*in" + r2 = f"{f.r2:.2f}" + else: + rel = f"~{f.output_mean:.0f}" + r2 = " -- " + lines.append(f" {name:<28} {f.n:>5} {f.method:>7} {rel:>16} {r2:>6} {f.resid_std:>8.0f}") + lines.append("") + # An honest read on whether the fit is worth trusting - the same spirit as + # the Go side surfacing CIs instead of bare point estimates. + weak = [name for name, f in model.fits.items() + if f.method == "mean" or f.r2 < 0.3] + if weak: + lines.append(" note: weak/absent input->output signal for: " + + ", ".join(sorted(weak))) + lines.append(" predictions fall back to the observed mean output; " + "treat them as rough.") + return "\n".join(lines) + + +def _render_report_quantile(model: Model) -> str: + """Report for a quantile model: the median and p95 lines side by side. + + These are the two the gate cares about — the median is the typical call, the + p95 is the tail the budget is checked against — so we surface both rather + than a single point and a symmetric band. + """ + lines: List[str] = [] + lines.append(f"Output-length model [quantile] (source: {model.source or 'unknown'})") + lines.append(f" {model.n_records} calls, {len(model.templates)} run templates, " + f"{len(model.fits)} model(s)") + lines.append("") + header = (f" {'model':<24} {'n':>5} {'method':>17} " + f"{'p50 out~in':>16} {'p95 out~in':>16} {'R1':>6}") + lines.append(header) + lines.append(" " + "-" * (len(header) - 2)) + for name in sorted(model.fits): + f = model.fits[name] + q = {ql.tau: ql for ql in f.quantiles} + p50 = q.get(0.5) + p95 = q.get(0.95) or (f.quantiles[-1] if f.quantiles else None) + p50s = f"{p50.intercept:.0f}+{p50.slope:.3f}*in" if p50 else "--" + p95s = f"{p95.intercept:.0f}+{p95.slope:.3f}*in" if p95 else "--" + r1 = f"{f.r2:.2f}" if f.method == "quantile_reg" else " -- " + lines.append(f" {name:<24} {f.n:>5} {f.method:>17} {p50s:>16} {p95s:>16} {r1:>6}") + lines.append("") + weak = [name for name, f in model.fits.items() + if f.method == "empirical_quantile"] + if weak: + lines.append(" note: too few points for quantile regression on: " + + ", ".join(sorted(weak))) + lines.append(" these use the observed marginal output quantiles " + "(flat in input) - still skewed, but no input signal.") + lines.append(" R1 is the Koenker-Machado pseudo-R2 for the median fit " + "(QR's goodness-of-fit).") + return "\n".join(lines) + + +def _cmd_predict(args: argparse.Namespace) -> int: + model = Model.load(args.model) + f = model.fit_for(args.model_name) + if f is None: + avail = ", ".join(sorted(model.fits)) or "(none)" + print(f"predict: no fit for model {args.model_name!r}; trace covered: {avail}", + file=sys.stderr) + return 1 + + x = args.input_tokens + print(f"model={args.model_name} input_tokens={x}") + if f.dist == "quantile": + # The two the gate cares about: typical call and tail. + p50 = f.quantile_at(x, 0.5) + p95 = f.quantile_at(x, 0.95) + print(f" predicted output tokens: p50 {p50:.0f}, p95 {p95:.0f} " + f"(method={f.method})") + mid, hi = p50, p95 + else: + mid = f.predict(x) + lo, hi = f.band(x) + print(f" predicted output tokens: {mid:.0f} " + f"(~95% band {lo:.0f}-{hi:.0f}, method={f.method})") + + if args.price_out is not None: + # Optional dollar estimate. We price only the completion side here: the + # sidecar's job is the unknown (output length); prompt cost is already + # known exactly from the input you supplied. Full per-call pricing lives + # in the Go `aggregate` stage, which emit-trace feeds. + cost_mid = mid / 1_000_000 * args.price_out + cost_hi = hi / 1_000_000 * args.price_out + print(f" est. output cost @ ${args.price_out}/Mtok: " + f"${cost_mid:.6f} (p95 ${cost_hi:.6f})") + return 0 + + +def _cmd_emit_trace(args: argparse.Namespace) -> int: + model = Model.load(args.model) + try: + records = emit( + model, + runs=args.runs, + input_scale=args.input_scale, + seed=args.seed, + scenario_filter=args.scenario, + run_prefix=args.run_prefix, + run_correlation=args.run_correlation, + ) + except ValueError as e: + print(f"emit-trace: {e}", file=sys.stderr) + return 1 + + n = write_trace(args.out, records) + print(f"emit-trace: wrote {n} predicted call(s) over {args.runs} run(s) " + f"(input_scale={args.input_scale}, run_correlation={args.run_correlation}, " + f"seed={args.seed}) -> {args.out}") + print(f" feed it to the Go gate, e.g.: augur aggregate --trace {args.out} | ...") + return 0 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="augur-predict", + description="Augur's predictive output-length sidecar: learn completion " + "length from a recorded trace, then estimate cost for inputs " + "you have not run.", + ) + p.add_argument("--version", action="version", version=f"augur-predict {__version__}") + sub = p.add_subparsers(dest="command", required=True) + + pf = sub.add_parser("fit", help="learn an output-length model from a trace") + pf.add_argument("--trace", required=True, help="recorded JSONL trace to learn from") + pf.add_argument("--out", default="model.json", help="model artifact to write") + pf.add_argument("--dist", choices=["gaussian", "quantile"], default="gaussian", + help="gaussian: OLS mean + symmetric spread (default, honest " + "baseline). quantile: conditional-quantile regression - " + "skewed, targets the p95 the gate uses.") + pf.set_defaults(func=_cmd_fit) + + pr = sub.add_parser("report", help="print fit quality for a model artifact") + pr.add_argument("--model", default="model.json", help="model artifact to read") + pr.set_defaults(func=_cmd_report) + + pp = sub.add_parser("predict", help="predict output length for one input size") + pp.add_argument("--model", default="model.json", help="model artifact to read") + pp.add_argument("--model-name", required=True, dest="model_name", + help="billed model to predict for (must appear in the trace)") + pp.add_argument("--input-tokens", required=True, type=float, dest="input_tokens", + help="prompt size to predict the completion length for") + pp.add_argument("--price-out", type=float, default=None, dest="price_out", + help="optional $/Mtok for output, to print an output-cost estimate") + pp.set_defaults(func=_cmd_predict) + + pe = sub.add_parser("emit-trace", + help="synthesise a predicted trace for the Go gate") + pe.add_argument("--model", default="model.json", help="model artifact to read") + pe.add_argument("--out", default="predicted-trace.jsonl", help="trace file to write") + pe.add_argument("--runs", type=int, default=20, + help="number of synthetic runs to generate") + pe.add_argument("--input-scale", type=float, default=1.0, dest="input_scale", + help="multiply every prompt size (and re-predict output); " + "the predictive analogue of --context-growth") + pe.add_argument("--scenario", default=None, + help="restrict to one scenario id from the templates") + pe.add_argument("--seed", type=int, default=0, + help="RNG seed; same seed reproduces the same trace") + pe.add_argument("--run-prefix", default="pred", dest="run_prefix", + help="prefix for synthetic run ids") + pe.add_argument("--run-correlation", type=float, default=0.0, dest="run_correlation", + help="0..1 verbosity shared across a run's calls (Gaussian " + "copula). 0 = independent (default); higher widens the " + "per-run cost spread the gate's p95 sees.") + pe.set_defaults(func=_cmd_emit_trace) + + return p + + +def main(argv: Optional[List[str]] = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + return args.func(args) diff --git a/sidecar/augur_predict/emit.py b/sidecar/augur_predict/emit.py new file mode 100644 index 0000000..f0ccba8 --- /dev/null +++ b/sidecar/augur_predict/emit.py @@ -0,0 +1,122 @@ +"""Synthesising a predicted trace from a fitted model. + +This is how the sidecar pays for itself: given a model learned from one recorded +run, produce a trace for inputs you have NOT executed — a bigger prompt, more +runs, a context-growth scenario — without spending a token. The synthetic trace +is ordinary JSONL, so ``augur aggregate | project | gate`` consume it with no +knowledge that a model, not a proxy, wrote it. That is the loose coupling the +SPEC asks for: the file is the whole interface. + +**Run-level correlation.** A naive synthesis samples every call independently, +which understates the variance of the per-run total — yet the gate aggregates to +per-run cost and then takes its p95, so that run-level spread is exactly what is +gated. Real runs are correlated: a verbose run is verbose across all its calls; a +retry storm correlates. We model this with a Gaussian copula. Each run draws one +latent ``z_run``; each call blends it with its own idiosyncratic draw: + + z_call = √ρ · z_run + √(1-ρ) · z_idio (still standard normal) + +``--run-correlation ρ`` ranges from 0 (independent, the original behaviour) to 1 +(every call in a run shares one percentile). The same ``z`` drives both modes: +gaussian uses it as the standardised residual, quantile maps it to u = Φ(z) and +inverts the conditional CDF — so the *skew* of the fitted tail survives. + +Determinism is deliberate. emit-trace seeds a numpy RNG so the same model and +flags reproduce the same trace — the record/replay ethos of the Go side carried +into the predictive path. Vary ``--seed`` to draw a different sample. +""" + +from __future__ import annotations + +import math +from typing import List, Optional + +import numpy as np + +from .model import Model +from .trace import Record + +# A fixed wall-clock stamp for synthetic rows. They were not observed at any real +# time, and labelling them so keeps a predicted trace honest and reproducible +# (no clock dependence, matching the Go replay path's stable run-ids). +_SYNTHETIC_TS = "1970-01-01T00:00:00Z" + + +def emit( + model: Model, + runs: int, + input_scale: float = 1.0, + seed: int = 0, + scenario_filter: Optional[str] = None, + run_prefix: str = "pred", + run_correlation: float = 0.0, +) -> List[Record]: + """Generate ``runs`` synthetic runs by resampling and rescaling templates. + + For each synthetic run we take an observed template (cycling through them in + order, so the scenario mix is preserved), scale every call's prompt by + ``input_scale``, then predict each call's output from the model at a latent + ``z`` that carries the run-level correlation (see module docstring). Sampling + the spread — not just the mean — is what keeps the predicted *distribution* + (and therefore the p95 the gate cares about) honest rather than collapsing + every run to the average. + + ``input_scale`` is the predictive analogue of the Go ``--context-growth`` + knob, but it does more: it also feeds the larger prompt back through the + output model, so a bigger ask predicts a longer answer instead of only a + costlier prompt. + """ + if runs <= 0: + return [] + if input_scale <= 0: + raise ValueError("input_scale must be positive") + if not (0.0 <= run_correlation <= 1.0): + raise ValueError("run_correlation must be in [0, 1]") + + templates = model.templates + if scenario_filter is not None: + templates = [t for t in templates if t.scenario_id == scenario_filter] + if not templates: + raise ValueError("no run templates to sample from" + + (f" for scenario {scenario_filter!r}" if scenario_filter else "")) + + rng = np.random.default_rng(seed) + a_shared = math.sqrt(run_correlation) + a_idio = math.sqrt(1.0 - run_correlation) + out: List[Record] = [] + width = max(4, len(str(runs))) + + for i in range(runs): + template = templates[i % len(templates)] + run_id = f"{run_prefix}-{i:0{width}d}" + z_run = float(rng.standard_normal()) + for call in template.calls: + fit = model.fit_for(call.model) + scaled_input = int(round(call.input_tokens * input_scale)) + # Cached tokens are a subset of input; scale them with it and never + # let them exceed the (possibly rounded) input, which would make the + # row fail cost.Usage.Validate on the Go side. + scaled_cached = min(scaled_input, int(round(call.cached_tokens * input_scale))) + + z = a_shared * z_run + a_idio * float(rng.standard_normal()) + if fit is None: + # A model present in the templates but with no successful calls + # to learn from: nothing to predict, emit a conservative zero. + predicted = 0 + else: + predicted = fit.sample(scaled_input, z) + + out.append(Record( + scenario_id=template.scenario_id, + run_id=run_id, + seq=call.seq, + model=call.model, + input_tokens=scaled_input, + output_tokens=predicted, + cached_tokens=scaled_cached, + latency_ms=0, + ts=_SYNTHETIC_TS, + endpoint=call.endpoint, + status=200, + )) + return out diff --git a/sidecar/augur_predict/model.py b/sidecar/augur_predict/model.py new file mode 100644 index 0000000..1eb0dfc --- /dev/null +++ b/sidecar/augur_predict/model.py @@ -0,0 +1,513 @@ +"""The predictive output-length model. + +Two distribution modes, one shared structure: + +* **gaussian** (default) — per model, ``output ~ a + b·input_tokens`` via OLS, + with residuals modelled as a symmetric Gaussian spread. The honest linear + baseline: the PreflightLLMCost direction, fit per billed model because + verbosity differs sharply across models. With too few points or no input + spread it degrades to predicting the mean output (``method == "mean"``) and + says so. + +* **quantile** — per model, a *grid* of conditional quantiles fit directly by + quantile regression (``output_τ ~ a_τ + b_τ·input``). This targets what the + gate actually uses (a conditional quantile, not the mean) and represents a + skewed, right-tailed output distribution instead of a symmetric band — the + cost surprise lives in that tail. The grid is a discretised conditional CDF: + ``emit-trace`` samples it by inverse transform. With too few points it falls + back to the *empirical marginal quantiles* of the output, which is still + asymmetric (and so more honest than the gaussian mean fallback). + +Either way we also capture each observed ``(scenario, run)`` as a *run template* +— its call-graph (which models, what input sizes, in what order) — so +``emit-trace`` can resample real run shapes rather than inventing structure. + +The model never imports the Go code; it only reads Records and writes a JSON +artifact. numpy does the OLS and the sampling; scipy's ``linprog`` solves the +quantile-regression LP (quantile mode only). +""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence + +import numpy as np + +from .trace import Record + +# z for an approximate 95% prediction band from the residual spread. The same +# 1.96 the Go projection uses for its normal-approximation intervals. +_Z95 = 1.959963984540054 + +# Below this many successful calls for a model we do not trust an OLS slope and +# fall back to predicting the mean output. +_MIN_FIT_POINTS = 8 + +# Quantile regression at the extreme taus needs more data than a median fit, so +# its fallback threshold is higher: below this we use empirical marginal +# quantiles instead of regressing. +_MIN_QR_POINTS = 12 + +# The conditional-quantile grid fit in quantile mode. 0.5 must be present (it is +# the point prediction); 0.05/0.95 bound the band the gate cares about. We stop +# at 0.95 — a p99 from a few dozen runs would be false precision. +DEFAULT_TAUS: tuple[float, ...] = (0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95) + + +def _phi(z: float) -> float: + """Standard-normal CDF, for turning a copula latent z into a uniform u.""" + return 0.5 * (1.0 + math.erf(z / math.sqrt(2.0))) + + +@dataclass +class QuantileLine: + """One conditional-quantile line: output_tau ≈ intercept + slope·input.""" + + tau: float + intercept: float + slope: float + + def at(self, input_tokens: float) -> float: + return self.intercept + self.slope * float(input_tokens) + + def to_json(self) -> dict: + return {"tau": self.tau, "intercept": self.intercept, "slope": self.slope} + + @classmethod + def from_json(cls, obj: dict) -> "QuantileLine": + return cls(tau=float(obj["tau"]), intercept=float(obj["intercept"]), + slope=float(obj["slope"])) + + +@dataclass +class ModelFit: + """The fitted output-length relationship for a single billed model. + + ``dist`` selects how the fields are read: + + * ``gaussian`` — ``intercept``/``slope`` are the OLS mean line; ``resid_std`` + is the symmetric prediction spread; ``method`` is ``"ols"`` or ``"mean"``. + * ``quantile`` — ``quantiles`` holds the conditional-quantile grid; + ``intercept``/``slope`` mirror the median line for convenience; ``method`` + is ``"quantile_reg"`` or ``"empirical_quantile"``. ``resid_std`` is the + spread of residuals about the median, kept only for the report. + """ + + model: str + n: int + method: str + intercept: float + slope: float + r2: float + resid_std: float + output_mean: float + output_std: float + input_mean: float + input_min: float + input_max: float + dist: str = "gaussian" + quantiles: List[QuantileLine] = field(default_factory=list) + + # --- prediction ------------------------------------------------------- + + def predict(self, input_tokens: float) -> float: + """Point prediction (clipped at zero). + + gaussian: the OLS mean. quantile: the median (tau=0.5) of the grid. + """ + if self.dist == "quantile": + return self.quantile_at(input_tokens, 0.5) + return max(0.0, self.intercept + self.slope * float(input_tokens)) + + def quantile_at(self, input_tokens: float, tau: float) -> float: + """The tau-quantile of output at this input (quantile mode only). + + Evaluates every grid line at the input, enforces monotonicity across + taus (independently-fit quantiles can cross), then interpolates at tau. + Outside the fitted tau range np.interp clamps to the ends — we do not + extrapolate a tail we never fit. + """ + if not self.quantiles: + # gaussian fit asked for a quantile: use the normal approximation. + from_z = _z_for_tau(tau) + return max(0.0, self.predict(input_tokens) + from_z * self.resid_std) + taus = np.array([q.tau for q in self.quantiles]) + vals = np.maximum.accumulate( + np.array([q.at(input_tokens) for q in self.quantiles])) + return max(0.0, float(np.interp(tau, taus, vals))) + + def band(self, input_tokens: float, z: float = _Z95) -> tuple[float, float]: + """Approximate ~95% band (lo, hi), both clipped at zero. + + gaussian: mean ± z·resid_std. quantile: the 0.05 and 0.95 grid lines — + an asymmetric band straight from the fitted tails. + """ + if self.dist == "quantile" and self.quantiles: + lo = self.quantile_at(input_tokens, self.quantiles[0].tau) + hi = self.quantile_at(input_tokens, self.quantiles[-1].tau) + return lo, hi + mid = self.predict(input_tokens) + return max(0.0, mid - z * self.resid_std), max(0.0, mid + z * self.resid_std) + + def sample(self, input_tokens: float, z: float) -> int: + """Draw one output length given a standard-normal latent ``z``. + + ``z`` carries the run-level correlation (see emit.py): the same z shared + across a run's calls makes a verbose run verbose throughout. gaussian + uses z directly as the standardised residual; quantile maps z→u=Φ(z) and + inverts the conditional CDF, so the *skew* of the fitted grid is + preserved rather than collapsed to a symmetric band. + """ + if self.dist == "quantile": + u = _phi(z) + return int(round(self.quantile_at(input_tokens, u))) + val = self.predict(input_tokens) + z * self.resid_std + return int(round(max(0.0, val))) + + # --- serialization ---------------------------------------------------- + + def to_json(self) -> dict: + d = { + "model": self.model, + "n": self.n, + "method": self.method, + "dist": self.dist, + "intercept": self.intercept, + "slope": self.slope, + "r2": self.r2, + "resid_std": self.resid_std, + "output_mean": self.output_mean, + "output_std": self.output_std, + "input_mean": self.input_mean, + "input_min": self.input_min, + "input_max": self.input_max, + } + if self.quantiles: + d["quantiles"] = [q.to_json() for q in self.quantiles] + return d + + @classmethod + def from_json(cls, obj: dict) -> "ModelFit": + return cls( + model=obj["model"], n=int(obj["n"]), method=obj["method"], + intercept=float(obj["intercept"]), slope=float(obj["slope"]), + r2=float(obj["r2"]), resid_std=float(obj["resid_std"]), + output_mean=float(obj["output_mean"]), output_std=float(obj["output_std"]), + input_mean=float(obj["input_mean"]), input_min=float(obj["input_min"]), + input_max=float(obj["input_max"]), + dist=obj.get("dist", "gaussian"), + quantiles=[QuantileLine.from_json(q) for q in obj.get("quantiles", [])], + ) + + +def _z_for_tau(tau: float) -> float: + """Inverse standard-normal CDF via a rational approximation (Acklam). + + Used only when a gaussian fit is asked for a quantile (the report's p95 on a + gaussian model). Accurate to ~1e-9 over (0,1), and avoids a scipy import on + the gaussian path. + """ + if tau <= 0.0: + return -math.inf + if tau >= 1.0: + return math.inf + a = [-3.969683028665376e+01, 2.209460984245205e+02, -2.759285104469687e+02, + 1.383577518672690e+02, -3.066479806614716e+01, 2.506628277459239e+00] + b = [-5.447609879822406e+01, 1.615858368580409e+02, -1.556989798598866e+02, + 6.680131188771972e+01, -1.328068155288572e+01] + c = [-7.784894002430293e-03, -3.223964580411365e-01, -2.400758277161838e+00, + -2.549732539343734e+00, 4.374664141464968e+00, 2.938163982698783e+00] + d = [7.784695709041462e-03, 3.224671290700398e-01, 2.445134137142996e+00, + 3.754408661907416e+00] + plow, phigh = 0.02425, 1 - 0.02425 + if tau < plow: + q = math.sqrt(-2 * math.log(tau)) + return (((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / \ + ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1) + if tau > phigh: + q = math.sqrt(-2 * math.log(1 - tau)) + return -(((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / \ + ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1) + q = tau - 0.5 + r = q * q + return (((((a[0]*r+a[1])*r+a[2])*r+a[3])*r+a[4])*r+a[5])*q / \ + (((((b[0]*r+b[1])*r+b[2])*r+b[3])*r+b[4])*r+1) + + +@dataclass +class Call: + """One call within a run template: the structure emit-trace replays.""" + + seq: int + model: str + input_tokens: int + cached_tokens: int + endpoint: str = "" + + def to_json(self) -> dict: + d = { + "seq": self.seq, + "model": self.model, + "input_tokens": self.input_tokens, + "cached_tokens": self.cached_tokens, + } + if self.endpoint: + d["endpoint"] = self.endpoint + return d + + @classmethod + def from_json(cls, obj: dict) -> "Call": + return cls( + seq=int(obj["seq"]), + model=obj["model"], + input_tokens=int(obj["input_tokens"]), + cached_tokens=int(obj.get("cached_tokens", 0)), + endpoint=obj.get("endpoint", ""), + ) + + +@dataclass +class RunTemplate: + """The observed call-graph of one (scenario, run): its sequence of calls.""" + + scenario_id: str + calls: List[Call] + + def to_json(self) -> dict: + return { + "scenario_id": self.scenario_id, + "calls": [c.to_json() for c in self.calls], + } + + @classmethod + def from_json(cls, obj: dict) -> "RunTemplate": + return cls( + scenario_id=obj.get("scenario_id", ""), + calls=[Call.from_json(c) for c in obj.get("calls", [])], + ) + + +@dataclass +class Model: + """The full artifact: per-model fits plus the observed run templates.""" + + version: int + n_records: int + fits: Dict[str, ModelFit] + templates: List[RunTemplate] + dist: str = "gaussian" + source: str = "" + + def to_json(self) -> dict: + return { + "version": self.version, + "source": self.source, + "dist": self.dist, + "n_records": self.n_records, + "models": {m: f.to_json() for m, f in self.fits.items()}, + "run_templates": [t.to_json() for t in self.templates], + } + + @classmethod + def from_json(cls, obj: dict) -> "Model": + return cls( + version=int(obj.get("version", 1)), + source=obj.get("source", ""), + dist=obj.get("dist", "gaussian"), + n_records=int(obj.get("n_records", 0)), + fits={m: ModelFit.from_json(f) for m, f in obj.get("models", {}).items()}, + templates=[RunTemplate.from_json(t) for t in obj.get("run_templates", [])], + ) + + def save(self, path: str) -> None: + with open(path, "w", encoding="utf-8") as f: + json.dump(self.to_json(), f, indent=2) + f.write("\n") + + @classmethod + def load(cls, path: str) -> "Model": + with open(path, "r", encoding="utf-8") as f: + return cls.from_json(json.load(f)) + + def fit_for(self, model: str) -> Optional[ModelFit]: + """The fit for a model, or None if the trace never exercised it.""" + return self.fits.get(model) + + +def _common_stats(xs: np.ndarray, ys: np.ndarray) -> dict: + n = int(xs.size) + return { + "n": n, + "output_mean": float(ys.mean()) if n else 0.0, + "output_std": float(ys.std(ddof=1)) if n > 1 else 0.0, + "input_mean": float(xs.mean()) if n else 0.0, + "input_min": float(xs.min()) if n else 0.0, + "input_max": float(xs.max()) if n else 0.0, + } + + +def _fit_gaussian(model: str, xs: np.ndarray, ys: np.ndarray) -> ModelFit: + """OLS mean line with a symmetric residual spread, or the mean fallback.""" + s = _common_stats(xs, ys) + n = s["n"] + enough = n >= _MIN_FIT_POINTS + has_spread = n > 1 and float(xs.std()) > 0.0 + if not (enough and has_spread): + return ModelFit( + model=model, method="mean", dist="gaussian", + intercept=s["output_mean"], slope=0.0, r2=0.0, + resid_std=s["output_std"], **s) + + slope, intercept = np.polyfit(xs, ys, 1) + pred = intercept + slope * xs + ss_res = float(np.sum((ys - pred) ** 2)) + ss_tot = float(np.sum((ys - ys.mean()) ** 2)) + r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0 + resid_std = float(np.sqrt(ss_res / (n - 2))) if n > 2 else 0.0 + return ModelFit( + model=model, method="ols", dist="gaussian", + intercept=float(intercept), slope=float(slope), r2=r2, + resid_std=resid_std, **s) + + +def _qr_lp(xs: np.ndarray, ys: np.ndarray, tau: float) -> tuple[float, float]: + """Linear quantile regression for one tau, as an exact LP. + + Quantile regression minimises the pinball loss, which is piecewise linear, + so it is a linear program. With residual r = y - (a + b·x) split into + positive/negative parts r = u⁺ - u⁻ (u⁺,u⁻ ≥ 0), the program is + + min Σ τ·u⁺ᵢ + (1-τ)·u⁻ᵢ + s.t. yᵢ - a - b·xᵢ = u⁺ᵢ - u⁻ᵢ + + over free (a, b) and non-negative (u⁺, u⁻). Solved with HiGHS via linprog. + """ + from scipy.optimize import linprog + + n = ys.size + X = np.column_stack([np.ones(n), xs]) # n×2 design (a, b) + c = np.concatenate([np.zeros(2), tau * np.ones(n), (1 - tau) * np.ones(n)]) + A_eq = np.hstack([X, np.eye(n), -np.eye(n)]) + bounds = [(None, None), (None, None)] + [(0, None)] * (2 * n) + res = linprog(c, A_eq=A_eq, b_eq=ys, bounds=bounds, method="highs") + if not res.success: + raise RuntimeError(f"quantile LP failed at tau={tau}: {res.message}") + a, b = res.x[0], res.x[1] + return float(a), float(b) + + +def _fit_quantile(model: str, xs: np.ndarray, ys: np.ndarray, + taus: Sequence[float]) -> ModelFit: + """Conditional-quantile grid by regression, or empirical marginal fallback.""" + s = _common_stats(xs, ys) + n = s["n"] + taus = sorted(taus) + enough = n >= _MIN_QR_POINTS + has_spread = n > 1 and float(xs.std()) > 0.0 + + if not (enough and has_spread): + # No trustworthy input→output trend: flat lines at the marginal output + # quantiles. Still asymmetric, so it beats a symmetric mean fallback. + lines = [QuantileLine(tau=t, intercept=float(np.quantile(ys, t)), slope=0.0) + for t in taus] + median = next(l for l in lines if abs(l.tau - 0.5) < 1e-9) + return ModelFit( + model=model, method="empirical_quantile", dist="quantile", + intercept=median.intercept, slope=0.0, r2=0.0, + resid_std=s["output_std"], quantiles=lines, **s) + + lines: List[QuantileLine] = [] + for t in taus: + try: + a, b = _qr_lp(xs, ys, t) + except RuntimeError: + # A single tau failing the LP degrades to its marginal quantile + # rather than dropping the whole grid. + a, b = float(np.quantile(ys, t)), 0.0 + lines.append(QuantileLine(tau=t, intercept=a, slope=b)) + + median = next(l for l in lines if abs(l.tau - 0.5) < 1e-9) + pred = median.intercept + median.slope * xs + resid = ys - pred + resid_std = float(resid.std(ddof=1)) if n > 1 else 0.0 + # Koenker–Machado pseudo-R¹ for the median fit: 1 − (pinball of the model) / + # (pinball of the unconditional median). The QR analogue of R², reported so + # the user can judge the median fit's strength. + r1 = _pseudo_r1(xs, ys, median, 0.5) + return ModelFit( + model=model, method="quantile_reg", dist="quantile", + intercept=median.intercept, slope=median.slope, r2=r1, + resid_std=resid_std, quantiles=lines, **s) + + +def _pinball(resid: np.ndarray, tau: float) -> float: + return float(np.sum(np.where(resid >= 0, tau * resid, (tau - 1) * resid))) + + +def _pseudo_r1(xs: np.ndarray, ys: np.ndarray, line: QuantileLine, tau: float) -> float: + v_model = _pinball(ys - (line.intercept + line.slope * xs), tau) + v_null = _pinball(ys - np.quantile(ys, tau), tau) + return 1.0 - v_model / v_null if v_null > 0 else 0.0 + + +def _templates(records: List[Record]) -> List[RunTemplate]: + """Group records into per-(scenario, run) call-graph templates. + + Order within a run follows ``seq`` so a replayed run reproduces the observed + call ordering — the fan-out/retry structure the cost depends on. Insertion + order of the runs themselves is preserved so emit-trace is deterministic. + """ + order: List[tuple[str, str]] = [] + groups: Dict[tuple[str, str], List[Record]] = {} + for r in records: + key = (r.scenario_id, r.run_id) + if key not in groups: + groups[key] = [] + order.append(key) + groups[key].append(r) + + out: List[RunTemplate] = [] + for key in order: + recs = sorted(groups[key], key=lambda r: r.seq) + calls = [ + Call(seq=r.seq, model=r.model, input_tokens=r.input_tokens, + cached_tokens=r.cached_tokens, endpoint=r.endpoint) + for r in recs + ] + out.append(RunTemplate(scenario_id=key[0], calls=calls)) + return out + + +def fit(records: List[Record], dist: str = "gaussian", + taus: Sequence[float] = DEFAULT_TAUS) -> Model: + """Fit the output-length model from a recorded trace. + + ``dist`` is ``"gaussian"`` (OLS mean + symmetric spread, the default) or + ``"quantile"`` (a conditional-quantile grid by regression). Only successful + calls feed the per-model fits (see Record.succeeded); every call, success or + not, contributes its structure to the run templates, because a failed-but- + billed call is part of the call graph emit-trace should reproduce. + """ + if dist not in ("gaussian", "quantile"): + raise ValueError(f"unknown dist {dist!r}: want 'gaussian' or 'quantile'") + + by_model: Dict[str, tuple[list, list]] = {} + for r in records: + if not r.succeeded(): + continue + xs, ys = by_model.setdefault(r.model, ([], [])) + xs.append(r.input_tokens) + ys.append(r.output_tokens) + + fits: Dict[str, ModelFit] = {} + for model, (xs, ys) in by_model.items(): + x = np.asarray(xs, dtype=float) + y = np.asarray(ys, dtype=float) + fits[model] = (_fit_quantile(model, x, y, taus) if dist == "quantile" + else _fit_gaussian(model, x, y)) + + return Model(version=1, n_records=len(records), fits=fits, + templates=_templates(records), dist=dist) diff --git a/sidecar/augur_predict/trace.py b/sidecar/augur_predict/trace.py new file mode 100644 index 0000000..6fd202a --- /dev/null +++ b/sidecar/augur_predict/trace.py @@ -0,0 +1,137 @@ +"""Loading and writing the Augur cost trace from Python. + +The trace is JSON Lines: one LLM call per line, the exact schema the Go proxy +emits (see trace/trace.go — ``trace.Record``). Field names are the contract, so +they are mirrored here verbatim; anything we add on the Python side must round- +trip back to a row the Go aggregator will accept. + +We keep a small dataclass rather than leaning on pandas: the trace is the only +data structure the sidecar touches, the files are small (a representative run, +not production telemetry), and a stdlib-only loader keeps the second toolchain +as light as the SPEC wants it. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass, field +from typing import Iterable, Iterator, List + + +@dataclass +class Record: + """One LLM call, mirroring trace.Record on the Go side. + + Only the fields the sidecar reads or writes are typed explicitly; the JSON + keys match the Go ``json:"..."`` tags so a Record marshals straight back + into a row ``augur aggregate`` can price. + """ + + scenario_id: str + run_id: str + seq: int + model: str + input_tokens: int + output_tokens: int + cached_tokens: int = 0 + latency_ms: int = 0 + ts: str = "" + endpoint: str = "" + status: int = 0 + + @classmethod + def from_json(cls, obj: dict) -> "Record": + """Build a Record from a parsed trace line, tolerating absent optionals. + + The proxy omits empty optionals (``omitempty`` on the Go side), so a + line may carry only the required token-accounting fields. Unknown keys + are ignored rather than raising: the trace schema may grow, and the + sidecar should not break on a field it does not model. + """ + return cls( + scenario_id=obj.get("scenario_id", ""), + run_id=obj.get("run_id", ""), + seq=int(obj.get("seq", 0)), + model=obj.get("model", ""), + input_tokens=int(obj.get("input_tokens", 0)), + output_tokens=int(obj.get("output_tokens", 0)), + cached_tokens=int(obj.get("cached_tokens", 0)), + latency_ms=int(obj.get("latency_ms", 0)), + ts=obj.get("ts", ""), + endpoint=obj.get("endpoint", ""), + status=int(obj.get("status", 0)), + ) + + def to_json(self) -> dict: + """Render to a dict with the Go JSON keys, dropping empty optionals. + + We mirror the proxy's ``omitempty`` behaviour for the optional fields so + an emitted trace is byte-comparable in spirit to a recorded one and does + not carry noise the Go reader would just ignore. + """ + out = { + "ts": self.ts, + "scenario_id": self.scenario_id, + "run_id": self.run_id, + "seq": self.seq, + "model": self.model, + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "cached_tokens": self.cached_tokens, + "latency_ms": self.latency_ms, + } + if self.endpoint: + out["endpoint"] = self.endpoint + if self.status: + out["status"] = self.status + return out + + def succeeded(self) -> bool: + """Whether this call is a normal success for output-length modelling. + + A status of 0 means the proxy left it unset (the common case for a + clean call); an explicit 2xx is also a success. Non-2xx rows are kept in + the trace on purpose (they burned input tokens) but they distort an + output-length fit — a 429 produced no completion — so the model excludes + them. They remain a *retry/fan-out* phenomenon, which the Go what-if + knobs already cover. + """ + return self.status == 0 or 200 <= self.status < 300 + + +def parse_lines(lines: Iterable[str]) -> Iterator[Record]: + """Parse an iterable of JSONL strings into Records, skipping blank lines. + + A malformed line is a hard error with its 1-based number, matching the Go + reader's stance: a corrupt trace must not silently shrink the dataset the + model learns from. + """ + for i, line in enumerate(lines, start=1): + line = line.strip() + if not line: + continue + try: + obj = json.loads(line) + except json.JSONDecodeError as e: + raise ValueError(f"trace: parsing line {i}: {e}") from e + yield Record.from_json(obj) + + +def load_trace(path: str) -> List[Record]: + """Read a JSONL trace file into a list of Records.""" + with open(path, "r", encoding="utf-8") as f: + return list(parse_lines(f)) + + +def write_trace(path: str, records: Iterable[Record]) -> int: + """Append records to a JSONL trace file, returning how many were written. + + Append (not truncate) mirrors the proxy: a trace is a ledger you add to. The + caller decides whether to point this at a fresh file or an existing one. + """ + n = 0 + with open(path, "a", encoding="utf-8") as f: + for r in records: + f.write(json.dumps(r.to_json()) + "\n") + n += 1 + return n diff --git a/sidecar/pyproject.toml b/sidecar/pyproject.toml new file mode 100644 index 0000000..386428a --- /dev/null +++ b/sidecar/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "augur-predict" +version = "0.1.0" +description = "Augur's optional predictive output-length sidecar (Hito 5)." +readme = "README.md" +requires-python = ">=3.10" +license = { text = "Apache-2.0" } +authors = [{ name = "Jesús Nuñez" }] +dependencies = ["numpy>=1.24", "scipy>=1.10"] + +[project.optional-dependencies] +dev = ["pytest>=7"] + +[project.scripts] +augur-predict = "augur_predict.cli:main" + +[tool.setuptools] +packages = ["augur_predict"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/sidecar/tests/__init__.py b/sidecar/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sidecar/tests/conftest.py b/sidecar/tests/conftest.py new file mode 100644 index 0000000..de0df9f --- /dev/null +++ b/sidecar/tests/conftest.py @@ -0,0 +1,10 @@ +"""Make the sidecar package importable when running pytest from sidecar/. + +Keeps the suite runnable with a bare ``pytest`` and no editable install, so the +checkpoint is one command on a clean checkout. +""" + +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) diff --git a/sidecar/tests/test_cli.py b/sidecar/tests/test_cli.py new file mode 100644 index 0000000..4c72dea --- /dev/null +++ b/sidecar/tests/test_cli.py @@ -0,0 +1,165 @@ +import json + +import numpy as np +import pytest + +from augur_predict.cli import main +from augur_predict.trace import Record, load_trace, write_trace + + +def _write_trace(path, n=40, noise=8.0, seed=5): + rng = np.random.default_rng(seed) + recs = [] + for i in range(n): + x = 100 + i * 40 + y = 40 + 0.25 * x + rng.normal(0, noise) + recs.append(Record("s1", f"r{i}", 0, "m", int(x), int(round(y)), + cached_tokens=int(x * 0.1))) + write_trace(str(path), recs) + + +def test_fit_then_report_then_predict(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + _write_trace(trace) + + assert main(["fit", "--trace", str(trace), "--out", str(model)]) == 0 + assert model.exists() + obj = json.loads(model.read_text()) + assert "m" in obj["models"] + assert obj["models"]["m"]["method"] == "ols" + + assert main(["report", "--model", str(model)]) == 0 + out = capsys.readouterr().out + assert "Output-length model" in out + assert "m" in out + + assert main(["predict", "--model", str(model), + "--model-name", "m", "--input-tokens", "1000"]) == 0 + out = capsys.readouterr().out + assert "predicted output tokens" in out + + +def test_predict_with_price_prints_cost(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + _write_trace(trace) + main(["fit", "--trace", str(trace), "--out", str(model)]) + capsys.readouterr() + + assert main(["predict", "--model", str(model), "--model-name", "m", + "--input-tokens", "1000", "--price-out", "0.6"]) == 0 + out = capsys.readouterr().out + assert "output cost" in out + + +def test_predict_unknown_model_fails(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + _write_trace(trace) + main(["fit", "--trace", str(trace), "--out", str(model)]) + capsys.readouterr() + + rc = main(["predict", "--model", str(model), "--model-name", "ghost", + "--input-tokens", "100"]) + assert rc == 1 + err = capsys.readouterr().err + assert "no fit" in err + + +def test_fit_empty_trace_fails(tmp_path, capsys): + trace = tmp_path / "empty.jsonl" + trace.write_text("") + rc = main(["fit", "--trace", str(trace), "--out", str(tmp_path / "m.json")]) + assert rc == 1 + assert "empty" in capsys.readouterr().err + + +def test_emit_trace_produces_consumable_jsonl(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + out = tmp_path / "pred.jsonl" + _write_trace(trace) + main(["fit", "--trace", str(trace), "--out", str(model)]) + capsys.readouterr() + + rc = main(["emit-trace", "--model", str(model), "--out", str(out), + "--runs", "15", "--input-scale", "1.5", "--seed", "9"]) + assert rc == 0 + recs = load_trace(str(out)) + assert len({r.run_id for r in recs}) == 15 + for r in recs: + assert r.cached_tokens <= r.input_tokens + assert r.output_tokens >= 0 + + +def test_emit_trace_round_trips_through_trace_loader(tmp_path): + """The emitted file must parse back as a valid trace (Go schema contract).""" + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + out = tmp_path / "pred.jsonl" + _write_trace(trace) + main(["fit", "--trace", str(trace), "--out", str(model)]) + main(["emit-trace", "--model", str(model), "--out", str(out), "--runs", "5"]) + + for line in out.read_text().splitlines(): + if not line.strip(): + continue + obj = json.loads(line) + # required keys the Go aggregator reads + for key in ("scenario_id", "run_id", "seq", "model", + "input_tokens", "output_tokens", "cached_tokens"): + assert key in obj + + +def _write_skewed_trace(path, n=120, seed=4): + rng = np.random.default_rng(seed) + recs = [] + for i in range(n): + x = int(rng.uniform(200, 2000)) + y = int(max(1, round((30 + 0.2 * x) * np.exp(rng.normal(0, 0.35))))) + recs.append(Record("s1", f"r{i}", 0, "m", x, y, cached_tokens=int(x * 0.1))) + write_trace(str(path), recs) + + +def test_quantile_fit_report_predict(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + _write_skewed_trace(trace) + + assert main(["fit", "--trace", str(trace), "--out", str(model), + "--dist", "quantile"]) == 0 + obj = json.loads(model.read_text()) + assert obj["dist"] == "quantile" + assert obj["models"]["m"]["method"] == "quantile_reg" + assert "quantiles" in obj["models"]["m"] + + assert main(["report", "--model", str(model)]) == 0 + out = capsys.readouterr().out + assert "[quantile]" in out + assert "p95 out~in" in out + + assert main(["predict", "--model", str(model), "--model-name", "m", + "--input-tokens", "1500", "--price-out", "10"]) == 0 + out = capsys.readouterr().out + assert "p50" in out and "p95" in out + + +def test_emit_with_run_correlation_flag(tmp_path, capsys): + trace = tmp_path / "trace.jsonl" + model = tmp_path / "model.json" + out = tmp_path / "pred.jsonl" + _write_skewed_trace(trace) + main(["fit", "--trace", str(trace), "--out", str(model), "--dist", "quantile"]) + capsys.readouterr() + + rc = main(["emit-trace", "--model", str(model), "--out", str(out), + "--runs", "20", "--run-correlation", "0.7", "--seed", "3"]) + assert rc == 0 + recs = load_trace(str(out)) + assert len({r.run_id for r in recs}) == 20 + + +def test_no_subcommand_errors(capsys): + with pytest.raises(SystemExit): + main([]) diff --git a/sidecar/tests/test_emit.py b/sidecar/tests/test_emit.py new file mode 100644 index 0000000..dbb49f6 --- /dev/null +++ b/sidecar/tests/test_emit.py @@ -0,0 +1,144 @@ +import numpy as np +import pytest + +from augur_predict.emit import emit +from augur_predict.model import fit +from augur_predict.trace import Record + + +def _model(noise=10.0, n=40): + rng = np.random.default_rng(3) + recs = [] + for i in range(n): + x = 100 + i * 40 + y = 40 + 0.25 * x + rng.normal(0, noise) + recs.append(Record("s1", f"r{i}", 0, "m", int(x), int(round(y)), + cached_tokens=int(x * 0.1))) + return fit(recs) + + +def test_emit_count_matches_runs(): + m = _model() + recs = emit(m, runs=10, seed=0) + run_ids = {r.run_id for r in recs} + assert len(run_ids) == 10 + + +def test_emit_is_deterministic_for_a_seed(): + m = _model() + a = emit(m, runs=8, seed=42) + b = emit(m, runs=8, seed=42) + assert [r.output_tokens for r in a] == [r.output_tokens for r in b] + + +def test_emit_seed_changes_the_sample(): + m = _model() + a = emit(m, runs=8, seed=1) + b = emit(m, runs=8, seed=2) + assert [r.output_tokens for r in a] != [r.output_tokens for r in b] + + +def test_emit_records_are_valid_for_go_aggregate(): + """cost.Usage.Validate on the Go side rejects negatives and cached>input.""" + m = _model() + for r in emit(m, runs=20, seed=7, input_scale=2.0): + assert r.input_tokens >= 0 + assert r.output_tokens >= 0 + assert r.cached_tokens >= 0 + assert r.cached_tokens <= r.input_tokens + assert r.status == 200 + assert r.model == "m" + assert r.scenario_id # non-empty + + +def test_input_scale_inflates_prompt_and_output(): + m = _model(noise=0.0) # deterministic line, no noise + base = emit(m, runs=4, seed=0, input_scale=1.0) + scaled = emit(m, runs=4, seed=0, input_scale=2.0) + # same seed, same templates: compare matched calls + assert scaled[0].input_tokens > base[0].input_tokens + # bigger prompt predicts a longer completion (the slope is positive) + assert scaled[0].output_tokens > base[0].output_tokens + + +def test_emit_zero_runs_is_empty(): + assert emit(_model(), runs=0) == [] + + +def test_emit_rejects_nonpositive_scale(): + with pytest.raises(ValueError, match="positive"): + emit(_model(), runs=2, input_scale=0) + + +def test_scenario_filter_restricts_templates(): + recs = [Record("s1", "r1", 0, "m", 100, 50), + Record("s2", "r1", 0, "m", 200, 60)] + m = fit(recs) + out = emit(m, runs=4, seed=0, scenario_filter="s2") + assert {r.scenario_id for r in out} == {"s2"} + + +def test_unknown_scenario_filter_raises(): + m = _model() + with pytest.raises(ValueError, match="nope"): + emit(m, runs=2, scenario_filter="nope") + + +def test_run_ids_use_prefix(): + out = emit(_model(), runs=3, seed=0, run_prefix="whatif") + assert all(r.run_id.startswith("whatif-") for r in out) + + +def _two_call_model(noise=60.0, n=60): + """A two-call run whose output variance is dominated by residual noise, not + by input variation — so the run-level correlation (which acts on the + residual) is the main driver of the per-run total spread and the effect is + visible rather than swamped by between-run input differences.""" + rng = np.random.default_rng(9) + recs = [] + for i in range(n): + # fixed inputs across runs: the only run-to-run variation is the noise + recs.append(Record("s1", f"r{i}", 0, "m", 600, + int(max(1, 200 + rng.normal(0, noise))))) + recs.append(Record("s1", f"r{i}", 1, "m", 600, + int(max(1, 200 + rng.normal(0, noise))))) + return fit(recs) + + +def _per_run_totals(records): + totals = {} + for r in records: + totals[r.run_id] = totals.get(r.run_id, 0) + r.output_tokens + return np.array(list(totals.values()), dtype=float) + + +def test_run_correlation_widens_per_run_total_variance(): + """The reason it exists: shared verbosity inflates the per-run total spread, + which is what the gate's p95 is taken over.""" + m = _two_call_model() + indep = _per_run_totals(emit(m, runs=400, seed=1, run_correlation=0.0)) + corr = _per_run_totals(emit(m, runs=400, seed=1, run_correlation=0.9)) + assert corr.var() > indep.var() * 1.2 + + +def test_run_correlation_out_of_range_raises(): + with pytest.raises(ValueError, match="run_correlation"): + emit(_model(), runs=2, run_correlation=1.5) + with pytest.raises(ValueError, match="run_correlation"): + emit(_model(), runs=2, run_correlation=-0.1) + + +def test_emit_quantile_model_is_valid_and_deterministic(): + rng = np.random.default_rng(2) + recs = [] + for i in range(120): + x = int(rng.uniform(200, 2000)) + y = int(max(1, round((30 + 0.2 * x) * np.exp(rng.normal(0, 0.35))))) + recs.append(Record("s1", f"r{i}", 0, "m", x, y, cached_tokens=int(x * 0.1))) + m = fit(recs, dist="quantile") + a = emit(m, runs=20, seed=5, input_scale=1.5) + b = emit(m, runs=20, seed=5, input_scale=1.5) + assert [r.output_tokens for r in a] == [r.output_tokens for r in b] + for r in a: + assert r.output_tokens >= 0 + assert r.cached_tokens <= r.input_tokens diff --git a/sidecar/tests/test_model.py b/sidecar/tests/test_model.py new file mode 100644 index 0000000..969d7f1 --- /dev/null +++ b/sidecar/tests/test_model.py @@ -0,0 +1,132 @@ +import numpy as np +import pytest + +from augur_predict.model import Model, ModelFit, fit, _MIN_FIT_POINTS +from augur_predict.trace import Record + + +def _linear_records(model="m", a=40.0, b=0.25, n=40, noise=0.0, seed=1): + """n successful calls whose output is a known linear function of input.""" + rng = np.random.default_rng(seed) + recs = [] + for i in range(n): + x = 100 + i * 50 + y = a + b * x + (rng.normal(0, noise) if noise else 0.0) + recs.append(Record("s1", f"r{i}", 0, model, int(x), int(round(y)))) + return recs + + +def test_ols_recovers_known_slope_and_intercept(): + m = fit(_linear_records(a=40.0, b=0.25, noise=0.0)) + f = m.fit_for("m") + assert f.method == "ols" + # outputs are integer token counts, so the recovered line is exact only up + # to the rounding of y — a few thousandths on the slope, well within noise. + assert f.slope == pytest.approx(0.25, abs=1e-2) + assert f.intercept == pytest.approx(40.0, abs=1.0) + assert f.r2 == pytest.approx(1.0, abs=1e-3) + assert f.resid_std == pytest.approx(0.0, abs=1.0) + + +def test_ols_predict_matches_line(): + f = fit(_linear_records(a=40.0, b=0.25, noise=0.0)).fit_for("m") + assert f.predict(1000) == pytest.approx(40 + 0.25 * 1000, abs=1.0) + + +def test_noisy_fit_has_lower_r2_and_positive_resid_std(): + f = fit(_linear_records(noise=30.0)).fit_for("m") + assert f.method == "ols" + assert 0.0 < f.r2 <= 1.0 + assert f.resid_std > 0.0 + + +def test_fallback_to_mean_when_too_few_points(): + recs = _linear_records(n=_MIN_FIT_POINTS - 1) + f = fit(recs).fit_for("m") + assert f.method == "mean" + assert f.slope == 0.0 + outputs = [r.output_tokens for r in recs] + assert f.intercept == pytest.approx(np.mean(outputs)) + + +def test_fallback_to_mean_when_no_input_spread(): + # Enough points but every input identical -> a slope would be noise. + recs = [Record("s", f"r{i}", 0, "m", 500, 100 + i) for i in range(20)] + f = fit(recs).fit_for("m") + assert f.method == "mean" + assert f.slope == 0.0 + + +def test_predict_clipped_at_zero(): + # Negative intercept, tiny input -> raw line goes negative, must clip. + recs = [Record("s", f"r{i}", 0, "m", 1000 + i * 10, 50 + i) for i in range(20)] + f = fit(recs).fit_for("m") + f.intercept = -100.0 + f.slope = 0.01 + assert f.predict(0) == 0.0 + + +def test_band_is_ordered_and_nonnegative(): + f = fit(_linear_records(noise=30.0)).fit_for("m") + lo, hi = f.band(1000) + assert 0.0 <= lo <= hi + + +def test_failed_calls_excluded_from_fit_but_kept_in_templates(): + recs = _linear_records(n=20) + # add a 429 with wild output that would wreck the regression if included + recs.append(Record("s1", "rX", 1, "m", 300, 9999, status=429)) + m = fit(recs) + f = m.fit_for("m") + # slope stays close to the clean 0.25 because the 429 was excluded + assert f.slope == pytest.approx(0.25, abs=0.05) + # but the failed call still appears in a run template (call graph structure) + all_calls = [c for t in m.templates for c in t.calls] + assert any(c.input_tokens == 300 for c in all_calls) + + +def test_per_model_fits_are_independent(): + recs = _linear_records(model="cheap", a=10, b=0.1) + \ + _linear_records(model="verbose", a=200, b=0.5) + m = fit(recs) + assert m.fit_for("cheap").slope == pytest.approx(0.1, abs=1e-6) + assert m.fit_for("verbose").slope == pytest.approx(0.5, abs=1e-6) + + +def test_templates_group_by_scenario_run_and_order_by_seq(): + recs = [ + Record("s1", "r1", 1, "m", 10, 5), + Record("s1", "r1", 0, "m", 20, 6), + Record("s2", "r1", 0, "m", 30, 7), + ] + m = fit(recs) + assert len(m.templates) == 2 + t = next(t for t in m.templates if t.scenario_id == "s1") + assert [c.seq for c in t.calls] == [0, 1] # sorted by seq + + +def test_model_json_round_trip(tmp_path): + m = fit(_linear_records(noise=10.0)) + m.source = "trace.jsonl" + path = tmp_path / "model.json" + m.save(str(path)) + back = Model.load(str(path)) + assert back.source == "trace.jsonl" + assert back.n_records == m.n_records + f0, f1 = m.fit_for("m"), back.fit_for("m") + assert f1.slope == pytest.approx(f0.slope) + assert f1.intercept == pytest.approx(f0.intercept) + assert f1.r2 == pytest.approx(f0.r2) + assert len(back.templates) == len(m.templates) + + +def test_fit_for_unknown_model_is_none(): + m = fit(_linear_records()) + assert m.fit_for("nope") is None + + +def test_empty_trace_yields_empty_model(): + m = fit([]) + assert m.n_records == 0 + assert m.fits == {} + assert m.templates == [] diff --git a/sidecar/tests/test_quantile.py b/sidecar/tests/test_quantile.py new file mode 100644 index 0000000..cc7620d --- /dev/null +++ b/sidecar/tests/test_quantile.py @@ -0,0 +1,150 @@ +import numpy as np +import pytest + +from augur_predict.model import ( + fit, Model, _phi, _z_for_tau, _MIN_QR_POINTS, DEFAULT_TAUS, +) +from augur_predict.trace import Record + + +def _skewed_records(model="m", n=120, seed=4): + """Right-skewed output: lognormal completion whose scale grows with input. + + Both the skew and the input-dependent spread are things the gaussian band + can't represent but the quantile grid can. + """ + rng = np.random.default_rng(seed) + recs = [] + for i in range(n): + x = 200 + rng.uniform(0, 1800) + # median output ~ 30 + 0.2x, multiplicative lognormal noise (skewed), + # spread widening with x. + med = 30 + 0.2 * x + y = med * np.exp(rng.normal(0, 0.3 + 0.0002 * x)) + recs.append(Record("s1", f"r{i}", 0, model, int(x), int(max(1, round(y))))) + return recs + + +def test_phi_and_inverse_are_consistent(): + for tau in (0.05, 0.25, 0.5, 0.75, 0.95): + assert _phi(_z_for_tau(tau)) == pytest.approx(tau, abs=1e-6) + + +def test_quantile_fit_uses_regression_with_enough_data(): + m = fit(_skewed_records(), dist="quantile") + f = m.fit_for("m") + assert m.dist == "quantile" + assert f.dist == "quantile" + assert f.method == "quantile_reg" + assert [q.tau for q in f.quantiles] == sorted(DEFAULT_TAUS) + + +def test_quantiles_are_ordered_p95_above_p50_above_p05(): + f = fit(_skewed_records(), dist="quantile").fit_for("m") + for x in (300, 1000, 1900): + p05 = f.quantile_at(x, 0.05) + p50 = f.quantile_at(x, 0.5) + p95 = f.quantile_at(x, 0.95) + assert p05 <= p50 <= p95 + + +def test_quantile_band_is_asymmetric_for_skewed_data(): + """The whole point: the upper tail is farther from the median than the lower + tail, which a symmetric gaussian band cannot express.""" + f = fit(_skewed_records(), dist="quantile").fit_for("m") + x = 1500 + p05 = f.quantile_at(x, 0.05) + p50 = f.quantile_at(x, 0.5) + p95 = f.quantile_at(x, 0.95) + upper = p95 - p50 + lower = p50 - p05 + assert upper > lower * 1.2 # clearly right-skewed + + +def test_quantile_p95_exceeds_gaussian_p95_on_skewed_data(): + """A heavy upper tail is exactly what gaussian underestimates.""" + recs = _skewed_records() + x = 1500 + qf = fit(recs, dist="quantile").fit_for("m") + gf = fit(recs, dist="gaussian").fit_for("m") + q95 = qf.quantile_at(x, 0.95) + g95 = gf.predict(x) + 1.6448536 * gf.resid_std # gaussian ~p95 + assert q95 > g95 + + +def test_quantile_at_clamps_outside_fitted_range(): + f = fit(_skewed_records(), dist="quantile").fit_for("m") + x = 1000 + lo_end = f.quantile_at(x, f.quantiles[0].tau) + hi_end = f.quantile_at(x, f.quantiles[-1].tau) + assert f.quantile_at(x, 0.001) == pytest.approx(lo_end) + assert f.quantile_at(x, 0.999) == pytest.approx(hi_end) + + +def test_sample_is_monotone_in_z(): + f = fit(_skewed_records(), dist="quantile").fit_for("m") + x = 1000 + lows = f.sample(x, -1.5) + mids = f.sample(x, 0.0) + highs = f.sample(x, 1.5) + assert lows <= mids <= highs + + +def test_empirical_quantile_fallback_when_few_points(): + recs = _skewed_records(n=_MIN_QR_POINTS - 2) + f = fit(recs, dist="quantile").fit_for("m") + assert f.method == "empirical_quantile" + # flat in input (no trustworthy slope) ... + assert all(q.slope == 0.0 for q in f.quantiles) + # ... but still asymmetric (skewed marginal) + p50 = f.quantile_at(0, 0.5) + p95 = f.quantile_at(0, 0.95) + p05 = f.quantile_at(0, 0.05) + assert (p95 - p50) > (p50 - p05) + + +def test_empirical_fallback_when_no_input_spread(): + recs = [Record("s", f"r{i}", 0, "m", 500, int(50 + i)) for i in range(30)] + f = fit(recs, dist="quantile").fit_for("m") + assert f.method == "empirical_quantile" + + +def test_pseudo_r1_in_range(): + f = fit(_skewed_records(), dist="quantile").fit_for("m") + assert 0.0 <= f.r2 <= 1.0 + + +def test_quantile_model_json_round_trip(tmp_path): + m = fit(_skewed_records(), dist="quantile") + m.source = "trace.jsonl" + path = tmp_path / "model.json" + m.save(str(path)) + back = Model.load(str(path)) + assert back.dist == "quantile" + f0, f1 = m.fit_for("m"), back.fit_for("m") + assert f1.method == f0.method + assert len(f1.quantiles) == len(f0.quantiles) + for a, b in zip(f0.quantiles, f1.quantiles): + assert b.tau == pytest.approx(a.tau) + assert b.slope == pytest.approx(a.slope) + assert b.intercept == pytest.approx(a.intercept) + + +def test_gaussian_model_loads_without_dist_field(tmp_path): + """Backward compat: an old artifact with no 'dist' key reads as gaussian.""" + import json + m = fit(_skewed_records(), dist="gaussian") + obj = m.to_json() + del obj["dist"] + for f in obj["models"].values(): + f.pop("dist", None) + path = tmp_path / "old.json" + path.write_text(json.dumps(obj)) + back = Model.load(str(path)) + assert back.dist == "gaussian" + assert back.fit_for("m").dist == "gaussian" + + +def test_unknown_dist_raises(): + with pytest.raises(ValueError, match="unknown dist"): + fit(_skewed_records(), dist="poisson") diff --git a/sidecar/tests/test_trace.py b/sidecar/tests/test_trace.py new file mode 100644 index 0000000..240703d --- /dev/null +++ b/sidecar/tests/test_trace.py @@ -0,0 +1,92 @@ +import json + +import pytest + +from augur_predict.trace import Record, parse_lines, load_trace, write_trace + + +def test_from_json_required_and_optional_fields(): + line = { + "ts": "2026-06-22T10:00:00Z", + "scenario_id": "s1", "run_id": "r1", "seq": 2, + "model": "gpt-4o-mini", + "input_tokens": 1200, "output_tokens": 300, "cached_tokens": 400, + "latency_ms": 850, "endpoint": "/v1/chat/completions", "status": 200, + } + r = Record.from_json(line) + assert r.scenario_id == "s1" + assert r.seq == 2 + assert r.input_tokens == 1200 + assert r.cached_tokens == 400 + assert r.status == 200 + + +def test_from_json_tolerates_missing_optionals_and_unknown_keys(): + r = Record.from_json({ + "scenario_id": "s", "run_id": "r", "seq": 0, "model": "m", + "input_tokens": 10, "output_tokens": 5, + "some_future_field": "ignored", + }) + assert r.cached_tokens == 0 + assert r.status == 0 + assert r.endpoint == "" + + +def test_to_json_drops_empty_optionals_like_go_omitempty(): + r = Record("s", "r", 0, "m", 10, 5) + obj = r.to_json() + assert "endpoint" not in obj + assert "status" not in obj + # required token accounting always present + assert obj["input_tokens"] == 10 + assert obj["output_tokens"] == 5 + assert obj["cached_tokens"] == 0 + + +def test_to_json_keeps_set_optionals(): + r = Record("s", "r", 0, "m", 10, 5, status=429, endpoint="/v1/x") + obj = r.to_json() + assert obj["status"] == 429 + assert obj["endpoint"] == "/v1/x" + + +def test_succeeded_classification(): + assert Record("s", "r", 0, "m", 1, 1, status=0).succeeded() + assert Record("s", "r", 0, "m", 1, 1, status=200).succeeded() + assert Record("s", "r", 0, "m", 1, 1, status=299).succeeded() + assert not Record("s", "r", 0, "m", 1, 1, status=429).succeeded() + assert not Record("s", "r", 0, "m", 1, 1, status=500).succeeded() + + +def test_parse_lines_skips_blanks(): + lines = ['{"scenario_id":"s","run_id":"r","seq":0,"model":"m","input_tokens":1,"output_tokens":2}', + "", " "] + recs = list(parse_lines(lines)) + assert len(recs) == 1 + + +def test_parse_lines_malformed_is_hard_error_with_line_number(): + with pytest.raises(ValueError, match="line 2"): + list(parse_lines(["{}", "not json"])) + + +def test_load_and_write_round_trip(tmp_path): + src = [ + Record("s1", "r1", 0, "m", 100, 50, cached_tokens=10), + Record("s1", "r1", 1, "m", 200, 80), + ] + path = tmp_path / "t.jsonl" + n = write_trace(str(path), src) + assert n == 2 + back = load_trace(str(path)) + assert len(back) == 2 + assert back[0].input_tokens == 100 + assert back[0].cached_tokens == 10 + assert back[1].seq == 1 + + +def test_write_appends_not_truncates(tmp_path): + path = tmp_path / "t.jsonl" + write_trace(str(path), [Record("s", "r", 0, "m", 1, 1)]) + write_trace(str(path), [Record("s", "r", 1, "m", 1, 1)]) + assert len(load_trace(str(path))) == 2