diff --git a/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py b/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py index 435ddfc899..266c9b740d 100644 --- a/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py +++ b/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py @@ -193,6 +193,14 @@ def _split_agent_messages( "gen_ai.usage.cache_creation.input_tokens", "ag.metrics.unit.tokens.cache_creation", ), + # Producer marker for how `gen_ai.usage.input_tokens` counts cached tokens. The + # OpenTelemetry contract says it includes them, which is what cost estimation + # assumes when this is absent; a producer whose count excludes them (the Agenta + # agent runner, which reports the cache buckets separately) sends false. + ( + "agenta.usage.input_tokens_includes_cache", + "ag.meta.usage.input_tokens_includes_cache", + ), # Producer-reported cost. This is the aggregate total for the span's whole run # (see the agent SDK's record_usage), so it maps to `cumulative`, not # `incremental`: as an incremental value the tree roll-up would add it on top of diff --git a/api/oss/src/core/tracing/utils/trees.py b/api/oss/src/core/tracing/utils/trees.py index 6f05ccd93e..415e940a27 100644 --- a/api/oss/src/core/tracing/utils/trees.py +++ b/api/oss/src/core/tracing/utils/trees.py @@ -1,7 +1,7 @@ from collections import OrderedDict -from typing import Dict, List, Optional +from typing import Any, Dict, List, NamedTuple, Optional -from litellm import cost_calculator +from litellm import cost_calculator, provider_list from oss.src.utils.logging import get_module_logger from oss.src.core.shared.dtos import Trace, Traces @@ -258,103 +258,77 @@ def _connect_tree_dfs( parent_span.spans = None -def _has_reported_cumulative(span: OTelFlatSpan) -> bool: +def _costs_bucket(span: OTelFlatSpan, bucket: str) -> Optional[dict]: if not isinstance(span.attributes, dict): - return False + return None - node = span.attributes - for key in ("ag", "metrics", "costs", "cumulative"): + node: Any = span.attributes + for key in ("ag", "metrics", "costs", bucket): if not isinstance(node, dict): - return False + return None node = node.get(key) - return isinstance(node, dict) and "total" in node + return node if isinstance(node, dict) else None -def cumulate_costs( - spans_id_tree: OrderedDict, - spans_idx: Dict[str, OTelFlatSpan], -) -> None: - def _get_incremental(span: OTelFlatSpan): - _costs = { - "prompt": 0.0, - "completion": 0.0, - "total": 0.0, - } +def _has_reported_cumulative(span: OTelFlatSpan) -> bool: + bucket = _costs_bucket(span, "cumulative") - if span.attributes is None: - return _costs + return bucket is not None and "total" in bucket - attr: dict = span.attributes - return { - "prompt": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("incremental", {}) - .get("prompt", 0.0) - ), - "completion": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("incremental", {}) - .get("completion", 0.0) - ), - "total": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("incremental", {}) - .get("total", 0.0) - ), - } +class _Costs(NamedTuple): + """Cost triple plus whether anything in this subtree actually measured a cost. - def _get_cumulative(span: OTelFlatSpan): - _costs = { - "prompt": 0.0, - "completion": 0.0, - "total": 0.0, - } + INVARIANT: `measured` is what decides whether the roll-up writes, never the + amounts. A measured 0.0 (a fully cached turn, a free model) is a fact and must + reach the ancestors; an absent measurement must leave them without the attribute + rather than claiming a zero nobody observed. + """ - if span.attributes is None: - return _costs + prompt: float = 0.0 + completion: float = 0.0 + total: float = 0.0 + measured: bool = False - attr: dict = span.attributes - return { - "prompt": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("cumulative", {}) - .get("prompt", 0.0) - ), - "completion": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("cumulative", {}) - .get("completion", 0.0) - ), - "total": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("costs", {}) - .get("cumulative", {}) - .get("total", 0.0) - ), - } +def _read_costs(span: OTelFlatSpan, bucket: str) -> _Costs: + values = _costs_bucket(span, bucket) - def _accumulate(a: dict, b: dict): - return { - "prompt": a.get("prompt", 0.0) + b.get("prompt", 0.0), - "completion": a.get("completion", 0.0) + b.get("completion", 0.0), - "total": a.get("total", 0.0) + b.get("total", 0.0), - } + if values is None: + return _Costs() - def _set_cumulative(span: OTelFlatSpan, costs: dict): + def _amount(key: str) -> float: + value = values.get(key, 0.0) + return float(value) if isinstance(value, (int, float)) else 0.0 + + return _Costs( + prompt=_amount("prompt"), + completion=_amount("completion"), + total=_amount("total"), + measured=True, + ) + + +def cumulate_costs( + spans_id_tree: OrderedDict, + spans_idx: Dict[str, OTelFlatSpan], +) -> None: + def _get_incremental(span: OTelFlatSpan) -> _Costs: + return _read_costs(span, "incremental") + + def _get_cumulative(span: OTelFlatSpan) -> _Costs: + return _read_costs(span, "cumulative") + + def _accumulate(a: _Costs, b: _Costs) -> _Costs: + return _Costs( + prompt=a.prompt + b.prompt, + completion=a.completion + b.completion, + total=a.total + b.total, + measured=a.measured or b.measured, + ) + + def _set_cumulative(span: OTelFlatSpan, costs: _Costs): if span.attributes is None: span.attributes = {} @@ -368,30 +342,32 @@ def _set_cumulative(span: OTelFlatSpan, costs: dict): if _has_reported_cumulative(span): return - if ( - costs.get("prompt", 0.0) != 0.0 - or costs.get("completion", 0.0) != 0.0 - or costs.get("total", 0.0) != 0.0 + if not costs.measured: + return + + if "ag" not in span.attributes or not isinstance( + span.attributes["ag"], + dict, ): - if "ag" not in span.attributes or not isinstance( - span.attributes["ag"], - dict, - ): - span.attributes["ag"] = {} + span.attributes["ag"] = {} - if "metrics" not in span.attributes["ag"] or not isinstance( - span.attributes["ag"]["metrics"], - dict, - ): - span.attributes["ag"]["metrics"] = {} + if "metrics" not in span.attributes["ag"] or not isinstance( + span.attributes["ag"]["metrics"], + dict, + ): + span.attributes["ag"]["metrics"] = {} - if "costs" not in span.attributes["ag"]["metrics"] or not isinstance( - span.attributes["ag"]["metrics"]["costs"], - dict, - ): - span.attributes["ag"]["metrics"]["costs"] = {} + if "costs" not in span.attributes["ag"]["metrics"] or not isinstance( + span.attributes["ag"]["metrics"]["costs"], + dict, + ): + span.attributes["ag"]["metrics"]["costs"] = {} - span.attributes["ag"]["metrics"]["costs"]["cumulative"] = costs + span.attributes["ag"]["metrics"]["costs"]["cumulative"] = { + "prompt": costs.prompt, + "completion": costs.completion, + "total": costs.total, + } _cumulate_tree_dfs( spans_id_tree, @@ -697,6 +673,132 @@ def _token_count(value) -> int: return 0 +# The incremental token buckets that carry a price. `total` is deliberately excluded: it +# prices nothing on its own, so a bucket holding only a total is not a measurement this +# path can turn into a cost. +PRICEABLE_TOKEN_BUCKETS = ("prompt", "completion", "cache_read", "cache_creation") + + +def _has_token_measurement(token_metrics: dict) -> bool: + """Whether the span reported a token count at all, as opposed to counting zero. + + Missing and zero are different facts. A fully cached turn that reports `prompt = 0` + measured zero and is priceable; a span carrying no token bucket measured nothing. + """ + return any(token_metrics.get(key) is not None for key in PRICEABLE_TOKEN_BUCKETS) + + +# The providers litellm can price by name. A provider identity outside this set is a +# customer's own connection slug, not a public catalog we can charge against. +KNOWN_PRICING_PROVIDERS = frozenset( + str(getattr(provider, "value", provider)).lower() for provider in provider_list +) + +# Agenta provider identities litellm has no name for, each mapped to the litellm provider +# whose published prices genuinely apply. An entry is a deliberate statement that this +# identity always denotes a public catalog, never a customer's own deployment; every +# identity absent from the map is treated as a custom connection and left unpriced. +TRUSTED_PRICING_PROVIDERS: Dict[str, str] = { + # Pi's and codex's id for OpenAI's ChatGPT/Codex subscription. The models are OpenAI's + # own (gpt-5.x, served from chatgpt.com/backend-api), so OpenAI's list prices apply. + "openai-codex": "openai", +} + + +def _dict(node: dict, key: str) -> dict: + value = node.get(key) + + return value if isinstance(value, dict) else {} + + +def _text(value) -> Optional[str]: + return value.strip() if isinstance(value, str) and value.strip() else None + + +def _input_tokens_include_cache(ag_meta: dict) -> bool: + """Whether the span's prompt bucket already counts its cached tokens. + + Default: True, the OpenTelemetry GenAI meaning of `gen_ai.usage.input_tokens`. + Producers whose input count *excludes* cache (the Agenta agent runner, which emits + the cache buckets separately) say so with `agenta.usage.input_tokens_includes_cache + = false`. The default is deliberately the inclusive one: runner and API ship as + independently versioned artifacts, so an old runner will post to a new API, and on + that skew the inclusive reading undercounts (the pre-existing bug) instead of + charging the cache twice. + """ + marker = _dict(ag_meta, "usage").get("input_tokens_includes_cache") + + if isinstance(marker, bool): + return marker + + if isinstance(marker, str): + normalized = marker.strip().lower() + if normalized in ("true", "1"): + return True + if normalized in ("false", "0"): + return False + + return True + + +def _pricing_prompt_tokens( + *, + input_tokens_include_cache: bool, + uncached_input_tokens: int, + cache_read_input_tokens: int, + cache_creation_input_tokens: int, +) -> int: + """The prompt count litellm expects: always inclusive of the cache buckets. + + litellm's generic calculator derives ordinary input by *subtracting* the cache + details from the prompt count it is given, so a producer whose input count already + includes them is passed through untouched and an exclusive one is summed up first. + """ + if input_tokens_include_cache: + return uncached_input_tokens + + return uncached_input_tokens + cache_read_input_tokens + cache_creation_input_tokens + + +def _is_priceable_provider(provider: str) -> bool: + """Whether a reported provider identity names a catalog litellm can price. + + Resolving through the trusted map before the membership test is what makes each + trusted entry load-bearing: a value that is not itself a litellm provider fails the + test, so a typo withholds the estimate instead of silently granting public prices. + """ + identity = provider.lower() + + return TRUSTED_PRICING_PROVIDERS.get(identity, identity) in KNOWN_PRICING_PROVIDERS + + +def _is_served_by_a_custom_connection(ag_meta: dict) -> bool: + """Whether the span was served by something whose prices we cannot know. + + A managed custom model is selected as `/` but the tracer + stamps only the bare model id, so a customer deployment named after a public model + would otherwise be charged that public model's price. The connection identity + survives as the provider attribute (the runner puts the slug in `gen_ai.system`), + and third-party instrumentation additionally reports a base URL or endpoint for a + self-hosted gateway. Either signal means no priceable identity exists. + + A custom endpoint outranks a trusted provider identity: a gateway in front of a + public catalog charges its own prices, so the identity no longer implies the + catalog's. + """ + request = _dict(ag_meta, "request") + if _text(request.get("base_url")) or _text(request.get("endpoint")): + return True + + provider = ( + _text(_dict(ag_meta, "provider").get("name")) + or _text(ag_meta.get("provider")) + or _text(ag_meta.get("system")) + ) + + return provider is not None and not _is_priceable_provider(provider) + + def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): for span in span_idx.values(): if ( @@ -709,19 +811,33 @@ def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): ag_meta: dict = ag_attr.get("meta", {}) ag_data: dict = ag_attr.get("data", {}) - # The agent runner sets the response model only for codex; every other - # harness sets only the request model, so without that fallback those - # spans are never priced at all. + # Every model name on the span is a bare id, so a custom connection + # disqualifies all of them equally: the Pi tracer stamps a response model on + # every assistant message, including one a customer's own connection served, + # so guarding only the request model would price that connection at public + # rates through the response model. A name we cannot attribute to a catalog + # is priced confidently and wrongly, which is worse than no estimate. model = ( - ag_meta.get("response", {}).get("model") - or ag_meta.get("request", {}).get("model") - or ag_data.get("parameters", {}).get("model") + None + if _is_served_by_a_custom_connection(ag_meta) + else ( + ag_meta.get("response", {}).get("model") + or ag_meta.get("request", {}).get("model") + or ag_data.get("parameters", {}).get("model") + ) ) token_metrics: dict = ( ag_attr.get("metrics", {}).get("tokens", {}).get("incremental", {}) ) + # A span that reported no token count measured nothing. Estimating it anyway + # writes a zero-cost dictionary, and the presence of that dictionary is the + # `measured` signal the roll-up propagates, so every ancestor would claim the + # run was measured and free. + if not _has_token_measurement(token_metrics): + continue + uncached_input_tokens = _token_count(token_metrics.get("prompt")) cache_read_input_tokens = _token_count(token_metrics.get("cache_read")) cache_creation_input_tokens = _token_count( @@ -729,16 +845,11 @@ def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): ) completion_tokens = _token_count(token_metrics.get("completion")) - # INVARIANT: the `prompt` bucket carries *exclusive* input (gen_ai.usage. - # input_tokens is raw uncached input), while litellm's generic calculator - # derives ordinary input by subtracting the cache details from the prompt - # count it is given, so it expects an *inclusive* count. If the canonical - # `prompt` bucket ever becomes inclusive at ingest, this sum double counts - # and must be dropped; it is kept in one place so that is a one-line change. - inclusive_prompt_tokens = ( - uncached_input_tokens - + cache_read_input_tokens - + cache_creation_input_tokens + inclusive_prompt_tokens = _pricing_prompt_tokens( + input_tokens_include_cache=_input_tokens_include_cache(ag_meta), + uncached_input_tokens=uncached_input_tokens, + cache_read_input_tokens=cache_read_input_tokens, + cache_creation_input_tokens=cache_creation_input_tokens, ) try: diff --git a/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py b/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py index 7dcc55c662..382c2d8edd 100644 --- a/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py +++ b/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py @@ -708,6 +708,18 @@ def test_cache_read_tokens_mapped(self, adapter): assert features.metrics["unit.tokens.cache_read"] == 2304 + @pytest.mark.parametrize("includes_cache", [True, False]) + def test_input_tokens_cache_marker_mapped(self, adapter, includes_cache): + # Tells cost estimation whether `gen_ai.usage.input_tokens` already counts the + # cache buckets; absent, OpenTelemetry's meaning (it does) is assumed. + bag = _make_bag( + {"agenta.usage.input_tokens_includes_cache": includes_cache}, + ) + features = SpanFeatures() + adapter.process(bag, features) + + assert features.meta["usage.input_tokens_includes_cache"] is includes_cache + def test_reported_cost_mapped_to_cumulative_not_incremental(self, adapter): # gen_ai.usage.cost is the aggregate for the whole run, so it lands on # `cumulative`; as an `incremental` value the tree roll-up would add it on top diff --git a/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py b/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py index 549967ab67..294b3063e1 100644 --- a/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py +++ b/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py @@ -150,7 +150,7 @@ def test_reported_cost_becomes_cumulative_total_after_ingest(): def test_reported_cost_survives_rollup_when_no_model_span_is_priceable( fixed_token_pricing, ): - # The roll-up writes `cumulative` only when it computes a non-zero total, so an + # The roll-up writes `cumulative` only for a subtree that measured something, so an # agent run whose model spans price to nothing must keep the reported figure. span_idx = _ingest( [ @@ -185,29 +185,36 @@ def test_reported_cost_wins_over_recomputed_child_costs(fixed_token_pricing): assert _cumulative_costs(model_span)["total"] == pytest.approx(0.00021) -def test_reported_cost_propagates_to_an_ancestor_that_reports_nothing( +def test_reported_cost_is_carried_by_the_agent_span_across_ingest_batches( fixed_token_pricing, ): - # The SDK's workflow root does not always report a cost; it must still show the - # agent subtree's, which is what the trace list and playground read. - root = _otel_span( - span_id=ROOT_SPAN_ID, - span_name="workflow", - attributes={"ag.type.node": "workflow"}, + """The real topology: the SDK and the runner ship this trace in two OTLP requests. + + Roll-up is per request (`TracingService.ingest_span_dtos` calls it on the batch it + was handed), so no batch ever holds both the workflow root and the runner subtree. + What this pins is therefore the opposite of a single-batch test: the reported cost + settles on the agent span, which is the *root of its own batch*, and the workflow + root gets no cost at ingest. Anything that needs a whole-trace total has to + aggregate stored spans; it cannot expect the roll-up to have crossed the requests. + """ + sdk_batch = _ingest( + [ + _otel_span( + span_id=ROOT_SPAN_ID, + span_name="workflow", + attributes={"ag.type.node": "workflow"}, + ) + ] ) - span_idx = _ingest( + runner_batch = _ingest( [ - root, _agent_span(parent_id=ROOT_SPAN_ID, start_offset_s=1), _unpriceable_model_span(parent_id=AGENT_SPAN_ID, start_offset_s=2), ] ) - root_span = span_idx["workflow"] - agent_span = span_idx["invoke_agent"] - - assert _cumulative_costs(root_span)["total"] == 0.4237 - assert _cumulative_costs(agent_span)["total"] == 0.4237 + assert _cumulative_costs(sdk_batch["workflow"]) == {} + assert _cumulative_costs(runner_batch["invoke_agent"])["total"] == 0.4237 def test_reported_cost_is_not_double_counted_when_parent_and_child_both_report(): diff --git a/api/oss/tests/pytest/unit/tracing/test_cost_calculation.py b/api/oss/tests/pytest/unit/tracing/test_cost_calculation.py index 588deba672..4524282676 100644 --- a/api/oss/tests/pytest/unit/tracing/test_cost_calculation.py +++ b/api/oss/tests/pytest/unit/tracing/test_cost_calculation.py @@ -1,8 +1,16 @@ +from datetime import datetime, timedelta, timezone from typing import Optional import pytest from litellm import cost_calculator +from oss.src.apis.fastapi.otlp.utils.processing import parse_from_otel_span_dto +from oss.src.core.otel.dtos import ( + OTelContextDTO, + OTelSpanDTO, + OTelSpanKind, + OTelStatusCode, +) from oss.src.core.tracing.dtos import OTelFlatSpan, SpanType from oss.src.core.tracing.utils.trees import ( KNOWN_PRICING_PROVIDERS, @@ -15,6 +23,28 @@ OPENAI_MODEL = "gpt-5.3-codex" ANTHROPIC_MODEL = "claude-sonnet-4-6" +# One real Claude turn: 13,463 of the 13,556 input tokens were served from cache. +# An OpenTelemetry producer reports input_tokens=13,556 (cache included); the agent +# runner reports prompt=93 (cache excluded) plus the cache buckets. Both describe the +# same turn and must price the same. +INCLUSIVE_INPUT = 13_556 +CACHE_READ = 13_463 +EXCLUSIVE_INPUT = INCLUSIVE_INPUT - CACHE_READ + +# Literal prices from the pinned litellm (1.92.0). They are hard-coded on purpose: an +# oracle that recomputes them through the same litellm call pins the wiring but not the +# prices, so it cannot tell a tenfold overcharge from a correct estimate. +CACHED_TURN_PROMPT_COST = { + OPENAI_MODEL: 0.002518775, + ANTHROPIC_MODEL: 0.0043179, +} +# What the same turn costs if the cache buckets are added on top of an already-inclusive +# input count, i.e. the double count this pricing path must never produce. +DOUBLE_COUNTED_PROMPT_COST = { + OPENAI_MODEL: 0.026079025, + ANTHROPIC_MODEL: 0.0447069, +} + def _span( *, @@ -24,6 +54,11 @@ def _span( response_model: Optional[str] = None, request_model: Optional[str] = None, parameters_model: Optional[str] = None, + provider: Optional[str] = None, + system: Optional[str] = None, + base_url: Optional[str] = None, + endpoint: Optional[str] = None, + input_tokens_includes_cache: Optional[bool] = None, prompt: Optional[int] = None, completion: Optional[int] = None, cache_read: Optional[int] = None, @@ -33,8 +68,23 @@ def _span( meta: dict = {} if response_model is not None: meta["response"] = {"model": response_model} + + request: dict = {} if request_model is not None: - meta["request"] = {"model": request_model} + request["model"] = request_model + if base_url is not None: + request["base_url"] = base_url + if endpoint is not None: + request["endpoint"] = endpoint + if request: + meta["request"] = request + + if provider is not None: + meta["provider"] = {"name": provider} + if system is not None: + meta["system"] = system + if input_tokens_includes_cache is not None: + meta["usage"] = {"input_tokens_includes_cache": input_tokens_includes_cache} data: dict = {} if parameters_model is not None: @@ -93,16 +143,149 @@ def _cumulative_costs(span: OTelFlatSpan) -> dict: ) -def _expected(model: str, *, prompt: int, completion: int, read: int, creation: int): - """Oracle: what litellm charges for the inclusive-prompt shape we mean to send.""" - prompt_cost, completion_cost = cost_calculator.cost_per_token( - model=model, - prompt_tokens=prompt + read + creation, - completion_tokens=completion, - cache_read_input_tokens=read, - cache_creation_input_tokens=creation, +@pytest.fixture +def litellm_calls(monkeypatch): + """Capture the exact arguments handed to litellm, and price through the real one.""" + calls: list[dict] = [] + real = cost_calculator.cost_per_token + + def _cost_per_token(**kwargs): + calls.append(kwargs) + return real(**kwargs) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _cost_per_token, ) - return prompt_cost, completion_cost + + return calls + + +# ── the producer contract for cached input ────────────────────────────────── + + +@pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) +def test_an_inclusive_producer_is_priced_at_its_input_count(model): + """OTel says gen_ai.usage.input_tokens already counts cached tokens.""" + span = _span( + response_model=model, + input_tokens_includes_cache=True, + prompt=INCLUSIVE_INPUT, + completion=0, + cache_read=CACHE_READ, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span)["prompt"] == pytest.approx( + CACHED_TURN_PROMPT_COST[model] + ) + + +@pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) +def test_an_unmarked_producer_is_treated_as_inclusive(model): + """The default must be OTel's meaning: third-party instrumentation sends no marker.""" + span = _span( + response_model=model, + prompt=INCLUSIVE_INPUT, + completion=0, + cache_read=CACHE_READ, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span)["prompt"] == pytest.approx( + CACHED_TURN_PROMPT_COST[model] + ) + assert _incremental_costs(span)["prompt"] != pytest.approx( + DOUBLE_COUNTED_PROMPT_COST[model] + ) + + +@pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) +def test_an_exclusive_producer_sums_its_cache_buckets_in(model): + """The runner's shape: prompt is raw uncached input, cache reported separately.""" + span = _span( + response_model=model, + input_tokens_includes_cache=False, + prompt=EXCLUSIVE_INPUT, + completion=0, + cache_read=CACHE_READ, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span)["prompt"] == pytest.approx( + CACHED_TURN_PROMPT_COST[model] + ) + + +@pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) +def test_both_producer_contracts_price_the_same_turn_identically(model): + inclusive = _span( + span_id="inclusive", + response_model=model, + input_tokens_includes_cache=True, + prompt=INCLUSIVE_INPUT, + completion=0, + cache_read=CACHE_READ, + ) + exclusive = _span( + span_id="exclusive", + response_model=model, + input_tokens_includes_cache=False, + prompt=EXCLUSIVE_INPUT, + completion=0, + cache_read=CACHE_READ, + ) + + calculate_costs({s.span_id: s for s in (inclusive, exclusive)}) + + assert _incremental_costs(inclusive)["total"] == pytest.approx( + _incremental_costs(exclusive)["total"] + ) + + +def test_litellm_receives_the_prompt_and_cache_arguments_it_expects(litellm_calls): + """Pin the tuple: litellm subtracts the cache details from the prompt count.""" + inclusive = _span( + span_id="inclusive", + response_model=OPENAI_MODEL, + input_tokens_includes_cache=True, + prompt=INCLUSIVE_INPUT, + completion=7, + cache_read=CACHE_READ, + cache_creation=4_096, + ) + exclusive = _span( + span_id="exclusive", + response_model=OPENAI_MODEL, + input_tokens_includes_cache=False, + prompt=EXCLUSIVE_INPUT, + completion=7, + cache_read=CACHE_READ, + cache_creation=4_096, + ) + + calculate_costs({"a": inclusive}) + calculate_costs({"b": exclusive}) + + assert litellm_calls == [ + { + "model": OPENAI_MODEL, + "prompt_tokens": INCLUSIVE_INPUT, + "completion_tokens": 7, + "cache_read_input_tokens": CACHE_READ, + "cache_creation_input_tokens": 4_096, + }, + { + "model": OPENAI_MODEL, + "prompt_tokens": EXCLUSIVE_INPUT + CACHE_READ + 4_096, + "completion_tokens": 7, + "cache_read_input_tokens": CACHE_READ, + "cache_creation_input_tokens": 4_096, + }, + ] @pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) @@ -119,6 +302,7 @@ def _expected(model: str, *, prompt: int, completion: int, read: int, creation: def test_calculate_costs_prices_cache_buckets(model, read, creation): span = _span( response_model=model, + input_tokens_includes_cache=False, prompt=1, completion=20, cache_read=read or None, @@ -127,8 +311,12 @@ def test_calculate_costs_prices_cache_buckets(model, read, creation): calculate_costs({span.span_id: span}) - prompt_cost, completion_cost = _expected( - model, prompt=1, completion=20, read=read, creation=creation + prompt_cost, completion_cost = cost_calculator.cost_per_token( + model=model, + prompt_tokens=1 + read + creation, + completion_tokens=20, + cache_read_input_tokens=read, + cache_creation_input_tokens=creation, ) costs = _incremental_costs(span) @@ -140,7 +328,13 @@ def test_calculate_costs_prices_cache_buckets(model, read, creation): @pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) def test_cached_input_is_not_dropped_from_the_estimate(model): """Regression for #5540: cache-read tokens used to be priced as if absent.""" - span = _span(response_model=model, prompt=1, completion=20, cache_read=25_182) + span = _span( + response_model=model, + input_tokens_includes_cache=False, + prompt=1, + completion=20, + cache_read=25_182, + ) calculate_costs({span.span_id: span}) @@ -153,9 +347,19 @@ def test_cached_input_is_not_dropped_from_the_estimate(model): @pytest.mark.parametrize("model", [OPENAI_MODEL, ANTHROPIC_MODEL]) def test_ordinary_input_survives_alongside_cache_reads(model): - """The prompt bucket is exclusive, so it must be added to, not replaced by, cache.""" - with_ordinary = _span(response_model=model, prompt=93, cache_read=13_463) - without_ordinary = _span(response_model=model, prompt=0, cache_read=13_463) + """An exclusive prompt bucket is added to, not replaced by, the cache buckets.""" + with_ordinary = _span( + response_model=model, + input_tokens_includes_cache=False, + prompt=EXCLUSIVE_INPUT, + cache_read=CACHE_READ, + ) + without_ordinary = _span( + response_model=model, + input_tokens_includes_cache=False, + prompt=0, + cache_read=CACHE_READ, + ) calculate_costs({"a": with_ordinary, "b": without_ordinary}) @@ -165,18 +369,16 @@ def test_ordinary_input_survives_alongside_cache_reads(model): ) +# ── model lookup ──────────────────────────────────────────────────────────── + + def test_request_model_is_used_when_response_model_is_absent(): span = _span(request_model=OPENAI_MODEL, prompt=1_000, completion=100) calculate_costs({span.span_id: span}) - prompt_cost, completion_cost = _expected( - OPENAI_MODEL, prompt=1_000, completion=100, read=0, creation=0 - ) - - assert _incremental_costs(span)["total"] == pytest.approx( - prompt_cost + completion_cost - ) + # litellm 1.92.0: 1,000 prompt + 100 completion tokens of gpt-5.3-codex. + assert _incremental_costs(span)["total"] == pytest.approx(0.00175 + 0.0014) def test_response_model_wins_over_request_model(): @@ -189,13 +391,7 @@ def test_response_model_wins_over_request_model(): calculate_costs({span.span_id: span}) - prompt_cost, completion_cost = _expected( - OPENAI_MODEL, prompt=1_000, completion=100, read=0, creation=0 - ) - - assert _incremental_costs(span)["total"] == pytest.approx( - prompt_cost + completion_cost - ) + assert _incremental_costs(span)["total"] == pytest.approx(0.00175 + 0.0014) def test_legacy_parameters_model_still_works(): @@ -232,6 +428,191 @@ def test_non_llm_span_types_are_skipped(): assert _incremental_costs(span) == {} +# ── custom connections must not borrow a public price ─────────────────────── + + +@pytest.mark.parametrize( + "custom", + [ + {"system": "acme-gateway"}, + {"provider": "acme-gateway"}, + {"base_url": "https://llm.acme.internal/v1"}, + {"endpoint": "https://llm.acme.internal/v1/chat/completions"}, + ], + ids=["slug-in-system", "slug-in-provider", "base-url", "endpoint"], +) +def test_a_priceable_request_model_on_a_custom_connection_is_not_priced(custom): + """A customer deployment named `gpt-5.3-codex` is not OpenAI's `gpt-5.3-codex`. + + The connection slug never reaches `gen_ai.request.model` (the tracer stamps the bare + model id), so the only evidence of who served the call is the provider identity or a + custom endpoint. Without a catalog we can attribute the name to, no estimate beats a + confident wrong one. + """ + span = _span(request_model=OPENAI_MODEL, prompt=1_000, completion=100, **custom) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +@pytest.mark.parametrize("provider", ["openai", "anthropic", "OpenAI"]) +def test_a_request_model_on_a_known_provider_is_still_priced(provider): + span = _span( + request_model=OPENAI_MODEL, system=provider, prompt=1_000, completion=100 + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span)["total"] == pytest.approx(0.00175 + 0.0014) + + +@pytest.mark.parametrize( + "custom", + [ + {"system": "acme-gateway"}, + {"provider": "acme-gateway"}, + {"base_url": "https://llm.acme.internal/v1"}, + {"endpoint": "https://llm.acme.internal/v1/chat/completions"}, + ], + ids=["slug-in-system", "slug-in-provider", "base-url", "endpoint"], +) +def test_a_response_model_on_a_custom_connection_is_not_priced_either(custom): + """The guard covers the model the lookup actually prefers. + + The Pi tracer stamps `gen_ai.response.model` on every assistant message, custom + connections included, so a guard that only withheld the request model would price a + customer's own deployment at public rates through the response model instead. + """ + span = _span( + response_model=OPENAI_MODEL, + request_model=OPENAI_MODEL, + prompt=1_000, + completion=100, + **custom, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +def test_a_legacy_parameters_model_on_a_custom_connection_is_not_priced(): + span = _span( + parameters_model=OPENAI_MODEL, + system="acme-gateway", + prompt=1_000, + completion=100, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +# ── trusted provider identities ───────────────────────────────────────────── + + +@pytest.mark.parametrize("provider_field", ["system", "provider"]) +def test_a_trusted_provider_identity_is_priced(provider_field): + """`openai-codex` is not a litellm provider; the map says it means OpenAI's catalog. + + Codex needs an explicit pricing identity. The mere presence of a response model is + not that identity — every Pi assistant message has one, custom connections included. + """ + span = _span( + response_model=OPENAI_MODEL, + request_model=OPENAI_MODEL, + prompt=1_000, + completion=100, + **{provider_field: "openai-codex"}, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span)["total"] == pytest.approx(0.00175 + 0.0014) + + +def test_a_trusted_provider_behind_a_custom_endpoint_is_not_priced(): + """A gateway charges its own prices, so the endpoint outranks the trusted identity.""" + span = _span( + response_model=OPENAI_MODEL, + system="openai-codex", + base_url="https://llm.acme.internal/v1", + prompt=1_000, + completion=100, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +def test_every_trusted_provider_maps_to_a_litellm_provider(): + """The map's values are load-bearing: a typo must withhold prices, not grant them.""" + assert TRUSTED_PRICING_PROVIDERS + assert all( + litellm_provider in KNOWN_PRICING_PROVIDERS + for litellm_provider in TRUSTED_PRICING_PROVIDERS.values() + ) + assert all(key == key.lower() for key in TRUSTED_PRICING_PROVIDERS) + + +def test_an_unlisted_agent_provider_is_not_trusted(): + span = _span( + response_model=OPENAI_MODEL, + system="openai-codex-lookalike", + prompt=1_000, + completion=100, + ) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +# ── a missing measurement is not a measured zero ──────────────────────────── + + +def test_a_model_span_with_no_token_metrics_is_not_priced(): + """No token bucket means nothing was measured; a zero-cost estimate would lie.""" + span = _span(response_model=OPENAI_MODEL) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +@pytest.mark.parametrize( + "bucket", + ["prompt", "completion", "cache_read", "cache_creation"], +) +def test_a_single_zero_token_bucket_still_counts_as_measured(bucket): + """Zero is a fact a producer reported; only an absent bucket is missing.""" + span = _span(response_model=OPENAI_MODEL, **{bucket: 0}) + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == { + "prompt": 0.0, + "completion": 0.0, + "total": 0.0, + } + + +def test_a_token_bucket_holding_only_a_total_is_not_priceable(): + """A total prices nothing on its own, so it must not manufacture a measured zero.""" + span = _span(response_model=OPENAI_MODEL) + span.attributes["ag"]["metrics"] = {"tokens": {"incremental": {"total": 1_100}}} + + calculate_costs({span.span_id: span}) + + assert _incremental_costs(span) == {} + + +# ── roll-up ───────────────────────────────────────────────────────────────── + + def test_reported_zero_cost_is_not_overwritten_by_an_estimate(): """Missing and measured-zero are different facts: a free/fully-cached turn is 0.""" parent = _span( @@ -253,6 +634,45 @@ def test_reported_zero_cost_is_not_overwritten_by_an_estimate(): assert _cumulative_costs(spans["child"])["total"] > 0.0 +def test_reported_zero_cost_propagates_to_a_parent_that_reports_nothing(): + """A measured zero is a measurement; the ancestor must show 0, not nothing.""" + parent = _span(span_id="parent", span_type=SpanType.AGENT) + child = _span( + span_id="child", + parent_id="parent", + span_type=SpanType.AGENT, + reported_cost=0.0, + ) + + spans = {s.span_id: s for s in calculate_and_propagate_metrics([parent, child])} + + assert _cumulative_costs(spans["parent"]) == { + "prompt": 0.0, + "completion": 0.0, + "total": 0.0, + } + + +def test_a_subtree_that_measured_nothing_leaves_the_parent_without_an_attribute(): + """The other half of the rule: never turn a missing measurement into a zero. + + The child is a priced span type carrying a priceable model, so estimation really runs + on it. A tool child would have been skipped before estimation and left this blind. + """ + parent = _span(span_id="parent", span_type=SpanType.AGENT) + child = _span( + span_id="child", + parent_id="parent", + span_type=SpanType.CHAT, + response_model=OPENAI_MODEL, + ) + + spans = {s.span_id: s for s in calculate_and_propagate_metrics([parent, child])} + + assert _cumulative_costs(spans["child"]) == {} + assert _cumulative_costs(spans["parent"]) == {} + + def test_reported_nonzero_cost_is_not_replaced_by_an_estimate(): parent = _span( span_id="parent", @@ -289,3 +709,58 @@ def test_missing_cost_is_filled_by_the_rollup(): == _cumulative_costs(spans["child"])["total"] > 0.0 ) + + +# ── the marker on the wire ────────────────────────────────────────────────── + +_TRACE_ID = "a" * 32 +_SPAN_ID = "d" * 16 +_START = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _ingest_one(attributes: dict) -> OTelFlatSpan: + otel_span = OTelSpanDTO( + context=OTelContextDTO(trace_id=f"0x{_TRACE_ID}", span_id=f"0x{_SPAN_ID}"), + parent=None, + name="chat", + kind=OTelSpanKind.SPAN_KIND_INTERNAL, + start_time=_START, + end_time=_START + timedelta(seconds=1), + status_code=OTelStatusCode.STATUS_CODE_OK, + attributes=attributes, + ) + + (span,) = calculate_and_propagate_metrics([parse_from_otel_span_dto(otel_span)]) + + return span + + +@pytest.mark.parametrize( + "marker,expected", + [ + (None, CACHED_TURN_PROMPT_COST[ANTHROPIC_MODEL]), + (True, CACHED_TURN_PROMPT_COST[ANTHROPIC_MODEL]), + (False, DOUBLE_COUNTED_PROMPT_COST[ANTHROPIC_MODEL]), + ], + ids=["absent-defaults-to-inclusive", "inclusive", "exclusive"], +) +def test_the_wire_marker_selects_the_pricing_contract(marker, expected): + """Same wire numbers, two producer contracts, two legitimate prices. + + With `input_tokens = 13,556` the inclusive reading is the whole turn and the + exclusive reading means 13,556 fresh tokens *plus* 13,463 cached ones — a different, + larger turn. The marker is what tells them apart; absent, OTel's meaning wins. + """ + attributes = { + "gen_ai.operation.name": "chat", + "gen_ai.response.model": ANTHROPIC_MODEL, + "gen_ai.usage.input_tokens": INCLUSIVE_INPUT, + "gen_ai.usage.output_tokens": 0, + "gen_ai.usage.cache_read.input_tokens": CACHE_READ, + } + if marker is not None: + attributes["agenta.usage.input_tokens_includes_cache"] = marker + + span = _ingest_one(attributes) + + assert _incremental_costs(span)["prompt"] == pytest.approx(expected)