diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 5d1f2ee..039a521 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,16 +57,26 @@ ) 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, + to_usage_dict, ) __all__ = [ # noqa: RUF022 @@ -100,7 +111,18 @@ "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", + "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/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/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/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..da27fc1 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -2,12 +2,15 @@ import json import re +from dataclasses import dataclass, field from typing import Any, Literal from .types import ( AiConfigRep, GraphNode, + InputTokenDetails, ProviderHandler, + UsageDict, VariationMeta, _HandlerFn, _StreamFn, @@ -49,27 +52,339 @@ 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} +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. + + ``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. + + 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 + 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..bde037b --- /dev/null +++ b/packages/client/tests/test_span_usage.py @@ -0,0 +1,446 @@ +"""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, + to_usage_dict, +) + + +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_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): + __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 + + +# ─── 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