diff --git a/agent_core/loop_types.py b/agent_core/loop_types.py index e9a5fee..db1f64e 100644 --- a/agent_core/loop_types.py +++ b/agent_core/loop_types.py @@ -8,7 +8,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, field from functools import partial -from typing import Any, Literal, Protocol, cast, runtime_checkable +from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable from agent_core.messages import Message @@ -18,6 +18,60 @@ WALL_DEADLINE_MONOTONIC_KEY = "wall_deadline_monotonic" +class UsageMetadataExtras(TypedDict, total=False): + """Host-specific usage aliases and provenance. + + ``extract_usage`` never emits these; they exist so a host that stamps + its own aliases onto a usage mapping can still describe the result as + a :class:`UsageMetadata`. ApodexHarness' budget observer reads the + ``input_tokens`` / ``output_tokens`` aliases, and its native clients + carry ``total_tokens`` on ``LLMResponse.usage``. + + Optional keys are NOT directly indexable under a type checker + (``reportTypedDictNotRequiredAccess``) — read them with ``.get(key, 0)``. + """ + + input_tokens: int + output_tokens: int + total_tokens: int + estimated: bool + + +class UsageMetadata(UsageMetadataExtras): + """What :func:`agent_core.runtime.loop.extract_usage` returns. + + The normalized fields below are required: every non-``None`` return + from ``extract_usage`` carries all of them, zero-filled when the + provider reported nothing, so a consumer can index them without + probing which response shape it was handed. + ``tests/test_llm_runtime_extract_usage.py`` pins that at runtime + against ``__required_keys__``. + + This is a *producer* type, deliberately narrower than the + ``usage`` fields on :class:`TurnContext` / :class:`LLMAttemptContext`. + Those stay ``Mapping[str, Any] | None`` because hosts build usage + mappings of their own — partial dicts in test doubles, alias-only + shapes, ``dict(event["usage"])`` re-wraps out of a ``dict[str, Any]`` + attempt event — and neither ``dict[str, int]`` nor ``dict[str, Any]`` + is assignable to a TypedDict. A consumer that wants the precise shape + annotates its own parameter as ``UsageMetadata``; the loop does not + force it on producers. Extra keys beyond those declared here are + likewise a host's business: a TypedDict cannot express "open" on + Python 3.12 (PEP 728 lands later), which is the other reason the + dataclass fields stay a plain mapping. + """ + + provider: str + model: str + prompt_tokens: int + completion_tokens: int + cache_read_tokens: int + cache_write_tokens: int + cached_tokens: int + cache_creation_tokens: int + reasoning_tokens: int + + def deadline_remaining_s(metadata: Mapping[str, Any] | None) -> float | None: """Return seconds to a structural lease or absolute soft deadline. @@ -135,7 +189,11 @@ class TurnContext: thinking: str tool_calls: list[dict[str, Any]] messages: list[Message] - usage: dict[str, int] | None + # Read-only mapping, not ``UsageMetadata``, on purpose: producers are + # hosts, and a TypedDict rejects the ``dict[str, int]`` / ``dict[str, Any]`` + # shapes they build. ``UsageMetadata`` documents what ``extract_usage`` + # puts here; annotate a consumer's own parameter with it for precision. + usage: Mapping[str, Any] | None metadata: dict[str, Any] # Reasoning recovered from tags leaked into visible content. leaked_reasoning: str = "" @@ -209,7 +267,8 @@ class LLMAttemptContext: recovery_action: str = "" duration_ms: int = 0 ttft_ms: int | None = None - usage: dict[str, int] | None = None + # See ``TurnContext.usage`` for why this is a plain mapping. + usage: Mapping[str, Any] | None = None finish_reason: str = "" visible_chars: int = 0 reasoning_chars: int = 0 @@ -631,6 +690,8 @@ async def notify_tool_result( "ToolCallIntervention", "ToolResult", "TurnContext", + "UsageMetadata", + "UsageMetadataExtras", "deadline_remaining_s", "drain_background_observers", "merge_interventions", diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py index 133dc48..826c210 100644 --- a/agent_core/runtime/loop/_response.py +++ b/agent_core/runtime/loop/_response.py @@ -6,6 +6,7 @@ from typing import Any from agent_core.llm import LLMResponse +from agent_core.loop_types import UsageMetadata from agent_core.messages import Message logger = logging.getLogger(__name__) @@ -118,7 +119,7 @@ def _pick_int(*candidates: Any) -> int: return 0 -def extract_usage(response: Any) -> dict[str, int | str] | None: +def extract_usage(response: Any) -> UsageMetadata | None: """Extract token usage from an LLM response, normalized to OpenAI shape. Returns a dict with keys ``provider`` / ``model`` / ``prompt_tokens`` / @@ -186,7 +187,7 @@ def extract_usage(response: Any) -> dict[str, int | str] | None: # the pre-split cache fields. cache_read = int(usage.get("cached_tokens", 0) or 0) cache_write = int(usage.get("cache_creation_tokens", 0) or 0) - out_dict: dict[str, int | str] = { + out_dict: UsageMetadata = { "provider": provider, "model": response.model or "", "prompt_tokens": inp, @@ -195,15 +196,13 @@ def extract_usage(response: Any) -> dict[str, int | str] | None: "cache_write_tokens": cache_write, "cached_tokens": cache_read + cache_write, "cache_creation_tokens": cache_write, + # Reasoning/thinking tokens (Anthropic extended thinking / + # OpenAI reasoning models). Part of completion_tokens, but + # surfaced separately for cost / analysis; the client's usage + # dict carries them. Present even as 0 — see the key-set + # invariant on the ``UsageMetadata`` return type. + "reasoning_tokens": int(usage.get("reasoning_tokens", 0) or 0), } - # Reasoning/thinking tokens (Anthropic extended thinking / OpenAI - # reasoning models). They are part of completion_tokens but surfaced - # separately for cost / analysis; the client's usage dict carries them. - # Always present, including as 0: this branch and the legacy shapes - # below must return the SAME key set, or a consumer indexing - # ``usage["reasoning_tokens"]`` works on one response object and - # raises KeyError on the other. - out_dict["reasoning_tokens"] = int(usage.get("reasoning_tokens", 0) or 0) return out_dict rmd = getattr(response, "response_metadata", None) or {} @@ -214,7 +213,12 @@ def extract_usage(response: Any) -> dict[str, int | str] | None: # own ``model_name`` lands on non-streaming responses but not on # streamed usage chunks). Falling through to it keeps streaming usage # attribution alive. - model = ( + # ``str(...)`` for the same reason ``provider`` gets it: these come off + # an untyped provider/gateway metadata dict, and ``UsageMetadata`` declares + # ``model`` as ``str``. A gateway echoing a non-string here — a nested + # dict, say — would otherwise put that object behind a field consumers + # format as text. + model = str( rmd.get("model_name") or rmd.get("model") or rmd.get("model_actually_used") @@ -228,7 +232,7 @@ def _build( cached: int, cache_create: int, reasoning: int, - ) -> dict[str, int | str]: + ) -> UsageMetadata: # ``cached`` carries cache READ; ``cache_create`` carries cache # WRITE. The legacy ``cached_tokens`` / ``cache_creation_tokens`` # keys are kept as a derived sum and an alias respectively so diff --git a/docs/llm-runtime-boundary.md b/docs/llm-runtime-boundary.md index 724982f..2fd0438 100644 --- a/docs/llm-runtime-boundary.md +++ b/docs/llm-runtime-boundary.md @@ -33,3 +33,42 @@ translate semantic retry intent into provider-specific request fields. Model profiles, provider client construction, tool parsing/execution, and the agent loop remain product-owned in this phase. + +## The usage contract + +`extract_usage` returns a `UsageMetadata` (`agent_core.loop_types`). Every +non-`None` return carries all nine normalized keys — `provider`, `model`, +`prompt_tokens`, `completion_tokens`, `cache_read_tokens`, +`cache_write_tokens`, `cached_tokens`, `cache_creation_tokens`, +`reasoning_tokens` — zero-filled when the provider reported nothing, whichever +of the three response shapes it parsed. That is a runtime invariant, not just a +declaration: a TypedDict validates nothing at import time, so +`tests/test_llm_runtime_extract_usage.py` pins each branch's key set against +`UsageMetadata.__required_keys__`. Consumers can index those nine directly. + +The `usage` fields on `TurnContext` and `LLMAttemptContext` are deliberately +*wider* — `Mapping[str, Any] | None`, not the TypedDict. Products, not +AgentCore, construct those contexts, and the shapes they hand over are ones a +TypedDict rejects outright: + +- partial literals in test doubles — `usage={"prompt_tokens": 182_000}`; +- alias-only shapes — `usage={"input_tokens": …, "output_tokens": …}`, which + ApodexHarness' budget observer reads; +- `dict(event["usage"])` re-wraps out of a `dict[str, Any]` attempt event, + which the ApodexHarness loop stamps `provider` / `model` onto before + constructing `LLMAttemptContext`. + +Neither `dict[str, int]` nor `dict[str, Any]` is assignable to a TypedDict, in +strict *or* standard mode, so narrowing these fields would break the product +loops the normalized contract exists to serve. `Mapping` rather than +`dict[str, Any]` because it is covariant in its value type: it accepts all of +the above *and* a `UsageMetadata`, which a `dict[str, Any]` field would reject. +A consumer wanting the precise shape annotates its own parameter +`UsageMetadata`; the boundary does not force that on producers. Hosts stamping +extra keys is likewise their business — a TypedDict cannot express "open" on +Python 3.12 (PEP 728 lands later). `UsageMetadataExtras` declares the aliases +hosts are known to add, for anyone who wants to describe such a mapping as a +`UsageMetadata`. + +`tests/test_loop_types.py` asserts these fields stay an open mapping, so a +later well-meant tightening fails loudly. diff --git a/tests/test_llm_runtime_extract_usage.py b/tests/test_llm_runtime_extract_usage.py index 4f81620..e26cc95 100644 --- a/tests/test_llm_runtime_extract_usage.py +++ b/tests/test_llm_runtime_extract_usage.py @@ -31,7 +31,10 @@ from types import SimpleNamespace +import pytest + from agent_core.llm import LLMResponse +from agent_core.loop_types import UsageMetadata from agent_core.runtime.loop._response import _pick_int from agent_core.runtime.loop.llm_client import extract_usage @@ -603,3 +606,71 @@ def __getitem__(self, k): assert u["prompt_tokens"] == 7 assert u["completion_tokens"] == 3 assert u["model"] == "x" + + +# --- key-set invariant --------------------------------------------------- +# +# The three response shapes are parsed by three separate branches, and a +# consumer indexing ``usage["reasoning_tokens"]`` must not work on one +# response object and raise KeyError on another. ``UsageMetadata`` states +# that invariant in the type system; a TypedDict validates nothing at +# runtime, so pin it here against the declared required keys — the type and +# all three branches move together or this fails. + + +_NATIVE = LLMResponse( + content="x", + model="m", + usage={"prompt_tokens": 1, "completion_tokens": 2}, + response_metadata={"provider_actually_used": "p"}, +) +_CANONICAL = SimpleNamespace( + usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={"model_name": "m"}, +) +_RAW = SimpleNamespace( + usage_metadata=None, + response_metadata={ + "token_usage": {"prompt_tokens": 1, "completion_tokens": 2}, + "model_name": "m", + }, +) + + +@pytest.mark.parametrize( + ("label", "response"), + [("native", _NATIVE), ("canonical", _CANONICAL), ("raw", _RAW)], +) +def test_every_branch_returns_exactly_the_required_key_set( + label: str, response: object, +) -> None: + usage = extract_usage(response) + assert usage is not None, label + assert set(usage) == UsageMetadata.__required_keys__, label + + +def test_model_is_coerced_to_str_on_the_legacy_path() -> None: + """``UsageMetadata`` declares ``model: str``; the source dict is untyped. + + A gateway echoing a non-string ``model_name`` would land that object + verbatim behind a field consumers format as text. ``provider`` was + already coerced; ``model`` was the asymmetric one. + """ + usage = extract_usage( + _resp( + usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={"model_name": {"id": "oops"}}, + ) + ) + assert usage is not None + assert isinstance(usage["model"], str) + assert usage["model"] == "{'id': 'oops'}" + + +def test_model_falls_back_to_empty_string_not_none() -> None: + usage = extract_usage( + _resp(usage_metadata={"input_tokens": 1, "output_tokens": 2}, + response_metadata={}) + ) + assert usage is not None + assert usage["model"] == "" diff --git a/tests/test_loop_types.py b/tests/test_loop_types.py index 4ad1979..6822368 100644 --- a/tests/test_loop_types.py +++ b/tests/test_loop_types.py @@ -8,6 +8,8 @@ from __future__ import annotations import asyncio +from collections.abc import Mapping +from typing import Any, get_type_hints import pytest @@ -19,6 +21,45 @@ merge_interventions = lt.merge_interventions +def test_usage_metadata_contract_marks_normalized_keys_required(): + assert lt.UsageMetadata.__required_keys__ == frozenset( + { + "provider", + "model", + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_write_tokens", + "cached_tokens", + "cache_creation_tokens", + "reasoning_tokens", + } + ) + assert lt.UsageMetadata.__optional_keys__ == frozenset( + {"input_tokens", "output_tokens", "total_tokens", "estimated"} + ) + assert "UsageMetadata" in lt.__all__ + assert "UsageMetadataExtras" in lt.__all__ + + +@pytest.mark.parametrize("cls_name", ["TurnContext", "LLMAttemptContext"]) +def test_observer_usage_field_stays_an_open_mapping(cls_name: str): + """``usage`` must NOT be typed as the ``UsageMetadata`` TypedDict. + + Hosts are the producers of these contexts and they hand over shapes a + TypedDict rejects outright: partial literals in test doubles + (``{"prompt_tokens": 182_000}``), alias-only budget shapes + (``{"input_tokens": .., "output_tokens": ..}``), and + ``dict(event["usage"])`` re-wraps whose value type is ``Any``. None of + those are assignable to a TypedDict, under strict OR standard mode, so + narrowing this field breaks the exact product loops the normalized + contract was introduced to unblock. ``UsageMetadata`` stays the + producer-side return type of ``extract_usage`` instead. + """ + usage_type = get_type_hints(getattr(lt, cls_name))["usage"] + assert usage_type == Mapping[str, Any] | None + + def _tool_result(name: str = "t", result: str = "r"): return ToolResult( name=name,