From c6d6011998ca33850bc908000213e1a91315c6b4 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:11:36 -0400 Subject: [PATCH 1/3] feat(client): cache-aware token accounting and a reusable span lifecycle Six handlers each hand-write their own usage attributes today, which is how they came to disagree. This adds the shared writers they will move onto, so the numbers have one author. parse_usage grows cache handling. It read three key-pair spellings and dropped every cache field, so Anthropic undercounted input by the whole cached portion of a call: a turn that reads 19,971 tokens from cache and writes 3,580 more reports input_tokens: 3, and 3 is what we billed a dashboard on. It now folds every accepted spelling, including Bedrock Converse's cacheWriteInputTokens, and reports the breakdown in input_details. The fold stays provider-blind, which is the contract handlers must respect: return raw usage with the cache fields intact, or an input figure that already includes cache with the fields omitted. Returning a pre-folded input alongside the fields double-counts. add_cached_tokens_to_input applies the Anthropic-shaped rule at the call site, and lang_chain_span_usage the LangChain-shaped one, because the direction differs per provider and centralising it would silently double-count for two providers out of three. SpanUsage is the type that means "the folding is already done". set_usage_span_attributes writes all seven attributes every time, zeros included, because an absent attribute drops a span from every query that groups on usage, which reads as "no cached tokens" rather than "this handler forgot to say". It also owns the two OpenLLMetry aliases, which previously lived beside the completion text and were computed off Anthropic's cache-excluding input field, so they disagreed with the canonical numbers on the same span. RunUsage counts whether any turn reported usage rather than testing the total for zero, so a failed run can put its partial spend on the root while a run that never completed a call correctly says nothing. All-zero attributes would assert the run cost nothing. number_or_zero replaces bare int(...), which raised on None. An emitted NaN is worse than an emitted 0, because the metric guard tests `> 0` and that is false for NaN, so the metric vanishes instead of reading low. end_span_once makes the streaming teardown idempotent and marks an abandoned stream without failing it. Stopping early is a normal thing for a consumer to do, and LaunchDarkly's own metrics record neither success nor error for it, so ERROR would put two dashboards in disagreement about one run. It tracks id(span) because an OTel span is not guaranteed hashable. UsageDict gains input_details, which broke graph.py's UsageDict(**dict) splat. Now built with named arguments, so the next member added here cannot silently arrive from a dict with no business filling it. --- .../src/launchdarkly_ai_server/__init__.py | 20 + .../src/launchdarkly_ai_server/graph.py | 10 +- .../src/launchdarkly_ai_server/types.py | 15 + .../src/launchdarkly_ai_server/utils.py | 308 +++++++++++++- packages/client/tests/test_span_usage.py | 382 ++++++++++++++++++ 5 files changed, 721 insertions(+), 14 deletions(-) create mode 100644 packages/client/tests/test_span_usage.py diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 5d1f2ee..feaeaec 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -34,6 +34,7 @@ HandlerResult, HandlerStreamEvent, InitClientOptions, + InputTokenDetails, JudgeResult, JudgeRunResult, JudgeTask, @@ -56,15 +57,24 @@ ) from .types_validation import parse_ai_config from .utils import ( + RunUsage, + SpanUsage, + add_cached_tokens_to_input, create_handler, + create_run_usage, + end_span_once, + lang_chain_span_usage, make_track_data, normalize_mode, + number_or_zero, parse_json_with_possible_fences, parse_template, parse_usage, set_ld_span_attributes, + set_model_identity_attributes, set_openllmetry_completion, set_openllmetry_prompt, + set_usage_span_attributes, to_ld_context, ) @@ -100,7 +110,17 @@ "StreamDoneEvent", "StreamEvent", "TrackData", + "InputTokenDetails", + "RunUsage", + "SpanUsage", "UsageDict", + "add_cached_tokens_to_input", + "create_run_usage", + "end_span_once", + "lang_chain_span_usage", + "number_or_zero", + "set_model_identity_attributes", + "set_usage_span_attributes", "VariationMeta", # utils "create_handler", diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index 0a397cf..46d606f 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -668,7 +668,15 @@ async def invoke( return ProviderGraphResponse( response=final_response, - usage=UsageDict(**total_usage), + # Named rather than splatted, so a new UsageDict member cannot silently arrive + # here from a dict that has no business filling it. Graph totals carry no cache + # breakdown: they are a sum across nodes, and the per-node detail is on the node's + # own spans. + usage=UsageDict( + input=total_usage["input"], + output=total_usage["output"], + total=total_usage["total"], + ), judge_results=judge_results, ) diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 55e863f..57f8134 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -192,11 +192,26 @@ def has_stream(self) -> bool: # --------------------------------------------------------------------------- +@dataclass +class InputTokenDetails: + """The cache breakdown behind an inclusive ``input`` figure. + + Present only when the provider reported at least one cache field. ``uncached + cache_read + + cache_creation`` equals :attr:`UsageDict.input`. + """ + + uncached: int = 0 + cache_read: int = 0 + cache_creation: int = 0 + + @dataclass class UsageDict: input: int = 0 output: int = 0 total: int = 0 + #: Cache breakdown, when the provider reported one. See :class:`InputTokenDetails`. + input_details: InputTokenDetails | None = None @dataclass diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 64648dc..81e83fd 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -2,6 +2,7 @@ import json import re +from dataclasses import dataclass, field from typing import Any, Literal from .types import ( @@ -49,27 +50,308 @@ def normalize_mode(mode: str | None) -> Literal["agent", "messages"]: return "agent" if mode == "agent" else "messages" -_USAGE_KEY_PAIRS = [ - ("input_tokens", "output_tokens"), - ("inputTokens", "outputTokens"), - ("input", "output"), -] +def number_or_zero(value: Any) -> int: + """Coerces a provider-reported token count to a finite number, defaulting to 0. + + Provider SDKs report usage as loosely typed bags where a field may be absent, ``None``, or (in + streaming paths) a partially populated value. Every count that reaches a span or a usage total + should pass through here. + An emitted 0 is bad. An emitted ``NaN`` is worse, because the metric guard tests whether the + total is greater than zero, and that test is false for ``NaN``, so the metric is dropped + silently rather than reported low. -def parse_usage(usage: dict[str, Any]) -> dict[str, int]: + Replaces the bare ``int(...)`` this module used to do, which raised on ``None``. """ - Normalises token counts from three possible key-pair shapes. - Always computes ``total`` from ``input + output``; never trusts a ``total`` - field from the raw object. + if value is None or isinstance(value, bool): + return 0 + try: + parsed = float(value) + except (TypeError, ValueError): + return 0 + if parsed != parsed or parsed in (float("inf"), float("-inf")): + return 0 + return int(parsed) + + +#: Provider usage field names, tried in order; the first row whose input *and* output keys are both +#: present wins. Cache fields list every spelling we accept, because providers disagree (Anthropic +#: ``cache_creation_input_tokens``, Bedrock Converse ``cacheWriteInputTokens``) and an unlisted +#: spelling is silently dropped from the total. +#: +#: Contract for handlers: this fold is provider-blind, so it always adds the cache fields on top of +#: the input field. A handler must therefore return *either* raw usage with the cache fields intact +#: (Anthropic-style reporting, where cache is genuinely additional), *or* an input figure that +#: already includes cache with the cache fields omitted (OpenAI- and LangChain-style reporting). +#: Returning a pre-folded input *alongside* the cache fields double-counts the cached portion. +_USAGE_KEY_PAIRS: list[tuple[str, str, tuple[str, ...], tuple[str, ...]]] = [ + ( + "input_tokens", + "output_tokens", + ("cache_read_input_tokens",), + ("cache_creation_input_tokens",), + ), + ( + "inputTokens", + "outputTokens", + ("cacheReadInputTokens",), + ("cacheCreationInputTokens", "cacheWriteInputTokens"), + ), + ("input", "output", (), ()), +] + + +def _read_cache_field(usage: dict[str, Any], keys: tuple[str, ...]) -> int: + """Sums every accepted spelling of a cache field, so an alias never silently reads as 0.""" + return sum(number_or_zero(usage.get(key)) for key in keys) + + +def parse_usage(usage: dict[str, Any]) -> dict[str, Any]: + """Normalises token counts from three possible key-pair shapes, folding cache tokens in. + + Always computes ``total`` from ``input + output``; never trusts a ``total`` field from the raw + object. A provider-supplied total can include tokens that appear in neither input nor output, + which would make the figure derivable on five handlers and not on the sixth. + + ``input`` is the inclusive total: every accepted cache spelling is added on top of the reported + input figure. See the contract note on ``_USAGE_KEY_PAIRS``. + + When the matched row defines cache keys and at least one of them is present, the result also + carries ``input_details`` with ``uncached``, ``cache_read`` and ``cache_creation``. + + An unrecognised bag returns all zeros rather than raising. """ - for input_key, output_key in _USAGE_KEY_PAIRS: + for input_key, output_key, cache_read_keys, cache_creation_keys in _USAGE_KEY_PAIRS: if input_key in usage and output_key in usage: - inp = int(usage[input_key]) - out = int(usage[output_key]) - return {"input": inp, "output": out, "total": inp + out} + cache_read = _read_cache_field(usage, cache_read_keys) + cache_creation = _read_cache_field(usage, cache_creation_keys) + uncached = number_or_zero(usage[input_key]) + inp = uncached + cache_read + cache_creation + out = number_or_zero(usage[output_key]) + result: dict[str, Any] = {"input": inp, "output": out, "total": inp + out} + has_details = bool(cache_read_keys or cache_creation_keys) and any( + key in usage for key in (*cache_read_keys, *cache_creation_keys) + ) + if has_details: + result["input_details"] = { + "uncached": uncached, + "cache_read": cache_read, + "cache_creation": cache_creation, + } + return result return {"input": 0, "output": 0, "total": 0} +@dataclass +class SpanUsage: + """The provider-neutral token counts a span reports, after the caller applied its cache rule. + + ``input`` is the *inclusive* total: whether cache tokens were already counted inside the + provider's input figure (OpenAI, LangChain) or reported alongside it and folded in by the caller + (Anthropic), by the time a value reaches this type the folding is done. + + That is the whole point of the type existing. See :func:`add_cached_tokens_to_input`. + """ + + input: int = 0 + output: int = 0 + cache_read: int = 0 + cache_creation: int = 0 + + +def add_cached_tokens_to_input(raw_usage: dict[str, Any]) -> SpanUsage: + """Adds a provider's cached-token counts into its input total. + + Providers describe the same call two different ways. Anthropic reports cache reads and cache + writes as buckets *beside* ``input_tokens``, so ``input_tokens`` counts only the new tokens: a + turn that read 19,971 tokens from cache and wrote 3,580 more reports ``input_tokens: 3``, when + the model actually processed 23,554. OpenAI and LangChain report one input figure that already + contains the cached tokens, with a subset breakdown alongside. + + This is for the first kind. Applying it to the second would count the cached tokens twice, which + is why the rule lives at the call site and not inside :func:`set_usage_span_attributes`. + + Not named for Anthropic: Bedrock Converse reports the same way, and the shape is the reason it + applies, not the vendor. + """ + cache_read = number_or_zero(raw_usage.get("cache_read_input_tokens")) + cache_creation = number_or_zero(raw_usage.get("cache_creation_input_tokens")) + return SpanUsage( + input=number_or_zero(raw_usage.get("input_tokens")) + + cache_read + + cache_creation, + output=number_or_zero(raw_usage.get("output_tokens")), + cache_read=cache_read, + cache_creation=cache_creation, + ) + + +@dataclass +class RunUsage: + """A run's accumulated token spend, for the ``invoke_agent`` root. + + The root is the only span carrying ``launchdarkly.*`` and the ``feature_flag`` event, so it is + the span a config-scoped query finds. Without a run total on it, that query returns nothing: + summing the children requires having already found them. + + Provider-blind on purpose. It sums four numbers and remembers whether anything was added. The + cache-folding rule that differs per provider has already been applied by the time a + :class:`SpanUsage` exists, so each handler maps its own bag first and this stays shared. + + The two Anthropic handlers deliberately do not use it. Their run total has to stay in Anthropic's + own field names with the cache buckets *unfolded*, because it is also the handler's return value + and :func:`parse_usage` folds it exactly once; handing back a :class:`SpanUsage` there would + count the cache twice. + """ + + total: SpanUsage = field(default_factory=SpanUsage) + _turns: int = 0 + + @property + def reported(self) -> bool: + """Whether any turn reported usage at all. + + The failure path needs this to tell "no call completed" from "a call completed and reported + zero". Only the second may be written to a span: all-zero attributes assert the run cost + nothing, which a run whose first call died mid-flight cannot claim, whereas an absent + attribute correctly says "unknown". + + Counted rather than derived by testing the total for zero, so a provider reporting a + genuinely empty bag stays distinguishable. + """ + return self._turns > 0 + + def add(self, turn: SpanUsage | None) -> None: + """Adds one turn. ``None`` is a no-op and does not count as reported.""" + if turn is None: + return + self._turns += 1 + self.total.input += turn.input + self.total.output += turn.output + self.total.cache_read += turn.cache_read + self.total.cache_creation += turn.cache_creation + + +def create_run_usage() -> RunUsage: + """A fresh run accumulator. Present so call sites read the same as the TypeScript SDK's.""" + return RunUsage() + + +def lang_chain_span_usage(usage: dict[str, Any] | None) -> SpanUsage | None: + """One LangChain ``usage_metadata`` bag as :class:`SpanUsage`, or ``None`` when it reports nothing. + + LangChain already includes cached tokens in ``input_tokens``, so nothing is folded here. Shared + because both LangChain handlers read the identical shape, and because ``None`` for an empty bag + is what keeps a turn the provider said nothing about from registering as reported: the callback + path hands over ``{}`` rather than nothing when a provider omits usage entirely. + """ + if not usage: + return None + if usage.get("input_tokens") is None and usage.get("output_tokens") is None: + return None + details = usage.get("input_token_details") or {} + return SpanUsage( + input=number_or_zero(usage.get("input_tokens")), + output=number_or_zero(usage.get("output_tokens")), + cache_read=number_or_zero(details.get("cache_read")), + cache_creation=number_or_zero(details.get("cache_creation")), + ) + + +def set_usage_span_attributes(span: Any, usage: SpanUsage) -> None: + """Writes the OpenTelemetry ``gen_ai.usage.*`` token attributes onto a span. + + Always writes all seven attributes, every time, including zeros. Consumers group and aggregate + on the complete set, and an *absent* attribute drops a span from those queries entirely, which + reads as "this call had no cached tokens" when it actually means "this handler forgot to say". A + provider with no cache-creation concept still reports 0. + + ``total`` is always ``input + output``. Providers that report their own total are deliberately + not trusted here. + + This helper does **not** know each provider's cache accounting, and must not learn it. Anthropic + reports cache tokens alongside input, so its callers fold them in; OpenAI and LangChain already + count cache tokens inside input, so their callers pass the reported figure through untouched. + Both arrive here as an inclusive ``input``. Centralising that rule would silently double-count + for two providers out of three. + + The last two are OpenLLMetry aliases for the same two numbers. They belong here rather than + beside the completion text so they cannot disagree with the canonical attributes above: + computed at a call site off the raw input field, the alias undercounted Anthropic by every + cached token, because on Anthropic that field excludes the cache. One writer, one number. + """ + inp = number_or_zero(usage.input) + out = number_or_zero(usage.output) + span.set_attribute("gen_ai.usage.input_tokens", inp) + span.set_attribute("gen_ai.usage.output_tokens", out) + span.set_attribute("gen_ai.usage.total_tokens", inp + out) + span.set_attribute( + "gen_ai.usage.cache_read.input_tokens", number_or_zero(usage.cache_read) + ) + span.set_attribute( + "gen_ai.usage.cache_creation.input_tokens", number_or_zero(usage.cache_creation) + ) + span.set_attribute("gen_ai.usage.prompt_tokens", inp) + span.set_attribute("gen_ai.usage.completion_tokens", out) + + +def set_model_identity_attributes( + span: Any, + provider_name: str, + request_model: str, + legacy_system: str | None = None, +) -> None: + """Writes the model identity attributes that every LLM span carries. + + Both spellings of the provider key are emitted on purpose. ``gen_ai.system`` is the pre-1.37 + semconv name and is what handlers shipped before the span hierarchy landed; + ``gen_ai.provider.name`` is the current name. Emitting only the new key would silently break + dashboards written against the old one, and emitting only the old one leaves us off-spec, so + both go out until the next major. + + ``legacy_system`` exists because the two keys do not always want the same value. The LangChain + handlers ship ``gen_ai.system = 'langchain'``, but ``gen_ai.provider.name`` means *who served + the model* and its semconv enum has no ``langchain`` member, so those handlers pass the real + provider for the new key and keep the framework name on the old one. + """ + span.set_attribute( + "gen_ai.system", provider_name if legacy_system is None else legacy_system + ) + span.set_attribute("gen_ai.provider.name", provider_name) + span.set_attribute("gen_ai.request.model", request_model) + + +def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None: + """Ends a span exactly once, even when the caller cannot know whether an earlier path ended it. + + The streaming handlers need this: a consumer that breaks out of ``async for``, or throws inside + the loop body, makes the generator run its ``finally`` without ever entering ``except``, so the + cleanup path and the success path can both reach the same span. Ending twice is silently ignored + by the OTel SDK but recorded as a diagnostic error, and would also hide a genuine leak, so the + guard is explicit. + + *tracker* holds ``id(span)`` rather than the span itself, because an OTel span is not guaranteed + hashable across implementations. + + An abandoned span is marked with ``launchdarkly.stream.abandoned`` and deliberately left at + ``UNSET`` rather than ``ERROR``. Stopping early is a normal thing for a consumer to do, such as + rendering enough of a response and moving on, and nothing failed. Marking it ``ERROR`` would + also put the trace at odds with LaunchDarkly's own metrics, which record neither a success nor + an error for an abandoned stream: two dashboards would disagree about the same run. The + attribute keeps abandonment findable without asserting a failure. + + Ending the span is not always enough. Two handlers also hold a vendor generator or run that must + be closed or cancelled in the same ``finally``. See TELEMETRY-CONTRACT.md section 6. + """ + key = id(span) + if key in tracker: + return + tracker.add(key) + if abandoned: + span.set_attribute("launchdarkly.stream.abandoned", True) + span.end() + + def parse_template(template: str, variables: dict[str, Any]) -> str: """ Replaces ``{{variable}}`` placeholders in *template* with values from diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py new file mode 100644 index 0000000..605a21d --- /dev/null +++ b/packages/client/tests/test_span_usage.py @@ -0,0 +1,382 @@ +"""Tests for the span usage and identity helpers. + +Covers TELEMETRY-CONTRACT.md sections 6 (span lifecycle), 8 (token accounting) and 9 (model +identity). + +The load-bearing tests here are the cache-direction ones. Folding cache tokens for a provider that +already includes them, or failing to fold for a provider that does not, produces a number that is +wrong by exactly the cached portion of every call, and nothing else catches it. +""" + +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server.utils import ( + SpanUsage, + add_cached_tokens_to_input, + create_run_usage, + end_span_once, + lang_chain_span_usage, + number_or_zero, + parse_usage, + set_model_identity_attributes, + set_usage_span_attributes, +) + + +class FakeSpan: + def __init__(self) -> None: + self.attributes: dict[str, Any] = {} + self.ended = 0 + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + def end(self) -> None: + self.ended += 1 + + +# ─── number_or_zero ────────────────────────────────────────────────────────── + + +class TestNumberOrZero: + def test_passes_a_number_through(self) -> None: + assert number_or_zero(42) == 42 + + def test_none_becomes_zero_rather_than_raising(self) -> None: + # The old bare int(...) raised here, taking the whole call down with it. + assert number_or_zero(None) == 0 + + def test_a_non_numeric_string_becomes_zero(self) -> None: + assert number_or_zero("abc") == 0 + + def test_a_numeric_string_is_read(self) -> None: + assert number_or_zero("17") == 17 + + def test_nan_becomes_zero(self) -> None: + # NaN is worse than 0: the metric guard tests `> 0`, which is false for NaN, so the metric + # would be dropped silently rather than reported low. + assert number_or_zero(float("nan")) == 0 + + def test_infinity_becomes_zero(self) -> None: + assert number_or_zero(float("inf")) == 0 + assert number_or_zero(float("-inf")) == 0 + + def test_a_bool_is_not_a_token_count(self) -> None: + assert number_or_zero(True) == 0 + + def test_a_float_truncates(self) -> None: + assert number_or_zero(3.7) == 3 + + +# ─── parse_usage ───────────────────────────────────────────────────────────── + + +class TestParseUsage: + def test_reads_the_snake_case_pair(self) -> None: + assert parse_usage({"input_tokens": 10, "output_tokens": 5}) == { + "input": 10, + "output": 5, + "total": 15, + } + + def test_reads_the_camel_case_pair(self) -> None: + assert parse_usage({"inputTokens": 10, "outputTokens": 5})["total"] == 15 + + def test_reads_the_bare_pair(self) -> None: + assert parse_usage({"input": 10, "output": 5})["total"] == 15 + + def test_first_matching_pair_wins(self) -> None: + result = parse_usage( + {"input_tokens": 1, "output_tokens": 2, "input": 100, "output": 200} + ) + assert result["total"] == 3 + + def test_never_trusts_a_provider_total(self) -> None: + result = parse_usage( + {"input_tokens": 3, "output_tokens": 7, "total_tokens": 999} + ) + assert result["total"] == 10 + + def test_an_unrecognised_bag_returns_zeros(self) -> None: + assert parse_usage({"foo": 1}) == {"input": 0, "output": 0, "total": 0} + + def test_a_none_field_does_not_raise(self) -> None: + assert parse_usage({"input_tokens": None, "output_tokens": 5})["input"] == 0 + + def test_folds_anthropic_cache_tokens_into_input(self) -> None: + # Anthropic reports cache beside input, so the real input is the sum of all three. + result = parse_usage( + { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + ) + assert result["input"] == 23554 + assert result["total"] == 23564 + + def test_reports_the_cache_breakdown_when_present(self) -> None: + result = parse_usage( + { + "input_tokens": 3, + "output_tokens": 1, + "cache_read_input_tokens": 10, + "cache_creation_input_tokens": 20, + } + ) + assert result["input_details"] == { + "uncached": 3, + "cache_read": 10, + "cache_creation": 20, + } + + def test_omits_the_breakdown_when_no_cache_field_is_present(self) -> None: + assert "input_details" not in parse_usage( + {"input_tokens": 1, "output_tokens": 2} + ) + + def test_omits_the_breakdown_for_the_bare_pair_which_has_no_cache_keys( + self, + ) -> None: + assert "input_details" not in parse_usage({"input": 1, "output": 2}) + + def test_accepts_the_bedrock_cache_write_alias(self) -> None: + # An unlisted spelling is silently dropped from the total, so every alias must be summed. + result = parse_usage( + {"inputTokens": 5, "outputTokens": 1, "cacheWriteInputTokens": 100} + ) + assert result["input"] == 105 + + def test_sums_both_camel_case_creation_spellings(self) -> None: + result = parse_usage( + { + "inputTokens": 0, + "outputTokens": 0, + "cacheCreationInputTokens": 10, + "cacheWriteInputTokens": 5, + } + ) + assert result["input"] == 15 + + +# ─── add_cached_tokens_to_input ────────────────────────────────────────────── + + +class TestAddCachedTokensToInput: + def test_adds_both_cache_buckets_on_top_of_input(self) -> None: + usage = add_cached_tokens_to_input( + { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + ) + assert usage.input == 23554 + assert usage.output == 10 + assert usage.cache_read == 19971 + assert usage.cache_creation == 3580 + + def test_an_absent_cache_field_is_zero_not_an_error(self) -> None: + usage = add_cached_tokens_to_input({"input_tokens": 5, "output_tokens": 2}) + assert (usage.input, usage.cache_read, usage.cache_creation) == (5, 0, 0) + + def test_an_absent_usage_bag_is_all_zeros(self) -> None: + usage = add_cached_tokens_to_input({}) + assert (usage.input, usage.output) == (0, 0) + + +# ─── lang_chain_span_usage ─────────────────────────────────────────────────── + + +class TestLangChainSpanUsage: + def test_does_not_add_cache_on_top_of_input(self) -> None: + # LangChain already counts cached tokens inside input_tokens. Adding them would double-count. + usage = lang_chain_span_usage( + { + "input_tokens": 100, + "output_tokens": 10, + "input_token_details": {"cache_read": 80, "cache_creation": 5}, + } + ) + assert usage is not None + assert usage.input == 100 + assert usage.cache_read == 80 + assert usage.cache_creation == 5 + + def test_an_empty_bag_is_none_not_zeros(self) -> None: + # None keeps a turn the provider said nothing about from registering as reported. + assert lang_chain_span_usage({}) is None + assert lang_chain_span_usage(None) is None + + def test_a_bag_with_no_token_keys_is_none(self) -> None: + assert lang_chain_span_usage({"input_token_details": {"cache_read": 5}}) is None + + def test_a_zero_count_is_still_reported(self) -> None: + usage = lang_chain_span_usage({"input_tokens": 0, "output_tokens": 0}) + assert usage is not None + assert usage.input == 0 + + def test_missing_details_default_to_zero(self) -> None: + usage = lang_chain_span_usage({"input_tokens": 5, "output_tokens": 1}) + assert usage is not None + assert usage.cache_read == 0 + + +# ─── The run accumulator ───────────────────────────────────────────────────── + + +class TestRunUsage: + def test_starts_unreported_and_at_zero(self) -> None: + run = create_run_usage() + assert run.reported is False + assert run.total.input == 0 + + def test_sums_across_turns(self) -> None: + run = create_run_usage() + run.add(SpanUsage(input=10, output=1, cache_read=2, cache_creation=3)) + run.add(SpanUsage(input=20, output=2, cache_read=4, cache_creation=5)) + assert run.total.input == 30 + assert run.total.output == 3 + assert run.total.cache_read == 6 + assert run.total.cache_creation == 8 + + def test_none_is_a_no_op_and_does_not_count_as_reported(self) -> None: + run = create_run_usage() + run.add(None) + assert run.reported is False + + def test_a_genuinely_empty_turn_still_counts_as_reported(self) -> None: + # This is the distinction the failure path needs: "reported zero" is not "reported nothing". + run = create_run_usage() + run.add(SpanUsage()) + assert run.reported is True + assert run.total.input == 0 + + def test_two_accumulators_do_not_share_state(self) -> None: + # A mutable default on the dataclass would make this fail. + first = create_run_usage() + first.add(SpanUsage(input=5)) + assert create_run_usage().total.input == 0 + + +# ─── set_usage_span_attributes ─────────────────────────────────────────────── + + +class TestSetUsageSpanAttributes: + def test_writes_all_seven_attributes(self) -> None: + span = FakeSpan() + set_usage_span_attributes( + span, SpanUsage(input=10, output=4, cache_read=6, cache_creation=1) + ) + assert span.attributes == { + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 4, + "gen_ai.usage.total_tokens": 14, + "gen_ai.usage.cache_read.input_tokens": 6, + "gen_ai.usage.cache_creation.input_tokens": 1, + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.completion_tokens": 4, + } + + def test_writes_zeros_rather_than_omitting_them(self) -> None: + # An absent attribute drops the span from every query that groups on usage, which reads as + # "no cached tokens" when it means "this handler forgot to say". + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage()) + assert len(span.attributes) == 7 + assert span.attributes["gen_ai.usage.cache_creation.input_tokens"] == 0 + + def test_total_is_derived_not_taken(self) -> None: + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage(input=7, output=3)) + assert span.attributes["gen_ai.usage.total_tokens"] == 10 + + def test_the_openllmetry_aliases_agree_with_the_canonical_keys(self) -> None: + # They disagreed once, because the alias was computed at a call site off Anthropic's + # cache-excluding input field. One writer, one number. + span = FakeSpan() + set_usage_span_attributes(span, SpanUsage(input=100, output=5, cache_read=80)) + assert ( + span.attributes["gen_ai.usage.prompt_tokens"] + == span.attributes["gen_ai.usage.input_tokens"] + ) + assert ( + span.attributes["gen_ai.usage.completion_tokens"] + == span.attributes["gen_ai.usage.output_tokens"] + ) + + +# ─── set_model_identity_attributes ─────────────────────────────────────────── + + +class TestSetModelIdentityAttributes: + def test_writes_both_provider_keys_with_the_same_value_by_default(self) -> None: + span = FakeSpan() + set_model_identity_attributes(span, "anthropic", "claude-3-5-sonnet") + assert span.attributes["gen_ai.system"] == "anthropic" + assert span.attributes["gen_ai.provider.name"] == "anthropic" + assert span.attributes["gen_ai.request.model"] == "claude-3-5-sonnet" + + def test_the_legacy_key_can_differ_from_the_current_one(self) -> None: + # The LangChain handlers keep the framework name on the old key, because the new key's + # semconv enum has no 'langchain' member. + span = FakeSpan() + set_model_identity_attributes( + span, "anthropic", "claude-3-5-sonnet", "langchain" + ) + assert span.attributes["gen_ai.system"] == "langchain" + assert span.attributes["gen_ai.provider.name"] == "anthropic" + + +# ─── end_span_once ─────────────────────────────────────────────────────────── + + +class TestEndSpanOnce: + def test_ends_the_span(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + assert span.ended == 1 + + def test_a_second_call_is_ignored(self) -> None: + # The streaming finally and the success path can both reach the same span. + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + end_span_once(span, tracker) + assert span.ended == 1 + + def test_marks_abandonment_without_asserting_failure(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker, abandoned=True) + assert span.attributes["launchdarkly.stream.abandoned"] is True + assert span.ended == 1 + + def test_does_not_mark_a_normal_end(self) -> None: + span: FakeSpan = FakeSpan() + tracker: set[int] = set() + end_span_once(span, tracker) + assert "launchdarkly.stream.abandoned" not in span.attributes + + def test_tracks_each_span_separately(self) -> None: + first, second = FakeSpan(), FakeSpan() + tracker: set[int] = set() + end_span_once(first, tracker) + end_span_once(second, tracker) + assert (first.ended, second.ended) == (1, 1) + + def test_an_unhashable_span_is_still_tracked(self) -> None: + # The tracker holds id(span), because an OTel span is not guaranteed hashable. + class Unhashable(FakeSpan): + __hash__ = None # type: ignore[assignment] + + span = Unhashable() + tracker: set[int] = set() + end_span_once(span, tracker) + end_span_once(span, tracker) + assert span.ended == 1 From 48b3f81bbdd0f80e5ffa80ac49bc7ebd01783b71 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:27:56 -0400 Subject: [PATCH 2/3] fix(client): carry the cache breakdown into the public UsageDict parse_usage started reporting a cache breakdown, but invoke and the judge runner both built UsageDict by hand from three keys, so input_details was always None on the blocking path while the streaming path handed back the nested dict. The two paths disagreed about the same run. Both now go through to_usage_dict, so the mapping has one author and cannot drift again. graph.py keeps building its own: a graph total is a sum across nodes and carries no per-call breakdown, and it says so. Found by Bugbot on #28. --- .../src/launchdarkly_ai_server/__init__.py | 2 + .../src/launchdarkly_ai_server/client.py | 8 +-- .../src/launchdarkly_ai_server/judges.py | 8 +-- .../src/launchdarkly_ai_server/utils.py | 27 +++++++++ packages/client/tests/test_span_usage.py | 58 +++++++++++++++++++ 5 files changed, 91 insertions(+), 12 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index feaeaec..039a521 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -76,6 +76,7 @@ set_openllmetry_prompt, set_usage_span_attributes, to_ld_context, + to_usage_dict, ) __all__ = [ # noqa: RUF022 @@ -121,6 +122,7 @@ "number_or_zero", "set_model_identity_attributes", "set_usage_span_attributes", + "to_usage_dict", "VariationMeta", # utils "create_handler", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 9e6ac4c..570ac0d 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -15,12 +15,12 @@ ProviderHandler, ProviderResponse, StreamEvent, - UsageDict, VariationMeta, ) from .utils import ( parse_json_with_possible_fences, select_handler, + to_usage_dict, ) @@ -112,11 +112,7 @@ async def invoke( else json.dumps(parsed_response) ) - usage_obj = UsageDict( - input=usage.get("input", 0), - output=usage.get("output", 0), - total=usage.get("total", 0), - ) + usage_obj = to_usage_dict(usage) if self._skip_judges: judge_tasks = await build_judge_tasks( diff --git a/packages/client/src/launchdarkly_ai_server/judges.py b/packages/client/src/launchdarkly_ai_server/judges.py index 9e644c2..898f864 100644 --- a/packages/client/src/launchdarkly_ai_server/judges.py +++ b/packages/client/src/launchdarkly_ai_server/judges.py @@ -13,7 +13,6 @@ NativeTool, ProviderHandler, TrackData, - UsageDict, ) from .utils import ( collapse_messages_to_instructions as _collapse_messages_to_instructions, @@ -22,6 +21,7 @@ normalize_mode, parse_json_with_possible_fences, to_ld_context, + to_usage_dict, ) @@ -409,11 +409,7 @@ def _matches(h: ProviderHandler) -> bool: reasoning = parsed.get("reasoning", "") raw_usage = result["usage"] - usage = UsageDict( - input=raw_usage.get("input", 0), - output=raw_usage.get("output", 0), - total=raw_usage.get("total", 0), - ) + usage = to_usage_dict(raw_usage) merged_track_data: TrackData = { **task.parent_track_data, diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 81e83fd..bf14b92 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -8,7 +8,9 @@ from .types import ( AiConfigRep, GraphNode, + InputTokenDetails, ProviderHandler, + UsageDict, VariationMeta, _HandlerFn, _StreamFn, @@ -142,6 +144,31 @@ def parse_usage(usage: dict[str, Any]) -> dict[str, Any]: return {"input": 0, "output": 0, "total": 0} +def to_usage_dict(usage: dict[str, Any]) -> UsageDict: + """Builds the public :class:`UsageDict` from a :func:`parse_usage` result. + + Shared because ``invoke`` and the judge runner both need it and both used to build the dataclass + by hand from three keys, which silently dropped the cache breakdown the moment ``parse_usage`` + started reporting one. A caller reading ``input_details`` off a blocking call got ``None`` while + the streaming path handed back the nested dict, so the two paths disagreed about the same run. + """ + details = usage.get("input_details") + return UsageDict( + input=usage.get("input", 0), + output=usage.get("output", 0), + total=usage.get("total", 0), + input_details=( + InputTokenDetails( + uncached=details.get("uncached", 0), + cache_read=details.get("cache_read", 0), + cache_creation=details.get("cache_creation", 0), + ) + if isinstance(details, dict) + else None + ), + ) + + @dataclass class SpanUsage: """The provider-neutral token counts a span reports, after the caller applied its cache rule. diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index 605a21d..dc3df9c 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -22,6 +22,7 @@ parse_usage, set_model_identity_attributes, set_usage_span_attributes, + to_usage_dict, ) @@ -380,3 +381,60 @@ class Unhashable(FakeSpan): end_span_once(span, tracker) end_span_once(span, tracker) assert span.ended == 1 + + +# ─── to_usage_dict ─────────────────────────────────────────────────────────── + + +class TestToUsageDict: + """The public UsageDict must carry everything parse_usage reported. + + Both call sites used to build the dataclass by hand from three keys, so the cache breakdown + vanished the moment parse_usage started reporting one: a caller reading input_details off a + blocking call got None while the streaming path handed back the nested dict. + """ + + def test_carries_the_three_totals(self) -> None: + usage = to_usage_dict({"input": 10, "output": 5, "total": 15}) + assert (usage.input, usage.output, usage.total) == (10, 5, 15) + + def test_carries_the_cache_breakdown_when_present(self) -> None: + usage = to_usage_dict( + { + "input": 23554, + "output": 10, + "total": 23564, + "input_details": { + "uncached": 3, + "cache_read": 19971, + "cache_creation": 3580, + }, + } + ) + assert usage.input_details is not None + assert usage.input_details.uncached == 3 + assert usage.input_details.cache_read == 19971 + assert usage.input_details.cache_creation == 3580 + + def test_is_none_when_the_provider_reported_no_cache(self) -> None: + assert ( + to_usage_dict({"input": 1, "output": 2, "total": 3}).input_details is None + ) + + def test_round_trips_a_parse_usage_result(self) -> None: + # The two functions are used together, so pin them together. + raw = { + "input_tokens": 3, + "output_tokens": 10, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + usage = to_usage_dict(parse_usage(raw)) + assert usage.input == 23554 + assert usage.input_details is not None + assert usage.input_details.cache_read == 19971 + + def test_a_missing_field_defaults_rather_than_raising(self) -> None: + usage = to_usage_dict({}) + assert (usage.input, usage.output, usage.total) == (0, 0, 0) + assert usage.input_details is None From 22ad56f1d34ef6abb6d7ad7f056dc670882fec64 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:29:08 -0400 Subject: [PATCH 3/3] fix(client): make end_span_once a no-op on a None span Handlers hold None for every span whenever the OpenTelemetry SDK is absent, and every other helper in this family already no-ops on it. This one did not, so a streaming cleanup path in a finally crashed with AttributeError on an install without the otel extra: the one place that should be hardest to break was the one place that was not guarded. Fixed in the shared helper rather than at six call sites, because all six handlers use it the same way and the next handler would hit the same edge. Found by Bugbot on #30. --- packages/client/src/launchdarkly_ai_server/utils.py | 6 ++++++ packages/client/tests/test_span_usage.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index bf14b92..da27fc1 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -369,7 +369,13 @@ def end_span_once(span: Any, tracker: set[int], abandoned: bool = False) -> None Ending the span is not always enough. Two handlers also hold a vendor generator or run that must be closed or cancelled in the same ``finally``. See TELEMETRY-CONTRACT.md section 6. + + A ``None`` span is a no-op, matching every other helper in this family. Handlers hold ``None`` + whenever the OpenTelemetry SDK is absent, and a cleanup path in a ``finally`` is the last place + that should have to remember it. """ + if span is None: + return key = id(span) if key in tracker: return diff --git a/packages/client/tests/test_span_usage.py b/packages/client/tests/test_span_usage.py index dc3df9c..bde037b 100644 --- a/packages/client/tests/test_span_usage.py +++ b/packages/client/tests/test_span_usage.py @@ -371,6 +371,12 @@ def test_tracks_each_span_separately(self) -> None: end_span_once(second, tracker) assert (first.ended, second.ended) == (1, 1) + def test_a_none_span_is_a_no_op(self) -> None: + # Handlers hold None whenever the OTel SDK is absent, and a cleanup path in a `finally` is + # the last place that should have to remember it. Every sibling helper already no-ops. + end_span_once(None, set()) + end_span_once(None, set(), abandoned=True) + def test_an_unhashable_span_is_still_tracked(self) -> None: # The tracker holds id(span), because an OTel span is not guaranteed hashable. class Unhashable(FakeSpan):