feat(metrics): populate token usage, add latency percentiles, record …

benchmarks/common/llm_client.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
 import json
 import logging
 import os
@@ -26,6 +27,39 @@
 
 T = TypeVar("T")
 
+# -----------------------------------------------------------------------------
+# Token accounting
+# -----------------------------------------------------------------------------
+# Both SDKs already report exact usage on every response, so nothing needs to be
+# re-tokenized locally. Accounting is per-task rather than per-client because one
+# LLMClient is shared by max_workers concurrent items: asyncio copies the context
+# at Task creation, so sibling items accumulate into their own bucket.
+
+_EMPTY_USAGE = {"prompt_tokens": 0, "completion_tokens": 0, "llm_calls": 0}
+
+_usage: contextvars.ContextVar[dict[str, int] | None] = contextvars.ContextVar("llm_usage", default=None)
+
+
+def start_usage() -> None:
+    """Begin token accounting for the current eval item. Call once per item."""
+    _usage.set(dict(_EMPTY_USAGE))
+
+
+def get_usage() -> dict[str, int]:
+    """Tokens billed to the current item. Zeros if accounting was never started."""
+    return dict(_usage.get() or _EMPTY_USAGE)
+
+
+def _record_usage(resp: Any) -> None:
+    """Accumulate SDK-reported usage; OpenAI and Anthropic name the fields differently."""
+    bucket = _usage.get()
+    usage = getattr(resp, "usage", None)
+    if bucket is None or usage is None:
+        return
+    bucket["prompt_tokens"] += getattr(usage, "prompt_tokens", 0) or getattr(usage, "input_tokens", 0) or 0
+    bucket["completion_tokens"] += getattr(usage, "completion_tokens", 0) or getattr(usage, "output_tokens", 0) or 0
+    bucket["llm_calls"] += 1
+
 
 class LLMClient:
     """Async LLM client with retry logic and rate limiting.
@@ -173,6 +207,7 @@ async def _generate_openai(self, system: str, user: str, temperature: float, max
                         ),
                         timeout=self.timeout,
                     )
+                _record_usage(resp)
                 content = resp.choices[0].message.content
                 if content is None:
                     logger.warning(
@@ -208,6 +243,7 @@ async def _generate_anthropic(self, system: str, user: str, temperature: float,
                         self._client.messages.create(**kwargs),
                         timeout=self.timeout,
                     )
+                _record_usage(resp)
                 content = resp.content[0].text if resp.content else ""
                 return content.strip()
             except asyncio.TimeoutError:
@@ -275,6 +311,7 @@ async def _generate_structured_openai(
                         ),
                         timeout=self.timeout,
                     )
+                _record_usage(resp)
                 raw = resp.choices[0].message.content
                 if not raw:
                     if attempt < self.max_retries - 1:
@@ -335,6 +372,7 @@ async def _generate_structured_anthropic(
                         self._client.messages.create(**kwargs),
                         timeout=self.timeout,
                     )
+                _record_usage(resp)
                 raw = resp.content[0].text if resp.content else ""
                 raw = raw.strip()
 

benchmarks/common/metrics.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import math
 import statistics
 from collections import defaultdict
 from typing import Any
@@ -143,6 +144,61 @@ def compute_overall_metrics(
     return metrics
 
 
+def _percentile(values: list[float], q: float) -> float:
+    """Nearest-rank percentile — unambiguous, and what o11y tooling reports."""
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    idx = max(0, math.ceil(q / 100 * len(ordered)) - 1)
+    return ordered[idx]
+
+
+def compute_cost_metrics(
+    evaluations: list[dict[str, Any]],
+    cutoff_label: str | None = None,
+) -> dict[str, float]:
+    """Latency percentiles and token means for a run.
+
+    Latency is reported as p50/p95 rather than a mean: retrieval latency is
+    right-skewed, so a mean is dominated by the tail and is not comparable to a
+    published p50. Token figures cover the answering model only (the context cost
+    of a query), not the judge, which does not exist in production.
+
+    Note that these latencies are measured under whatever concurrency the run
+    used, so they are only interpretable alongside ``max_workers`` in metadata.
+    """
+    # The benchmarks disagree on where this lives: locomo and beam write it at
+    # item top level, longmemeval nests it under "retrieval" (the schema
+    # location). Read both so this works on already-published results too.
+    search: list[float] = []
+    for e in evaluations:
+        ms = e.get("search_latency_ms")
+        if ms is None:
+            ms = (e.get("retrieval") or {}).get("search_latency_ms")
+        if ms:
+            search.append(ms)
+
+    prompt: list[float] = []
+    completion: list[float] = []
+    for e in evaluations:
+        src = (e.get("cutoff_results", {}) or {}).get(cutoff_label, {}) if cutoff_label else e
+        if src.get("prompt_tokens") is not None:
+            prompt.append(src["prompt_tokens"])
+            completion.append(src.get("completion_tokens") or 0)
+
+    out = {
+        "search_latency_p50_ms": round(_percentile(search, 50), 1),
+        "search_latency_p95_ms": round(_percentile(search, 95), 1),
+        "search_latency_mean_ms": round(statistics.mean(search), 1) if search else 0.0,
+        "n_latency_samples": len(search),
+    }
+    if prompt:
+        out["mean_prompt_tokens"] = round(statistics.mean(prompt), 1)
+        out["mean_completion_tokens"] = round(statistics.mean(completion), 1)
+        out["n_token_samples"] = len(prompt)
+    return out
+
+
 def compute_kendall_tau_b(predicted_order: list[int], reference_order: list[int]) -> float:
     """Compute Kendall tau-b rank correlation coefficient.
 

benchmarks/common/selfcheck.py

@@ -0,0 +1,120 @@
+"""
+Self-check for token accounting and cost metrics
+================================================
+
+No test framework — the repo has none and this needs no new dependency.
+
+    python -m benchmarks.common.selfcheck
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+from benchmarks.common.llm_client import _record_usage, get_usage, start_usage
+from benchmarks.common.metrics import _percentile, compute_cost_metrics
+
+
+class _FakeUsage:
+    def __init__(self, prompt: int, completion: int) -> None:
+        self.prompt_tokens = prompt
+        self.completion_tokens = completion
+
+
+class _FakeAnthropicUsage:
+    def __init__(self, prompt: int, completion: int) -> None:
+        self.input_tokens = prompt
+        self.output_tokens = completion
+
+
+class _FakeResp:
+    def __init__(self, usage: object) -> None:
+        self.usage = usage
+
+
+def check_percentile() -> None:
+    assert _percentile([], 50) == 0.0
+    assert _percentile([5], 50) == 5
+    # nearest-rank: p50 of 1..10 is the 5th value, p95 is the 10th
+    ten = list(range(1, 11))
+    assert _percentile(ten, 50) == 5, _percentile(ten, 50)
+    assert _percentile(ten, 95) == 10, _percentile(ten, 95)
+    assert _percentile(ten, 100) == 10
+    # unsorted input must not change the answer
+    assert _percentile(list(reversed(ten)), 50) == 5
+    print("ok  percentile")
+
+
+def check_provider_field_names() -> None:
+    start_usage()
+    _record_usage(_FakeResp(_FakeUsage(100, 10)))
+    _record_usage(_FakeResp(_FakeAnthropicUsage(50, 5)))
+    _record_usage(_FakeResp(None))  # a response without usage must not explode
+    u = get_usage()
+    assert u == {"prompt_tokens": 150, "completion_tokens": 15, "llm_calls": 2}, u
+    print("ok  openai + anthropic field names, missing usage tolerated")
+
+
+def check_untracked_is_zero_not_crash() -> None:
+    # A call made outside an eval item must be silently ignored, not counted
+    # against whichever item happens to be running.
+    import benchmarks.common.llm_client as lc
+
+    lc._usage.set(None)
+    _record_usage(_FakeResp(_FakeUsage(999, 999)))
+    assert get_usage() == {"prompt_tokens": 0, "completion_tokens": 0, "llm_calls": 0}
+    print("ok  untracked calls are not billed to anyone")
+
+
+async def check_concurrent_isolation() -> None:
+    """The real risk: one shared LLMClient, max_workers concurrent items."""
+
+    async def item(n: int) -> dict:
+        start_usage()
+        for _ in range(n):
+            await asyncio.sleep(0)  # force interleaving between records
+            _record_usage(_FakeResp(_FakeUsage(10, 1)))
+        return get_usage()
+
+    results = await asyncio.gather(*(item(n) for n in range(1, 11)))
+    for n, got in enumerate(results, start=1):
+        assert got["prompt_tokens"] == 10 * n, (n, got)
+        assert got["llm_calls"] == n, (n, got)
+    print("ok  concurrent items do not mix tokens")
+
+
+def check_cost_metrics() -> None:
+    evaluations = [
+        {
+            "retrieval": {"search_latency_ms": ms},
+            "cutoff_results": {"top_200": {"prompt_tokens": tok, "completion_tokens": 10}},
+        }
+        for ms, tok in zip([100, 200, 300, 400, 500], [1000, 2000, 3000, 4000, 5000])
+    ]
+    m = compute_cost_metrics(evaluations, cutoff_label="top_200")
+    assert m["search_latency_p50_ms"] == 300, m
+    assert m["search_latency_p95_ms"] == 500, m
+    assert m["search_latency_mean_ms"] == 300.0, m
+    assert m["mean_prompt_tokens"] == 3000.0, m
+    assert m["n_token_samples"] == 5, m
+
+    # A run made before this change carries no token fields; it must still report
+    # latency rather than failing or inventing zeros.
+    legacy = [{"retrieval": {"search_latency_ms": 100}}]
+    m2 = compute_cost_metrics(legacy, cutoff_label="top_200")
+    assert "mean_prompt_tokens" not in m2, m2
+    assert m2["search_latency_p50_ms"] == 100, m2
+    print("ok  cost metrics, and legacy results still aggregate")
+
+
+def main() -> None:
+    check_percentile()
+    check_provider_field_names()
+    check_untracked_is_zero_not_crash()
+    asyncio.run(check_concurrent_isolation())
+    check_cost_metrics()
+    print("\nall checks passed")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/locomo/run.py

@@ -42,9 +42,9 @@
 from dotenv import load_dotenv
 from tqdm import tqdm
 
-from benchmarks.common.llm_client import LLMClient
+from benchmarks.common.llm_client import LLMClient, get_usage, start_usage
 from benchmarks.common.mem0_client import Mem0Client, format_search_results
-from benchmarks.common.metrics import compute_overall_metrics
+from benchmarks.common.metrics import compute_cost_metrics, compute_overall_metrics
 from benchmarks.common.schema import (
     CutoffResult,
     EvalItem,
@@ -412,6 +412,8 @@ async def process_question(
     category = qa["category"]
     answer = str(qa["answer"])
 
+    start_usage()
+
     # --- Search ---
     start = time.monotonic()
     search_results = await mem0.search(question, user_id, top_k=top_k, score_debug=score_debug)
@@ -463,7 +465,11 @@ async def process_question(
 
         # Generate answer
         gen_prompt = get_answer_generation_prompt(question, sliced, reference_date=reference_date_human, user_profile=user_profile)
+        before = get_usage()
         generated_answer = await answerer.generate(system="", user=gen_prompt)
+        answer_usage = {
+            k: get_usage()[k] - before[k] for k in ("prompt_tokens", "completion_tokens")
+        }
         if "ANSWER:" in generated_answer:
             generated_answer = generated_answer.rsplit("ANSWER:", 1)[-1].strip()
 
@@ -492,9 +498,16 @@ async def process_question(
             "generated_answer": generated_answer,
             "memories_evaluated": len(sliced),
             "reason": raw.get("reasoning", "") if isinstance(raw, dict) else "",
+            # Context sent to the answering model at this cutoff. This is the
+            # "tokens per query" figure; judge tokens are deliberately excluded.
+            "prompt_tokens": answer_usage["prompt_tokens"],
+            "completion_tokens": answer_usage["completion_tokens"],
         }
 
     result["cutoff_results"] = cutoff_results
+    # Whole-item spend, answerer + judge across every cutoff. This is the cost of
+    # running the benchmark, which is a different quantity from the line above.
+    result["usage"] = get_usage()
     return result
 
 
@@ -652,6 +665,7 @@ def compute_locomo_metrics(evaluations: list[dict], cutoffs: list[int]) -> dict:
                 "avg_score": statistics.mean(scores) * 100 if scores else 0.0,
             },
             "by_category": cat_metrics,
+            "cost": compute_cost_metrics(evaluations, cutoff_label=label),
         }
     return metrics_by_cutoff
 
@@ -817,6 +831,7 @@ async def judge_one(qid: str, conv_idx: int, qi: int, qa: dict) -> None:
                 "top_k": args.top_k,
                 "top_k_cutoffs": [cutoff_label(c) for c in cutoffs],
                 "total_questions": len(all_evaluations),
+                "max_workers": args.max_workers,
                 "categories": categories,
                 "evaluate_only": True,
             },
@@ -963,6 +978,7 @@ async def process_conversation(conv_idx: int):
                     "top_k": args.top_k,
                     "top_k_cutoffs": [cutoff_label(c) for c in cutoffs],
                     "total_questions": len(all_evaluations),
+                "max_workers": args.max_workers,
                     "categories": categories,
                 },
                 "metrics_by_cutoff": metrics,