nabla-cli 0.1.0__py3-none-any.whl

nabla/recorder.py

@@ -0,0 +1,609 @@
+"""T4 — traffic recorder. See cli-plan.md T4.
+
+Three capture paths behind one interface (`record`):
+
+  1. "samples" — run the site's enclosing function in a fresh subprocess per
+     sample-input line, with OPENAI_BASE_URL pointed at a local recording
+     proxy. Subprocess because target modules bind an OpenAI client at
+     import time (see fixtures/demo-repo/app/*.py), so env vars must be set
+     *before* the module is imported.
+  2. "proxy" — run an arbitrary command (typically the target's pytest) as
+     a subprocess under the same proxy and harvest whatever traffic flows.
+     The fixture repo's own tests mock the SDK client directly, so this path
+     must detect "zero traffic" and say so explicitly (ZeroTrafficError)
+     instead of silently returning an empty profile.
+  3. "logs" — parse user-supplied JSONL of past request/response bodies.
+
+The recording proxy is a tiny threaded HTTP server. Its upstream behavior is
+injectable (a `responder` callable, or NABLA_FAKE_UPSTREAM=1 for a canned
+response) so tests never touch the network or need a real API key. Default
+behavior with no override forwards to https://api.openai.com with auth
+passthrough.
+
+Site attribution (matching a captured request back to `site`) uses a
+lightweight static heuristic: string literals reachable from the site's
+enclosing function (including module-level constants it references, and up
+to two hops into same-module helper functions it calls) are compared against
+the captured request's message content. This is deliberately conservative —
+a request that doesn't share enough of the site's static text is not
+attributed to it.
+"""
+
+from __future__ import annotations
+
+import ast
+import http.server
+import json
+import os
+import statistics
+import subprocess
+import sys
+import threading
+import time
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+from .contract import Example, RawSite
+
+
+@dataclass
+class RecordSource:
+    """Where examples come from: exactly one of the three paths."""
+    kind: str                     # "samples" | "proxy" | "logs"
+    samples_file: str | None = None   # JSONL of {"input_slots": {...}} or {"prompt": "..."}
+    logs_file: str | None = None      # JSONL of past request/response pairs
+    proxy_command: list[str] | None = None  # command to run under the proxy (e.g. pytest)
+
+
+class ZeroTrafficError(RuntimeError):
+    """Raised when the proxy path observed no OpenAI traffic — e.g. the
+    target's tests mock the SDK. Message must say this explicitly and point
+    at the samples/logs paths."""
+
+
+# ---------------------------------------------------------------------------
+# Recording proxy
+# ---------------------------------------------------------------------------
+
+@dataclass
+class _Capture:
+    request: dict
+    response: dict
+    latency_ms: int
+
+
+Responder = Callable[[dict, dict], dict]
+
+
+def _fake_upstream_responder(body: dict, headers: dict) -> dict:
+    """Canned upstream used under NABLA_FAKE_UPSTREAM=1 (or when a caller
+    passes no responder in a test context). Never touches the network, never
+    needs a real API key, always returns a valid chat-completions body with
+    usage numbers."""
+    messages = body.get("messages", [])
+    user_text = ""
+    for m in reversed(messages):
+        if m.get("role") == "user" and isinstance(m.get("content"), str):
+            user_text = m["content"]
+            break
+    completion_text = json.dumps({"fake": True, "echo_chars": len(user_text)})
+    prompt_tokens = max(1, len(user_text) // 4)
+    completion_tokens = max(1, len(completion_text) // 4)
+    return {
+        "id": "chatcmpl-fake0000000000000000000",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": body.get("model", "gpt-4o-mini"),
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": completion_text},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens,
+        },
+    }
+
+
+def _forward_to_openai_responder(body: dict, headers: dict) -> dict:
+    """Default upstream: forward to the real OpenAI API with auth
+    passthrough.
+
+    Overridable via env for recording without an OpenAI key:
+      NABLA_UPSTREAM_BASE_URL — any OpenAI-compatible endpoint (e.g. Gemini's
+        compat layer) used as a stand-in oracle;
+      NABLA_UPSTREAM_API_KEY  — replaces the passthrough Authorization;
+      NABLA_UPSTREAM_MODEL    — rewrites the request's model for the upstream
+        only. The target repo's own model kwarg is what lands in the
+        profile's sampling/baseline blocks, so the "before" cost story is
+        unaffected by the stand-in.
+    """
+    base_url = (os.environ.get("NABLA_UPSTREAM_BASE_URL") or "https://api.openai.com/v1").rstrip("/")
+    api_key = os.environ.get("NABLA_UPSTREAM_API_KEY")
+    auth = f"Bearer {api_key}" if api_key else headers.get("Authorization", "")
+    model_override = os.environ.get("NABLA_UPSTREAM_MODEL")
+    if model_override:
+        body = {**body, "model": model_override}
+    req = urllib.request.Request(
+        f"{base_url}/chat/completions",
+        data=json.dumps(body).encode("utf-8"),
+        method="POST",
+        headers={
+            "Content-Type": "application/json",
+            "Authorization": auth,
+            # Some providers' CDNs (e.g. Cloudflare in front of Groq) reject
+            # urllib's default Python-urllib User-Agent with a 403.
+            "User-Agent": "nabla-recorder/0.1",
+        },
+    )
+    # Rate limits (429) and transient upstream errors are expected during a
+    # long record run, especially on free-tier stand-in oracles: back off and
+    # retry, honoring Retry-After when the upstream sends one.
+    delays = [5, 10, 20, 30, 45]
+    for attempt, delay in enumerate([*delays, None]):
+        try:
+            with urllib.request.urlopen(req, timeout=60) as resp:
+                return json.loads(resp.read().decode("utf-8"))
+        except urllib.error.HTTPError as exc:
+            if delay is None or exc.code not in (429, 500, 502, 503):
+                # Surface the upstream's error body — "HTTP 400" alone is
+                # undiagnosable (e.g. Groq's json_validate_failed).
+                try:
+                    detail = exc.read().decode("utf-8", "replace")[:500]
+                except Exception:  # noqa: BLE001 - body read is best-effort
+                    detail = ""
+                raise RuntimeError(f"upstream HTTP {exc.code}: {detail or exc.reason}") from exc
+            retry_after = exc.headers.get("Retry-After") if exc.headers else None
+            try:
+                wait = max(float(retry_after), delay) if retry_after else delay
+            except ValueError:
+                wait = delay
+            time.sleep(min(wait, 60))
+    raise RuntimeError("unreachable")
+
+
+def _default_responder() -> Responder:
+    if os.environ.get("NABLA_FAKE_UPSTREAM") == "1":
+        return _fake_upstream_responder
+    return _forward_to_openai_responder
+
+
+class RecordingProxy:
+    """Local threaded HTTP server that records chat-completions traffic.
+
+    Binds 127.0.0.1 on an OS-assigned free port (port 0). Upstream behavior
+    is injectable via `responder`; if omitted it's resolved from
+    NABLA_FAKE_UPSTREAM at start() time.
+    """
+
+    def __init__(self, responder: Responder | None = None):
+        self.captured: list[_Capture] = []
+        self._lock = threading.Lock()
+        self._responder = responder or _default_responder()
+        self._server: http.server.ThreadingHTTPServer | None = None
+        self._thread: threading.Thread | None = None
+
+    @property
+    def port(self) -> int:
+        assert self._server is not None, "proxy not started"
+        return self._server.server_address[1]
+
+    @property
+    def base_url(self) -> str:
+        return f"http://127.0.0.1:{self.port}/v1"
+
+    def start(self) -> None:
+        proxy = self
+
+        class Handler(http.server.BaseHTTPRequestHandler):
+            def log_message(self, fmt, *args):  # noqa: A002 - stdlib signature
+                pass  # silence per-request logging to stderr
+
+            def do_POST(self):
+                length = int(self.headers.get("Content-Length") or 0)
+                raw = self.rfile.read(length) if length else b"{}"
+                try:
+                    body = json.loads(raw or b"{}")
+                except json.JSONDecodeError:
+                    body = {}
+
+                if not self.path.rstrip("/").endswith("chat/completions"):
+                    self.send_response(404)
+                    self.end_headers()
+                    return
+
+                t0 = time.perf_counter()
+                try:
+                    resp_body = proxy._responder(body, dict(self.headers))
+                except Exception as exc:  # noqa: BLE001 - report upstream failure to caller
+                    payload = json.dumps({"error": {"message": str(exc)}}).encode("utf-8")
+                    self.send_response(502)
+                    self.send_header("Content-Type", "application/json")
+                    self.send_header("Content-Length", str(len(payload)))
+                    self.end_headers()
+                    self.wfile.write(payload)
+                    return
+                latency_ms = max(1, round((time.perf_counter() - t0) * 1000))
+
+                with proxy._lock:
+                    proxy.captured.append(_Capture(request=body, response=resp_body, latency_ms=latency_ms))
+
+                payload = json.dumps(resp_body).encode("utf-8")
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(payload)))
+                self.end_headers()
+                self.wfile.write(payload)
+
+        self._server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        if self._server is not None:
+            self._server.shutdown()
+            self._server.server_close()
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+
+
+# ---------------------------------------------------------------------------
+# Attribution: static-literal matching between a site and a captured request
+# ---------------------------------------------------------------------------
+
+def _callee_name(func_node: ast.AST) -> str | None:
+    if isinstance(func_node, ast.Name):
+        return func_node.id
+    if isinstance(func_node, ast.Attribute):
+        return func_node.attr
+    return None
+
+
+def _extract_site_literals(file_path: str, symbol: str, max_hops: int = 2) -> set[str]:
+    """Static string literals reachable from `symbol`'s function body in
+    `file_path`: literals written directly in the function, module-level
+    string constants it references by name, and (up to `max_hops`) literals
+    from same-module helper functions it calls. Best-effort — used only to
+    build an attribution heuristic, never to fail hard."""
+    try:
+        source = Path(file_path).read_text(encoding="utf-8")
+        tree = ast.parse(source)
+    except (OSError, SyntaxError):
+        return set()
+
+    module_consts: dict[str, str] = {}
+    module_funcs: dict[str, ast.AST] = {}
+    for node in tree.body:
+        if (
+            isinstance(node, ast.Assign)
+            and len(node.targets) == 1
+            and isinstance(node.targets[0], ast.Name)
+            and isinstance(node.value, ast.Constant)
+            and isinstance(node.value.value, str)
+        ):
+            module_consts[node.targets[0].id] = node.value.value
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            module_funcs[node.name] = node
+
+    func_name = symbol.rstrip("()").split(".")[-1]
+    target = module_funcs.get(func_name)
+    if target is None:
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
+                target = node
+                break
+    if target is None:
+        return set()
+
+    literals: set[str] = set()
+    visited: set[str] = set()
+
+    def visit(fn_node: ast.AST, hops_left: int) -> None:
+        for node in ast.walk(fn_node):
+            if isinstance(node, ast.Constant) and isinstance(node.value, str):
+                if len(node.value.strip()) >= 6:
+                    literals.add(node.value)
+            elif isinstance(node, ast.Name) and node.id in module_consts:
+                literals.add(module_consts[node.id])
+            elif isinstance(node, ast.Call) and hops_left > 0:
+                callee = _callee_name(node.func)
+                if callee and callee in module_funcs and callee not in visited:
+                    visited.add(callee)
+                    visit(module_funcs[callee], hops_left - 1)
+
+    visit(target, max_hops)
+    return literals
+
+
+def _matches_site(request_body: dict, literals: set[str]) -> bool:
+    if not literals:
+        return False
+    messages = request_body.get("messages", [])
+    haystack = " ".join(
+        m.get("content", "") for m in messages if isinstance(m.get("content"), str)
+    )
+    haystack_n = " ".join(haystack.split())
+    considered = [" ".join(lit.split()) for lit in literals if len(lit.strip()) >= 6]
+    considered = [lit for lit in considered if lit]
+    if not considered:
+        return False
+    hits = sum(1 for lit in considered if lit in haystack_n)
+    return (hits / len(considered)) >= 0.34
+
+
+# ---------------------------------------------------------------------------
+# Example construction
+# ---------------------------------------------------------------------------
+
+def _capture_to_example(capture: _Capture, input_slots: dict[str, str]) -> Example:
+    request, response = capture.request, capture.response
+    messages = request.get("messages", []) or []
+    user_msg = ""
+    for m in reversed(messages):
+        if m.get("role") == "user" and isinstance(m.get("content"), str):
+            user_msg = m["content"]
+            break
+
+    completion = ""
+    choices = response.get("choices") or []
+    if choices:
+        msg = choices[0].get("message") or {}
+        completion = msg.get("content") or ""
+
+    usage = response.get("usage") or {}
+    if usage:
+        tokens_in = int(usage.get("prompt_tokens", 0) or 0)
+        tokens_out = int(usage.get("completion_tokens", 0) or 0)
+    else:
+        tokens_in = len(user_msg) // 4
+        tokens_out = len(completion) // 4
+
+    return Example(
+        input_slots=dict(input_slots),
+        prompt=user_msg,
+        completion=completion,
+        tokens_in=tokens_in,
+        tokens_out=tokens_out,
+        latency_ms=capture.latency_ms,
+    )
+
+
+def _read_jsonl(path: str) -> list[dict]:
+    lines: list[dict] = []
+    with open(path, "r", encoding="utf-8") as f:
+        for raw_line in f:
+            raw_line = raw_line.strip()
+            if not raw_line:
+                continue
+            lines.append(json.loads(raw_line))
+    return lines
+
+
+# ---------------------------------------------------------------------------
+# Samples path
+# ---------------------------------------------------------------------------
+
+_RUNNER_SRC = """
+import sys, json, importlib
+
+repo_path, rel_file, symbol, kwargs_json = sys.argv[1:5]
+sys.path.insert(0, repo_path)
+module_path = rel_file[:-3] if rel_file.endswith(".py") else rel_file
+module_name = module_path.replace("/", ".")
+mod = importlib.import_module(module_name)
+func_name = symbol.rstrip("()").split(".")[-1]
+func = getattr(mod, func_name)
+kwargs = json.loads(kwargs_json)
+func(**kwargs)
+sys.stdout.write("OK")
+"""
+
+
+def _run_site_in_subprocess(site: RawSite, repo_path: str, input_slots: dict, env: dict) -> None:
+    rel_file = site.file.replace("\\", "/")
+    args = [sys.executable, "-c", _RUNNER_SRC, repo_path, rel_file, site.symbol, json.dumps(input_slots)]
+    # Generous default: the recording proxy may spend minutes in upstream
+    # backoff-retry (rate-limited stand-in oracles), and the subprocess must
+    # outlive that. Override via NABLA_SAMPLE_TIMEOUT (seconds).
+    timeout = int(os.environ.get("NABLA_SAMPLE_TIMEOUT", "300"))
+    proc = subprocess.run(args, cwd=repo_path, env=env, capture_output=True, text=True, timeout=timeout)
+    if proc.returncode != 0:
+        raise RuntimeError(
+            f"Sample invocation of {site.symbol!r} in {site.file!r} failed "
+            f"(exit {proc.returncode}): {proc.stderr.strip() or proc.stdout.strip()}"
+        )
+
+
+def _record_samples(site: RawSite, source: RecordSource, repo_path: str) -> list[Example]:
+    if not source.samples_file:
+        raise ValueError("RecordSource.kind='samples' requires samples_file")
+    sample_lines = _read_jsonl(source.samples_file)
+    examples: list[Example] = []
+    if not sample_lines:
+        return examples
+
+    proxy = RecordingProxy()
+    proxy.start()
+    try:
+        env = os.environ.copy()
+        env["OPENAI_BASE_URL"] = proxy.base_url
+        env["OPENAI_API_KEY"] = env.get("OPENAI_API_KEY") or "sk-nabla-recorder-test"
+
+        # Optional pacing between samples (milliseconds) so rate-limited
+        # upstreams aren't hammered into 429 backoff on every call.
+        delay_ms = int(os.environ.get("NABLA_RECORD_DELAY_MS", "0"))
+        failures: list[tuple[int, str]] = []
+        total = len(sample_lines)
+
+        for i, line in enumerate(sample_lines):
+            input_slots = line.get("input_slots")
+            if input_slots is None:
+                raise ValueError(
+                    f"samples_file line missing 'input_slots': {line!r} "
+                    "(the samples path invokes the site's function with input_slots as kwargs)"
+                )
+            if delay_ms and i:
+                time.sleep(delay_ms / 1000)
+            with proxy._lock:
+                before = len(proxy.captured)
+            # One bad sample must not kill the harvest: record the failure,
+            # keep going, and fail hard only if every sample failed.
+            try:
+                _run_site_in_subprocess(site, repo_path, input_slots, env)
+            except (RuntimeError, subprocess.TimeoutExpired) as exc:
+                failures.append((i, str(exc)))
+                print(f"nabla record: sample {i + 1}/{total} FAILED: {str(exc)[:300]}",
+                      file=sys.stderr)
+                continue
+            with proxy._lock:
+                new_captures = list(proxy.captured[before:])
+            for capture in new_captures:
+                examples.append(_capture_to_example(capture, input_slots))
+            print(f"nabla record: sample {i + 1}/{total} -> {len(new_captures)} capture(s)",
+                  file=sys.stderr)
+    finally:
+        proxy.stop()
+
+    if failures and not examples:
+        raise RuntimeError(
+            f"all {len(failures)} sample invocations failed; first error: {failures[0][1]}"
+        )
+    if failures:
+        print(f"nabla record: {len(failures)}/{total} samples failed and were skipped",
+              file=sys.stderr)
+    return examples
+
+
+# ---------------------------------------------------------------------------
+# Proxy path
+# ---------------------------------------------------------------------------
+
+def _record_proxy(site: RawSite, source: RecordSource, repo_path: str) -> list[Example]:
+    if not source.proxy_command:
+        raise ValueError("RecordSource.kind='proxy' requires proxy_command")
+
+    proxy = RecordingProxy()
+    proxy.start()
+    try:
+        env = os.environ.copy()
+        env["OPENAI_BASE_URL"] = proxy.base_url
+        env["OPENAI_API_KEY"] = env.get("OPENAI_API_KEY") or "sk-nabla-recorder-test"
+        subprocess.run(source.proxy_command, cwd=repo_path, env=env, capture_output=True, text=True, timeout=300)
+    finally:
+        proxy.stop()
+
+    command_str = " ".join(source.proxy_command)
+
+    if not proxy.captured:
+        raise ZeroTrafficError(
+            f"Recording proxy observed zero OpenAI requests while running `{command_str}`. "
+            "The target's tests likely mock the OpenAI SDK client directly (e.g. via "
+            "unittest.mock.patch on client.chat.completions.create), so no HTTP traffic ever "
+            "reaches OPENAI_BASE_URL. Use the 'samples' or 'logs' RecordSource path instead."
+        )
+
+    literals = _extract_site_literals(str(Path(repo_path) / site.file), site.symbol)
+    matched = [c for c in proxy.captured if _matches_site(c.request, literals)]
+
+    if not matched:
+        raise ZeroTrafficError(
+            f"Recording proxy observed {len(proxy.captured)} OpenAI request(s) while running "
+            f"`{command_str}`, but none were attributable to site {site.symbol!r} in {site.file!r}. "
+            "The target's tests likely mock the OpenAI SDK for this site, so no attributable "
+            "traffic reaches OPENAI_BASE_URL. Use the 'samples' or 'logs' RecordSource path instead."
+        )
+
+    return [_capture_to_example(c, input_slots={}) for c in matched]
+
+
+# ---------------------------------------------------------------------------
+# Logs path
+# ---------------------------------------------------------------------------
+
+def _record_logs(site: RawSite, source: RecordSource, repo_path: str) -> list[Example]:
+    if not source.logs_file:
+        raise ValueError("RecordSource.kind='logs' requires logs_file")
+    examples: list[Example] = []
+    for line in _read_jsonl(source.logs_file):
+        request = line.get("request", {}) or {}
+        response = line.get("response", {}) or {}
+        latency_ms = int(line.get("latency_ms", 0) or 0)
+        capture = _Capture(request=request, response=response, latency_ms=latency_ms)
+        examples.append(_capture_to_example(capture, input_slots={}))
+    return examples
+
+
+# ---------------------------------------------------------------------------
+# Public interface
+# ---------------------------------------------------------------------------
+
+def record(site: RawSite, source: RecordSource, repo_path: str) -> list[Example]:
+    """Capture real (prompt, completion, tokens, latency) pairs for site.
+
+    Must: attribute captured requests back to the site (template match),
+    populate token counts and latency, raise ZeroTrafficError on an empty
+    proxy run instead of returning [].
+    """
+    if source.kind == "samples":
+        return _record_samples(site, source, repo_path)
+    if source.kind == "proxy":
+        return _record_proxy(site, source, repo_path)
+    if source.kind == "logs":
+        return _record_logs(site, source, repo_path)
+    raise ValueError(f"unknown RecordSource.kind: {source.kind!r}")
+
+
+# ---------------------------------------------------------------------------
+# Baseline stats
+# ---------------------------------------------------------------------------
+
+# Pricing as of 2026-07-18, USD per 1,000,000 tokens (input, output).
+# Approximate — confirm against platform.openai.com/pricing before relying
+# on this for real billing decisions; models not listed fall back to
+# cost 0.0 with pricing_unknown=True rather than a guess.
+_PRICING_PER_1M_USD: dict[str, tuple[float, float]] = {
+    "gpt-4o": (2.50, 10.00),
+    "gpt-4o-2024-08-06": (2.50, 10.00),
+    "gpt-4o-mini": (0.15, 0.60),
+    "gpt-4o-mini-2024-07-18": (0.15, 0.60),
+    "gpt-4.1": (2.00, 8.00),
+    "gpt-4.1-mini": (0.40, 1.60),
+    "gpt-4.1-nano": (0.10, 0.40),
+    "gpt-3.5-turbo": (0.50, 1.50),
+    "o1": (15.00, 60.00),
+    "o1-mini": (1.10, 4.40),
+    "o3-mini": (1.10, 4.40),
+}
+
+
+def baseline_stats(examples: list[Example], model: str) -> dict:
+    """Cost/latency baseline from observed traffic and the pinned pricing
+    table (hardcode current OpenAI prices with a date stamp)."""
+    calls_observed = len(examples)
+    if calls_observed == 0:
+        return {"cost_per_call_usd": 0.0, "p50_latency_ms": 0, "calls_observed": 0}
+
+    pricing = _PRICING_PER_1M_USD.get(model)
+    result: dict = {}
+    if pricing is None:
+        avg_cost = 0.0
+        result["pricing_unknown"] = True
+    else:
+        in_price, out_price = pricing
+        costs = [
+            (e.tokens_in / 1_000_000) * in_price + (e.tokens_out / 1_000_000) * out_price
+            for e in examples
+        ]
+        avg_cost = sum(costs) / len(costs)
+
+    result["cost_per_call_usd"] = round(avg_cost, 6)
+    result["p50_latency_ms"] = int(round(statistics.median(e.latency_ms for e in examples)))
+    result["calls_observed"] = calls_observed
+    return result