diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 11032675e4..20629b76b8 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -354,7 +354,9 @@ def __str__(self) -> str: ("anthropic.cache_creation_input_tokens", OtelAttr.CACHE_CREATION_INPUT_TOKENS), ("anthropic.cache_read_input_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS), ("openai.cached_input_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS), + ("openai.cache_write_tokens", OtelAttr.CACHE_CREATION_INPUT_TOKENS), ("prompt/cached_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS), + ("prompt/cache_write_tokens", OtelAttr.CACHE_CREATION_INPUT_TOKENS), ("openai.reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS), ("completion/reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS), ("reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS), diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 3ee0278afa..53e99dc139 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -2383,6 +2383,28 @@ def test_get_response_attributes_maps_legacy_usage_keys(): assert result[OtelAttr.REASONING_OUTPUT_TOKENS] == 34 +def test_get_response_attributes_maps_openai_cache_write_tokens(): + """Test _get_response_attributes maps the OpenAI cache write usage key to the OTel attribute.""" + from unittest.mock import Mock + + from agent_framework.observability import OtelAttr, _get_response_attributes + + response = Mock() + response.response_id = None + response.finish_reason = None + response.raw_representation = None + response.usage_details = { + "openai.cache_write_tokens": 1024, + "openai.cached_input_tokens": 512, + } + + attrs: dict[str, Any] = {} + result = _get_response_attributes(attrs, response) + + assert result[OtelAttr.CACHE_CREATION_INPUT_TOKENS] == 1024 + assert result[OtelAttr.CACHE_READ_INPUT_TOKENS] == 512 + + def test_get_response_attributes_capture_usage_false(): """Test _get_response_attributes skips usage when capture_usage is False.""" from unittest.mock import Mock diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 92cfd6671d..fce5ecdb3b 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -3379,6 +3379,10 @@ def _parse_usage_from_openai(self, usage: ResponseUsage) -> UsageDetails | None: total_token_count=usage.total_tokens, ) if usage.input_tokens_details: + cache_write_tokens = cast("int | None", getattr(usage.input_tokens_details, "cache_write_tokens", None)) + if cache_write_tokens is not None: + details["openai.cache_write_tokens"] = cache_write_tokens + details["cache_creation_input_token_count"] = cache_write_tokens cached_tokens = cast("int | None", getattr(usage.input_tokens_details, "cached_tokens", None)) if cached_tokens is not None: details["openai.cached_input_tokens"] = cached_tokens diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index 69dd1b8225..37fc968a68 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -854,6 +854,10 @@ def _parse_usage_from_openai(self, usage: CompletionUsage) -> UsageDetails: if usage.prompt_tokens_details: if tokens := usage.prompt_tokens_details.audio_tokens: details["prompt/audio_tokens"] = tokens + cache_write_tokens = cast("int | None", getattr(usage.prompt_tokens_details, "cache_write_tokens", None)) + if cache_write_tokens is not None: + details["prompt/cache_write_tokens"] = cache_write_tokens + details["cache_creation_input_token_count"] = cache_write_tokens if (tokens := usage.prompt_tokens_details.cached_tokens) is not None: details["prompt/cached_tokens"] = tokens details["cache_read_input_token_count"] = tokens diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index ad0e061560..d52157842e 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -4458,6 +4458,7 @@ def test_usage_details_with_cached_tokens() -> None: mock_usage.total_tokens = 275 mock_usage.input_tokens_details = MagicMock() mock_usage.input_tokens_details.cached_tokens = 25 + mock_usage.input_tokens_details.cache_write_tokens = None mock_usage.output_tokens_details = None details = client._parse_usage_from_openai(mock_usage) # type: ignore @@ -4468,6 +4469,46 @@ def test_usage_details_with_cached_tokens() -> None: assert details["cache_read_input_token_count"] == 25 +def test_usage_details_with_cache_write_tokens() -> None: + """Test _parse_usage_from_openai with cache write tokens.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + mock_usage = MagicMock() + mock_usage.input_tokens = 2000 + mock_usage.output_tokens = 60 + mock_usage.total_tokens = 2060 + mock_usage.input_tokens_details = MagicMock() + mock_usage.input_tokens_details.cached_tokens = 0 + mock_usage.input_tokens_details.cache_write_tokens = 1024 + mock_usage.output_tokens_details = None + + details = client._parse_usage_from_openai(mock_usage) # type: ignore + assert details is not None + details_dict = cast("dict[str, Any]", details) + assert details_dict["openai.cache_write_tokens"] == 1024 + assert details["cache_creation_input_token_count"] == 1024 + assert details["cache_read_input_token_count"] == 0 + + +def test_usage_details_omits_missing_cache_write_tokens() -> None: + """Test _parse_usage_from_openai omits cache write tokens when the provider does not report them.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + mock_usage = MagicMock() + mock_usage.input_tokens = 100 + mock_usage.output_tokens = 20 + mock_usage.total_tokens = 120 + mock_usage.input_tokens_details = MagicMock(spec=["cached_tokens"]) + mock_usage.input_tokens_details.cached_tokens = 10 + mock_usage.output_tokens_details = None + + details = client._parse_usage_from_openai(mock_usage) # type: ignore + assert details is not None + assert "openai.cache_write_tokens" not in details + assert "cache_creation_input_token_count" not in details + assert details["cache_read_input_token_count"] == 10 + + def test_usage_details_with_reasoning_tokens() -> None: """Test _parse_usage_from_openai with reasoning tokens.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -4498,12 +4539,15 @@ def test_usage_details_with_zero_cached_and_reasoning_tokens() -> None: mock_usage.total_tokens = 230 mock_usage.input_tokens_details = MagicMock() mock_usage.input_tokens_details.cached_tokens = 0 + mock_usage.input_tokens_details.cache_write_tokens = 0 mock_usage.output_tokens_details = MagicMock() mock_usage.output_tokens_details.reasoning_tokens = 0 details = client._parse_usage_from_openai(mock_usage) # type: ignore assert details is not None details_dict = cast("dict[str, Any]", details) + assert details_dict["openai.cache_write_tokens"] == 0 + assert details["cache_creation_input_token_count"] == 0 assert details_dict["openai.cached_input_tokens"] == 0 assert details["cache_read_input_token_count"] == 0 assert details_dict["openai.reasoning_tokens"] == 0 @@ -4520,11 +4564,14 @@ def test_usage_details_omits_missing_cached_and_reasoning_tokens() -> None: mock_usage.total_tokens = 230 mock_usage.input_tokens_details = MagicMock() mock_usage.input_tokens_details.cached_tokens = None + mock_usage.input_tokens_details.cache_write_tokens = None mock_usage.output_tokens_details = MagicMock() mock_usage.output_tokens_details.reasoning_tokens = None details = client._parse_usage_from_openai(mock_usage) # type: ignore assert details is not None + assert "openai.cache_write_tokens" not in details + assert "cache_creation_input_token_count" not in details assert "openai.cached_input_tokens" not in details assert "cache_read_input_token_count" not in details assert "openai.reasoning_tokens" not in details diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index 2c84631b88..0a48280a8d 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -1179,6 +1179,48 @@ def test_parse_usage_includes_standard_and_legacy_mapped_token_details() -> None assert details["cache_read_input_token_count"] == 0 +def test_parse_usage_with_cache_write_tokens() -> None: + """Test _parse_usage_from_openai maps cache write tokens to standard and legacy keys.""" + client = OpenAIChatCompletionClient(model="test-model", api_key="test-key") + + mock_usage = MagicMock() + mock_usage.prompt_tokens = 2000 + mock_usage.completion_tokens = 60 + mock_usage.total_tokens = 2060 + mock_usage.completion_tokens_details = None + mock_usage.prompt_tokens_details = MagicMock() + mock_usage.prompt_tokens_details.audio_tokens = None + mock_usage.prompt_tokens_details.cached_tokens = 0 + mock_usage.prompt_tokens_details.cache_write_tokens = 1024 + + details = client._parse_usage_from_openai(mock_usage) # type: ignore[arg-type] + + details_dict = cast("dict[str, Any]", details) + assert details_dict["prompt/cache_write_tokens"] == 1024 + assert details["cache_creation_input_token_count"] == 1024 + assert details["cache_read_input_token_count"] == 0 + + +def test_parse_usage_omits_missing_cache_write_tokens() -> None: + """Test _parse_usage_from_openai omits cache write tokens when the provider does not report them.""" + client = OpenAIChatCompletionClient(model="test-model", api_key="test-key") + + mock_usage = MagicMock() + mock_usage.prompt_tokens = 100 + mock_usage.completion_tokens = 20 + mock_usage.total_tokens = 120 + mock_usage.completion_tokens_details = None + mock_usage.prompt_tokens_details = MagicMock(spec=["audio_tokens", "cached_tokens"]) + mock_usage.prompt_tokens_details.audio_tokens = None + mock_usage.prompt_tokens_details.cached_tokens = 10 + + details = client._parse_usage_from_openai(mock_usage) # type: ignore[arg-type] + + assert "prompt/cache_write_tokens" not in details + assert "cache_creation_input_token_count" not in details + assert details["cache_read_input_token_count"] == 10 + + def test_streaming_chunk_with_usage_and_text( openai_unit_test_env: dict[str, str], ) -> None: diff --git a/python/samples/02-agents/providers/openai/client_prompt_caching.py b/python/samples/02-agents/providers/openai/client_prompt_caching.py index 38afdaf93e..4d35989002 100644 --- a/python/samples/02-agents/providers/openai/client_prompt_caching.py +++ b/python/samples/02-agents/providers/openai/client_prompt_caching.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import time from agent_framework import Content, Message from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions @@ -15,8 +16,10 @@ writes are billed on these models, so marking exactly where a reusable prefix ends lets you control what gets cached. -Two knobs work together: +Three knobs work together: +- ``prompt_cache_key`` on ``OpenAIChatOptions`` is required to use the improved + prompt caching. - ``prompt_cache_options`` on ``OpenAIChatOptions`` sets the request-wide policy. ``{"mode": "explicit"}`` disables the automatic breakpoint on the latest message, so only the breakpoints you place are used for cache reads and writes. @@ -24,6 +27,8 @@ reusable prefix on a specific content part. The content before a breakpoint must be at least 1024 tokens long to be cached. +The number of tokens written to the cache is shown in +``usage_details["cache_creation_input_token_count"]`` on the first response. Running the same prefix twice shows the cache hit through ``usage_details["cache_read_input_token_count"]`` on later responses. @@ -64,16 +69,21 @@ async def main() -> None: print("\033[92m=== OpenAI Chat Client Prompt Caching Example ===\033[0m\n") client = OpenAIChatClient[OpenAIChatOptions](model="gpt-5.6-luna") - options: OpenAIChatOptions = {"prompt_cache_options": {"mode": "explicit"}} + options: OpenAIChatOptions = { + "prompt_cache_options": {"mode": "explicit"}, + "prompt_cache_key": f"contoso_appliance_store-{time.time()}", + } questions = ["Do you sell refrigerators?", "What is the return policy contact?"] for turn, question in enumerate(questions, start=1): response = await client.get_response(build_messages(question), options=options) usage = response.usage_details or {} cached = usage.get("cache_read_input_token_count", 0) + cached_write = usage.get("cache_creation_input_token_count", 0) print(f"Turn {turn}: {question}") print(f" Answer: {response.text}") - print(f" Cached input tokens: {cached}\n") + print(f" Cached input tokens (read): {cached}\n") + print(f" Cached input tokens (created): {cached_write}\n") if turn < len(questions): # A freshly written cache entry becomes readable shortly after the request # completes; the brief pause keeps the next turn from racing this one.