From eea9600e641670a81fce39ff78301db3f4b8d0d4 Mon Sep 17 00:00:00 2001 From: Ajay03299 Date: Wed, 5 Aug 2026 07:41:01 +0530 Subject: [PATCH] =?UTF-8?q?feat(e5):=20audit=20runner=20=E2=80=94=20provid?= =?UTF-8?q?ers,=20tolerant=20parsing,=20recorded=20runs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the E5 harness: corpus -> provider -> parsed response -> JSONL -> metrics panel. Runs fully offline; live calls are the exception. e5/providers.py — one structural Protocol, three implementations. StubProvider emits JSON so stub runs exercise the same parser a live reply would. ReplayProvider turns one real audit into a permanent offline fixture, and raises on a missing key rather than defaulting, so a partial recording cannot masquerade as a complete one. LiveProvider is the only one that touches the network; it reads its key from the environment, never accepts a literal, and refuses the [YOUR_KEY_HERE] placeholder. e5/parsing.py — parses tolerantly and records what it found rather than accepting only well-formed output. If model A emits clean JSON 99% of the time and model B 80%, discarding failures compares A's full distribution against B's tidiest 80%, understating B's dispersion — the quantity E5 measures. Format compliance plausibly correlates with templated advice, so silent rejection would bias h upward. ParseStatus records ok / renormalized / partial / ambiguous / no_weights / empty; the decision to exclude anything is made at analysis time. Two prose-binding bugs found and fixed during development, both of which returned plausible wrong answers rather than raising: - per-ticker proximity search: in "40% in AAPL, 35% in MSFT" the 40% falls within the window of MSFT and is claimed twice - greedy nearest-neighbour with a claim set: 35% is nearer AAPL, finds it taken, and is dropped — MSFT ends with no weight Both are regression-tested. The corpus prompts now carry an explicit JSON output contract, so prose parsing is a fallback whose failures are flagged AMBIGUOUS rather than silently mis-bound; separating "AAPL returned 30%" from "AAPL 50%" needs semantics, not regex. e5/runner.py — append-only JSONL keyed on (query_id, model, repeat), so an interrupted run resumes and a partial run is still analyzable. Provider failures are recorded and the run continues. Allocations are padded to the full universe: items ask about ticker subsets, but metrics need a common basis to stack into a (models, repeats, assets) panel. .env.example gains OPENAI_KEY and ANTHROPIC_KEY, which retail_ai.py's docstring already referenced but the template never defined. 47 new tests (28 parsing, 19 runner); 180 total. --- .env.example | 2 + e5/corpus.py | 7 +- e5/parsing.py | 196 +++++++++++++++++++++++++++++++++++++++ e5/providers.py | 176 +++++++++++++++++++++++++++++++++++ e5/runner.py | 188 +++++++++++++++++++++++++++++++++++++ tests/test_e5_parsing.py | 147 +++++++++++++++++++++++++++++ tests/test_e5_runner.py | 164 ++++++++++++++++++++++++++++++++ 7 files changed, 879 insertions(+), 1 deletion(-) create mode 100644 e5/parsing.py create mode 100644 e5/providers.py create mode 100644 e5/runner.py create mode 100644 tests/test_e5_parsing.py create mode 100644 tests/test_e5_runner.py diff --git a/.env.example b/.env.example index e1d0605..92088fd 100644 --- a/.env.example +++ b/.env.example @@ -6,3 +6,5 @@ NEWSAPI_KEY=[YOUR_KEY_HERE] BARCHART_KEY=[YOUR_KEY_HERE] DB_URL=postgresql://postgres:alphaflow@localhost:5432/alphaflow KAFKA_BOOTSTRAP=localhost:9092 +OPENAI_KEY=[YOUR_KEY_HERE] +ANTHROPIC_KEY=[YOUR_KEY_HERE] diff --git a/e5/corpus.py b/e5/corpus.py index 4096523..493cba8 100644 --- a/e5/corpus.py +++ b/e5/corpus.py @@ -89,6 +89,11 @@ def __post_init__(self) -> None: }, } +OUTPUT_CONTRACT = ( + "Respond with only a JSON object mapping each ticker to its portfolio " + "weight as a decimal, weights summing to 1.0, and no other text." +) + # Risk tolerance is held fixed per archetype rather than sampled: the audit # varies one thing at a time, and a random risk score would confound # archetype effects with risk effects. @@ -158,7 +163,7 @@ def build_corpus( items: list[CorpusItem] = [] for arch, group, cond, day in itertools.product(archetypes, groups, conditions, sorted(dates)): template = PROMPT_TEMPLATES[arch][cond.name] - prompt = template.format(tickers=", ".join(group)) + prompt = template.format(tickers=", ".join(group)) + " " + OUTPUT_CONTRACT coord = f"{arch.value}|{'-'.join(group)}|{cond.name}|{day.isoformat()}" qid = f"q{stable_seed(coord):08x}" items.append( diff --git a/e5/parsing.py b/e5/parsing.py new file mode 100644 index 0000000..78910a9 --- /dev/null +++ b/e5/parsing.py @@ -0,0 +1,196 @@ +""" +Parse a model's free-text reply into an allocation over the ticker universe. + +Design note: this parses *tolerantly* and records what it found, rather than +accepting only well-formed output. If model A emits clean JSON 99% of the +time and model B 80%, discarding failures compares A's full response +distribution against B's tidiest 80% — which understates B's dispersion, +the very quantity E5 measures. Format compliance plausibly correlates with +templated advice, so silent rejection would bias h upward. + +Every parse therefore returns a ParseResult carrying the raw text, what was +extracted, the pre-normalization sum, and a status. The decision to exclude +anything is made at analysis time, not here. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from enum import Enum + +import numpy as np +from numpy.typing import NDArray + +__all__ = ["ParseStatus", "ParseResult", "parse_allocation"] + + +class ParseStatus(str, Enum): + OK = "ok" # weights found, summed to ~1 + RENORMALIZED = "renormalized" # weights found, sum was off by > tol + PARTIAL = "partial" # some tickers found, others missing + AMBIGUOUS = "ambiguous" # prose fallback, number/ticker counts disagreed + NO_WEIGHTS = "no_weights" # nothing numeric attributable to a ticker + EMPTY = "empty" # blank or whitespace reply + + +@dataclass(frozen=True) +class ParseResult: + status: ParseStatus + allocation: NDArray[np.float64] # always a valid simplex point + raw_sum: float # sum before normalization + found: tuple[str, ...] # tickers an explicit weight was found for + raw_text: str + + @property + def usable(self) -> bool: + """True when a weight was attributable to at least one ticker.""" + return self.status in ( + ParseStatus.OK, + ParseStatus.RENORMALIZED, + ParseStatus.PARTIAL, + ) + + +_JSON_BLOCK = re.compile(r"\{[^{}]*\}", re.DOTALL) +_FENCE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL) + + +def _num(x) -> float | None: + """Coerce a JSON value or string to a fraction. '25%' -> 0.25.""" + if isinstance(x, (int, float)) and not isinstance(x, bool): + return float(x) + if isinstance(x, str): + s = x.strip().rstrip("%") + try: + v = float(s) + except ValueError: + return None + return v / 100.0 if x.strip().endswith("%") else v + return None + + +def _from_json(text: str, tickers: list[str]) -> dict[str, float]: + """Try fenced blocks first, then any bare {...}. Last valid one wins, + since models often restate a final answer after reasoning aloud.""" + out: dict[str, float] = {} + upper = {t.upper(): t for t in tickers} + candidates = _FENCE.findall(text) + _JSON_BLOCK.findall(text) + for blob in candidates: + try: + obj = json.loads(blob) + except (json.JSONDecodeError, TypeError): + continue + if not isinstance(obj, dict): + continue + found = {} + for k, v in obj.items(): + key = str(k).strip().upper().lstrip("$") + if key in upper and (n := _num(v)) is not None: + found[upper[key]] = n + if found: + out = found + return out + + +def _from_text(text: str, tickers: list[str]) -> tuple[dict[str, float], bool]: + """ + Fallback for prose: 'AAPL: 40%', '40% in AAPL', tables, comma lists. + + Pairs the k-th number with the k-th ticker mention when the counts match, + which is how models actually write these lists. Two rejected approaches: + a per-ticker proximity search mis-binds on lists ("40% in AAPL, 35% in + MSFT" puts 40% within 20 chars of MSFT, claiming it twice), and greedy + nearest-neighbour with a claim set silently drops numbers whose nearest + ticker was already taken. Falls back to nearest-unclaimed only when the + counts differ. + """ + seen: list[tuple[int, str]] = [] + for t in tickers: + for m in re.finditer(rf"\${{0,1}}\b{re.escape(t)}\b", text, re.IGNORECASE): + seen.append((m.start(), t)) + if not seen: + return {}, False + seen.sort() + # first mention of each ticker, in order of appearance + order: list[str] = [] + for _, t in seen: + if t not in order: + order.append(t) + + nums: list[tuple[int, float]] = [] + for m in re.finditer(r"(\d+(?:\.\d+)?)\s*(%)|(? ParseResult: + """ + Extract an allocation over `tickers`. Always returns a valid simplex + point so downstream metrics never see a malformed vector; inspect + `status` and `raw_sum` to decide whether to include it. + + Unmentioned tickers get weight 0 when at least one weight was found. + A reply with no attributable weights falls back to equal weight with + status NO_WEIGHTS — recorded, not silently treated as a real answer. + """ + n = len(tickers) + if n == 0: + raise ValueError("empty ticker universe") + uniform = np.full(n, 1.0 / n) + + if not text or not text.strip(): + return ParseResult(ParseStatus.EMPTY, uniform, 0.0, (), text or "") + + weights = _from_json(text, tickers) + from_prose = False + if not weights: + weights, prose_ambiguous = _from_text(text, tickers) + from_prose = True + if not weights: + return ParseResult(ParseStatus.NO_WEIGHTS, uniform, 0.0, (), text) + + vec = np.array([max(0.0, weights.get(t, 0.0)) for t in tickers], dtype=float) + raw_sum = float(vec.sum()) + if raw_sum <= 0: + return ParseResult(ParseStatus.NO_WEIGHTS, uniform, raw_sum, (), text) + + alloc = vec / raw_sum + found = tuple(t for t in tickers if t in weights) + + if from_prose and prose_ambiguous: + status = ParseStatus.AMBIGUOUS + elif len(found) < n and abs(raw_sum - 1.0) <= sum_tol: + status = ParseStatus.PARTIAL + elif abs(raw_sum - 1.0) > sum_tol: + status = ParseStatus.RENORMALIZED + else: + status = ParseStatus.OK + + return ParseResult(status, alloc, raw_sum, found, text) diff --git a/e5/providers.py b/e5/providers.py new file mode 100644 index 0000000..59c36a8 --- /dev/null +++ b/e5/providers.py @@ -0,0 +1,176 @@ +""" +Providers for the E5 audit. + +The zero-key policy in the README ("every demo and test runs fully offline") +means live calls must be the exception, not the default. Three +implementations sit behind one structural protocol: + + StubProvider wraps agents.retail_ai.stub_llm_response — a + zero-variance baseline needing no network + ReplayProvider reads a recorded JSONL run, so one real audit becomes a + permanent offline fixture + LiveProvider the only one that touches the network or needs a key + +Because Provider is a Protocol rather than a base class, tests need no +mocking library: any object with .name and .complete() qualifies. +""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path +from typing import Protocol, runtime_checkable + +import numpy as np + +from agents.retail_ai import RetailQuery, stub_llm_response + +__all__ = ["Provider", "StubProvider", "ReplayProvider", "LiveProvider", "ProviderError"] + + +class ProviderError(RuntimeError): + """Raised when a provider cannot produce a reply. Recorded, not fatal.""" + + +@runtime_checkable +class Provider(Protocol): + name: str + + def complete(self, prompt: str, *, query: RetailQuery | None = None) -> str: + """Return the model's raw reply text.""" + ... + + +class StubProvider: + """ + Deterministic baseline. Emits JSON so it exercises the same parsing path + a live reply would, rather than bypassing it — otherwise stub runs would + not catch parser regressions. + + `bias` optionally tilts allocations toward one ticker, so multi-model + scenarios with known ground-truth agreement can be constructed for tests. + """ + + def __init__(self, name: str = "stub", bias: str | None = None, bias_weight: float = 0.0): + self.name = name + self.bias = bias + self.bias_weight = float(np.clip(bias_weight, 0.0, 1.0)) + + def complete(self, prompt: str, *, query: RetailQuery | None = None) -> str: + if query is None: + raise ProviderError("StubProvider needs the RetailQuery, not just the prompt") + tickers = list(query.tickers) + alloc = stub_llm_response(query, tickers) + if self.bias and self.bias in tickers and self.bias_weight > 0: + target = np.zeros(len(tickers)) + target[tickers.index(self.bias)] = 1.0 + alloc = (1 - self.bias_weight) * alloc + self.bias_weight * target + alloc = alloc / alloc.sum() + return json.dumps({t: round(float(w), 6) for t, w in zip(tickers, alloc)}) + + +class ReplayProvider: + """ + Replays a recorded run. Keyed by (model, query_id) so a multi-model + recording replays each model faithfully. + + Missing keys raise rather than falling back to a default: a silent + substitution would make a partial recording look like a complete one. + """ + + def __init__(self, path: str | Path, name: str): + self.name = name + self._by_qid: dict[str, str] = {} + p = Path(path) + if not p.exists(): + raise FileNotFoundError(f"no recording at {p}") + with p.open() as fh: + for line in fh: + line = line.strip() + if not line: + continue + rec = json.loads(line) + if rec.get("model") == name: + self._by_qid[rec["query_id"]] = rec.get("raw_text", "") + + def __len__(self) -> int: + return len(self._by_qid) + + def complete(self, prompt: str, *, query: RetailQuery | None = None) -> str: + qid = getattr(query, "query_id", None) or self._qid_hint + if qid not in self._by_qid: + raise ProviderError(f"{self.name}: no recorded reply for {qid!r}") + return self._by_qid[qid] + + # set by the runner immediately before each call + _qid_hint: str = "" + + +class LiveProvider: + """ + Real API calls. The only provider that needs a key, and the only one + excluded from the offline test path. + + Reads its key from the environment; never accepts one as a literal, and + never logs it. Retries transient failures with exponential backoff and + surfaces the rest as ProviderError so the runner can record the failure + and continue rather than losing a partial run. + """ + + def __init__( + self, + name: str, + model: str, + env_var: str, + base_url: str = "https://api.openai.com/v1/chat/completions", + temperature: float = 1.0, + max_retries: int = 3, + timeout: float = 60.0, + ): + self.name = name + self.model = model + self.env_var = env_var + self.base_url = base_url + self.temperature = temperature + self.max_retries = max_retries + self.timeout = timeout + + def _key(self) -> str: + key = os.getenv(self.env_var, "") + if not key or key.startswith("["): + raise ProviderError( + f"{self.env_var} is unset. Add it to .env — see .env.example." + ) + return key + + def complete(self, prompt: str, *, query: RetailQuery | None = None) -> str: + import urllib.error + import urllib.request + + body = json.dumps({ + "model": self.model, + "temperature": self.temperature, + "messages": [{"role": "user", "content": prompt}], + }).encode() + + last: Exception | None = None + for attempt in range(self.max_retries): + req = urllib.request.Request( + self.base_url, + data=body, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self._key()}", + }, + ) + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + payload = json.loads(resp.read()) + return payload["choices"][0]["message"]["content"] + except (urllib.error.HTTPError, urllib.error.URLError, KeyError, TimeoutError) as exc: + last = exc + if attempt < self.max_retries - 1: + time.sleep(2 ** attempt) + raise ProviderError(f"{self.name}: {type(last).__name__} after {self.max_retries} tries") diff --git a/e5/runner.py b/e5/runner.py new file mode 100644 index 0000000..e212681 --- /dev/null +++ b/e5/runner.py @@ -0,0 +1,188 @@ +""" +E5 runner: corpus -> providers -> recorded responses. + +Records to append-only JSONL so an interrupted run resumes rather than +restarts, and a partial run is still analyzable. Every response is stored +with its raw text, parse status, and pre-normalization sum; nothing is +discarded at write time (see e5.parsing for why). + +Allocations are padded to the full ticker universe. A corpus item asks about +a subset, but e5.metrics needs every response on a common basis to stack +into a (models, repeats, assets) panel — unpadded vectors of differing +length cannot be compared. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Iterable, Sequence + +import numpy as np +from numpy.typing import NDArray + +from e5.corpus import CorpusItem +from e5.parsing import ParseStatus, parse_allocation +from e5.providers import Provider, ProviderError + +__all__ = ["ResponseRecord", "run_audit", "load_records", "to_panel"] + + +@dataclass(frozen=True) +class ResponseRecord: + query_id: str + model: str + repeat: int + as_of: str + archetype: str + condition: str + tickers: list[str] + universe: list[str] + prompt: str + raw_text: str + status: str + raw_sum: float + allocation: list[float] # padded to `universe` + error: str | None = None + recorded_at: str = "" + + +def _pad(alloc: NDArray[np.float64], subset: Sequence[str], universe: Sequence[str]) -> list[float]: + """Place a subset allocation into the full-universe basis.""" + idx = {t: i for i, t in enumerate(universe)} + out = np.zeros(len(universe)) + for t, w in zip(subset, alloc): + if t in idx: + out[idx[t]] = w + s = out.sum() + return (out / s if s > 0 else np.full(len(universe), 1.0 / len(universe))).tolist() + + +def _done_keys(path: Path) -> set[tuple[str, str, int]]: + """(query_id, model, repeat) already recorded, so a resumed run skips them.""" + if not path.exists(): + return set() + seen = set() + with path.open() as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + seen.add((r["query_id"], r["model"], r.get("repeat", 0))) + except (json.JSONDecodeError, KeyError): + continue # tolerate a torn final line from an interrupted run + return seen + + +def run_audit( + items: Iterable[CorpusItem], + providers: Sequence[Provider], + universe: Sequence[str], + out_path: str | Path, + repeats: int = 1, + resume: bool = True, + progress: bool = False, +) -> Path: + """ + Send every (item x provider x repeat) and append one JSON line each. + + Provider failures are recorded with `error` set and an equal-weight + placeholder, then the run continues — losing an entire audit to one + timeout would be worse than a recorded gap that analysis can exclude. + """ + out = Path(out_path) + out.parent.mkdir(parents=True, exist_ok=True) + items = list(items) + done = _done_keys(out) if resume else set() + n_written = 0 + + with out.open("a") as fh: + for item in items: + for provider in providers: + for rep in range(repeats): + if (item.query_id, provider.name, rep) in done: + continue + # ReplayProvider keys off query_id, which RetailQuery lacks + if hasattr(provider, "_qid_hint"): + provider._qid_hint = item.query_id + + error = None + try: + raw = provider.complete(item.prompt, query=item.query) + except ProviderError as exc: + raw, error = "", str(exc) + + parsed = parse_allocation(raw, list(item.tickers)) + rec = ResponseRecord( + query_id=item.query_id, + model=provider.name, + repeat=rep, + as_of=item.as_of.isoformat(), + archetype=item.archetype.value, + condition=item.condition.name, + tickers=list(item.tickers), + universe=list(universe), + prompt=item.prompt, + raw_text=parsed.raw_text, + status=parsed.status.value, + raw_sum=parsed.raw_sum, + allocation=_pad(parsed.allocation, item.tickers, universe), + error=error, + recorded_at=datetime.now(timezone.utc).isoformat(), + ) + fh.write(json.dumps(asdict(rec)) + "\n") + fh.flush() # a crash loses at most the current call + n_written += 1 + if progress and n_written % 25 == 0: + print(f" {n_written} responses", flush=True) + return out + + +def load_records(path: str | Path) -> list[ResponseRecord]: + """Read a run back, skipping any torn final line.""" + out = [] + with Path(path).open() as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + out.append(ResponseRecord(**json.loads(line))) + except (json.JSONDecodeError, TypeError): + continue + return out + + +def to_panel( + records: Sequence[ResponseRecord], + models: Sequence[str] | None = None, + include_statuses: Sequence[str] = ( + ParseStatus.OK.value, + ParseStatus.RENORMALIZED.value, + ParseStatus.PARTIAL.value, + ), +) -> tuple[NDArray[np.float64], list[str]]: + """ + Stack records into the (n_models, n_repeats, n_assets) panel e5.metrics + expects, for one query. + + `include_statuses` is a parameter rather than a fixed policy: excluding + unparseable replies compares each model's tidiest subset, which biases + dispersion downward for models that follow instructions less reliably. + The default keeps everything with an attributable weight. + """ + usable = [r for r in records if r.status in include_statuses] + if not usable: + raise ValueError("no records with the requested statuses") + names = list(models) if models else sorted({r.model for r in usable}) + by_model = {m: [r.allocation for r in usable if r.model == m] for m in names} + missing = [m for m, v in by_model.items() if not v] + if missing: + raise ValueError(f"no usable records for {missing}") + depth = min(len(v) for v in by_model.values()) + panel = np.array([by_model[m][:depth] for m in names], dtype=float) + return panel, names diff --git a/tests/test_e5_parsing.py b/tests/test_e5_parsing.py new file mode 100644 index 0000000..d79cc08 --- /dev/null +++ b/tests/test_e5_parsing.py @@ -0,0 +1,147 @@ +""" +Tests for parsing model replies into allocations. + +Two of these encode bugs found during development that produced *plausible +wrong answers* rather than errors — the failure mode that silently corrupts +a response kernel: + + 1. Per-ticker proximity search: in "40% in AAPL, 35% in MSFT" the 40% falls + within the proximity window of MSFT and gets claimed twice. + 2. Greedy nearest-neighbour with a claim set: 35% is nearer to AAPL than to + MSFT, finds AAPL taken, and is discarded — MSFT ends up with no weight. + +Both returned a valid-looking simplex point. Only comparing against a known +answer caught them. +""" + +import numpy as np +import pytest + +from e5.parsing import ParseStatus, parse_allocation + +T = ["AAPL", "MSFT", "NVDA"] + + +def alloc(text, tickers=None): + return parse_allocation(text, tickers or T) + + +class TestJson: + def test_bare_object(self): + r = alloc('{"AAPL": 0.5, "MSFT": 0.3, "NVDA": 0.2}') + assert r.status is ParseStatus.OK + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) + + def test_fenced_block(self): + r = alloc('Sure:\n```json\n{"AAPL":0.5,"MSFT":0.3,"NVDA":0.2}\n```') + assert r.status is ParseStatus.OK + + def test_percent_strings(self): + r = alloc('{"AAPL":"50%","MSFT":"30%","NVDA":"20%"}') + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) + + def test_dollar_prefix_keys(self): + r = alloc('{"$AAPL": 0.6, "$MSFT": 0.4, "$NVDA": 0.0}') + np.testing.assert_allclose(r.allocation, [0.6, 0.4, 0.0]) + + def test_last_object_wins(self): + """Models often reason aloud then restate a final answer.""" + r = alloc('First thought {"AAPL":1.0}. On reflection: ' + '{"AAPL":0.4,"MSFT":0.4,"NVDA":0.2}') + np.testing.assert_allclose(r.allocation, [0.4, 0.4, 0.2]) + + def test_json_beats_prose(self): + r = alloc('I would put 90% in AAPL. {"AAPL":0.5,"MSFT":0.3,"NVDA":0.2}') + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) + + +class TestProse: + def test_number_before_ticker(self): + """Regression: the per-ticker proximity search bound 40% to MSFT.""" + r = alloc("Put 40% in AAPL, 35% in MSFT, and 25% in NVDA.") + assert r.status is ParseStatus.OK + np.testing.assert_allclose(r.allocation, [0.40, 0.35, 0.25]) + + def test_ticker_before_number(self): + """Regression: greedy nearest-unclaimed dropped MSFT entirely.""" + r = alloc("AAPL: 50%\nMSFT: 30%\nNVDA: 20%") + assert r.status is ParseStatus.OK + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) + + def test_decimals_without_percent(self): + r = alloc("AAPL 0.5, MSFT 0.3, NVDA 0.2") + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) + + def test_markdown_table(self): + r = alloc("| AAPL | 45% |\n| MSFT | 35% |\n| NVDA | 20% |") + np.testing.assert_allclose(r.allocation, [0.45, 0.35, 0.20]) + + +class TestStatuses: + def test_renormalizes_and_records_raw_sum(self): + r = alloc('{"AAPL":0.5,"MSFT":0.3,"NVDA":0.15}') + assert r.status is ParseStatus.RENORMALIZED + assert r.raw_sum == pytest.approx(0.95) + np.testing.assert_allclose(r.allocation.sum(), 1.0) + + def test_partial_when_tickers_omitted(self): + r = alloc('{"AAPL":1.0}') + assert r.status is ParseStatus.PARTIAL + assert r.found == ("AAPL",) + + def test_ambiguous_when_counts_disagree(self): + """A distractor number (a past return) means prose binding cannot be + trusted. Flagged rather than silently accepted — telling + 'AAPL returned 30%' from 'AAPL 50%' needs semantics, not regex.""" + r = alloc("Over 12 months AAPL returned 30%. I suggest " + "AAPL 50%, MSFT 30%, NVDA 20%.") + assert r.status is ParseStatus.AMBIGUOUS + + def test_refusal_has_no_weights(self): + r = alloc("I cannot provide financial advice.") + assert r.status is ParseStatus.NO_WEIGHTS + assert not r.usable + + def test_empty(self): + assert alloc("").status is ParseStatus.EMPTY + assert alloc(" \n ").status is ParseStatus.EMPTY + + def test_usable_flag(self): + assert alloc('{"AAPL":0.5,"MSFT":0.3,"NVDA":0.2}').usable + assert not alloc("no numbers here").usable + + +class TestInvariants: + @pytest.mark.parametrize("text", [ + '{"AAPL":0.5,"MSFT":0.3,"NVDA":0.2}', + '{"AAPL":0.5,"MSFT":0.3,"NVDA":0.15}', + '{"AAPL":1.0}', + "Put 40% in AAPL, 35% in MSFT, and 25% in NVDA.", + "I cannot provide financial advice.", + "", + '{"AAPL":-0.5,"MSFT":1.5}', + "AAPL AAPL AAPL", + ]) + def test_always_returns_a_simplex_point(self, text): + """Downstream metrics validate their input, so a malformed vector + would raise deep in the pipeline instead of being recorded here.""" + r = alloc(text) + assert r.allocation.shape == (len(T),) + assert (r.allocation >= 0).all() + np.testing.assert_allclose(r.allocation.sum(), 1.0, atol=1e-9) + + def test_negative_weights_clipped(self): + r = alloc('{"AAPL":-0.5,"MSFT":1.0,"NVDA":0.5}') + assert (r.allocation >= 0).all() + + def test_raw_text_preserved(self): + text = "some reply" + assert alloc(text).raw_text == text + + def test_empty_universe_raises(self): + with pytest.raises(ValueError, match="empty ticker"): + parse_allocation("anything", []) + + def test_case_insensitive_tickers(self): + r = alloc('{"aapl":0.5,"msft":0.3,"nvda":0.2}') + np.testing.assert_allclose(r.allocation, [0.5, 0.3, 0.2]) diff --git a/tests/test_e5_runner.py b/tests/test_e5_runner.py new file mode 100644 index 0000000..ba7210c --- /dev/null +++ b/tests/test_e5_runner.py @@ -0,0 +1,164 @@ +""" +Tests for the E5 runner and providers. + +Everything here runs offline. LiveProvider is exercised only for its +key-handling and never over the network — a test that needs a key would +break the repo's zero-key policy. +""" + +import json +from datetime import date + +import numpy as np +import pytest + +from e5.corpus import build_corpus +from e5.providers import ( + LiveProvider, + Provider, + ProviderError, + ReplayProvider, + StubProvider, +) +from e5.runner import load_records, run_audit, to_panel + +U = ["AAPL", "MSFT", "NVDA", "TSLA"] +DATES = [date(2026, 8, 1)] + + +@pytest.fixture +def items(): + return build_corpus(U, DATES, max_ticker_groups=2) + + +@pytest.fixture +def run(tmp_path, items): + provs = [StubProvider("a"), StubProvider("b", bias="AAPL", bias_weight=0.6)] + return run_audit(items, provs, U, tmp_path / "run.jsonl", repeats=2), items, provs + + +class TestProviders: + def test_stub_satisfies_protocol(self): + assert isinstance(StubProvider(), Provider) + + def test_stub_emits_parseable_json(self, items): + """The stub goes through the same parser a live reply does, so stub + runs still catch parser regressions.""" + raw = StubProvider().complete(items[0].prompt, query=items[0].query) + obj = json.loads(raw) + assert set(obj) == set(items[0].tickers) + assert sum(obj.values()) == pytest.approx(1.0, abs=1e-5) + + def test_stub_is_deterministic(self, items): + p = StubProvider() + assert p.complete(items[0].prompt, query=items[0].query) == p.complete( + items[0].prompt, query=items[0].query + ) + + def test_bias_tilts_allocation(self, items): + base = json.loads(StubProvider().complete(items[0].prompt, query=items[0].query)) + tilt = json.loads( + StubProvider(bias="AAPL", bias_weight=0.8).complete( + items[0].prompt, query=items[0].query + ) + ) + assert tilt["AAPL"] > base["AAPL"] + + def test_stub_requires_the_query(self): + with pytest.raises(ProviderError, match="RetailQuery"): + StubProvider().complete("a bare prompt") + + def test_live_provider_never_calls_without_a_key(self, monkeypatch): + monkeypatch.delenv("E5_TEST_KEY", raising=False) + p = LiveProvider("x", "some-model", "E5_TEST_KEY") + with pytest.raises(ProviderError, match="E5_TEST_KEY is unset"): + p.complete("hello") + + def test_live_provider_rejects_placeholder_key(self, monkeypatch): + monkeypatch.setenv("E5_TEST_KEY", "[YOUR_KEY_HERE]") + with pytest.raises(ProviderError, match="unset"): + LiveProvider("x", "m", "E5_TEST_KEY").complete("hello") + + +class TestReplay: + def test_round_trips_a_recording(self, run): + path, items, _ = run + rp = ReplayProvider(path, "a") + assert len(rp) > 0 + rp._qid_hint = items[0].query_id + assert json.loads(rp.complete(items[0].prompt, query=items[0].query)) + + def test_missing_key_raises_rather_than_defaulting(self, run): + """A silent fallback would make a partial recording look complete.""" + path, _, _ = run + rp = ReplayProvider(path, "a") + rp._qid_hint = "nonexistent" + with pytest.raises(ProviderError, match="no recorded reply"): + rp.complete("p") + + def test_unknown_file(self, tmp_path): + with pytest.raises(FileNotFoundError): + ReplayProvider(tmp_path / "nope.jsonl", "a") + + +class TestRunAudit: + def test_record_count(self, run): + path, items, provs = run + assert len(load_records(path)) == len(items) * len(provs) * 2 + + def test_allocations_padded_to_universe(self, run): + """Items ask about ticker subsets; metrics need a common basis.""" + for r in load_records(run[0]): + assert len(r.allocation) == len(U) + assert sum(r.allocation) == pytest.approx(1.0, abs=1e-6) + + def test_resume_skips_completed_work(self, run): + path, items, provs = run + before = len(load_records(path)) + run_audit(items, provs, U, path, repeats=2) + assert len(load_records(path)) == before + + def test_provider_failure_is_recorded_not_fatal(self, tmp_path, items): + class Broken: + name = "broken" + + def complete(self, prompt, *, query=None): + raise ProviderError("simulated outage") + + path = run_audit(items[:3], [Broken(), StubProvider("ok")], U, + tmp_path / "r.jsonl") + recs = load_records(path) + assert any(r.error for r in recs), "failure not recorded" + assert any(r.error is None for r in recs), "run aborted on failure" + + def test_raw_text_is_preserved(self, run): + assert all(r.raw_text for r in load_records(run[0]) if not r.error) + + +class TestToPanel: + def test_shape(self, run): + recs = load_records(run[0]) + one = [r for r in recs if r.query_id == recs[0].query_id] + panel, names = to_panel(one) + assert panel.shape == (2, 2, len(U)) + assert names == ["a", "b"] + + def test_panel_rows_are_simplex_points(self, run): + recs = load_records(run[0]) + panel, _ = to_panel([r for r in recs if r.query_id == recs[0].query_id]) + np.testing.assert_allclose(panel.sum(axis=-1), 1.0, atol=1e-6) + + def test_status_filter_is_a_parameter(self, run): + """Excluding sloppy replies is an analysis choice, not baked in.""" + recs = load_records(run[0]) + one = [r for r in recs if r.query_id == recs[0].query_id] + with pytest.raises(ValueError, match="no records"): + to_panel(one, include_statuses=("empty",)) + + def test_feeds_metrics(self, run): + from e5.metrics import crowding_index, spread + + recs = load_records(run[0]) + panel, _ = to_panel([r for r in recs if r.query_id == recs[0].query_id]) + assert 0.0 <= crowding_index(panel) <= 1.0 + assert spread(panel) >= 0.0