From 70bda74425948a86087dee2103abbb236193624e Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Mon, 10 Aug 2026 07:59:47 -0700 Subject: [PATCH 1/3] Fix Mistral cached token usage Map prompt cache hits from Mistral chat usage into the standard usage details. Add regression coverage for regular and streaming responses. --- .../agent_framework_mistral/_chat_client.py | 5 ++++ .../tests/mistral/test_mistral_chat_client.py | 30 ++++++++++++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py index 18f8ecd17f..237f66bbf4 100644 --- a/python/packages/mistral/agent_framework_mistral/_chat_client.py +++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py @@ -853,6 +853,11 @@ def _parse_usage(self, usage: Mapping[str, Any] | None) -> UsageDetails | None: details["output_token_count"] = value if (value := usage.get("total_tokens")) is not None: details["total_token_count"] = value + prompt_tokens_details = usage.get("prompt_tokens_details") + if isinstance(prompt_tokens_details, Mapping): + if (value := prompt_tokens_details.get("cached_tokens")) is not None: + details["prompt/cached_tokens"] = value + details["cache_read_input_token_count"] = value return details or None # endregion diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py index 5249ea1d9d..5862a7a526 100644 --- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py +++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py @@ -220,6 +220,27 @@ async def test_get_response_basic() -> None: assert server.last_request["messages"] == [{"role": "user", "content": "hi"}] +async def test_get_response_includes_cached_input_tokens() -> None: + client, _ = make_client( + json_response( + make_response_payload( + content="hello", + usage={ + "prompt_tokens": 100, + "completion_tokens": 7, + "total_tokens": 107, + "prompt_tokens_details": {"cached_tokens": 80}, + }, + ) + ) + ) + + response = await client.get_response([Message("user", ["hi"])]) + + assert response.usage_details is not None + assert response.usage_details["cache_read_input_token_count"] == 80 + + @pytest.mark.parametrize( ("status_code", "expected_exception"), [ @@ -684,7 +705,12 @@ async def test_streaming_response() -> None: make_chunk_payload(content="lo"), make_chunk_payload( finish_reason="stop", - usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + usage={ + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + "prompt_tokens_details": {"cached_tokens": 2}, + }, ), ) ) @@ -700,6 +726,8 @@ async def test_streaming_response() -> None: "input_token_count": 3, "output_token_count": 2, "total_token_count": 5, + "prompt/cached_tokens": 2, + "cache_read_input_token_count": 2, } assert server.last_request["stream"] is True From a1c1b44b05259f5fd9689c4e0572d8e3cd24327b Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Mon, 10 Aug 2026 08:38:09 -0700 Subject: [PATCH 2/3] Validate Mistral cached token usage --- .../agent_framework_mistral/_chat_client.py | 7 +++--- .../tests/mistral/test_mistral_chat_client.py | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py index 237f66bbf4..f3cd5ba420 100644 --- a/python/packages/mistral/agent_framework_mistral/_chat_client.py +++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py @@ -855,9 +855,10 @@ def _parse_usage(self, usage: Mapping[str, Any] | None) -> UsageDetails | None: details["total_token_count"] = value prompt_tokens_details = usage.get("prompt_tokens_details") if isinstance(prompt_tokens_details, Mapping): - if (value := prompt_tokens_details.get("cached_tokens")) is not None: - details["prompt/cached_tokens"] = value - details["cache_read_input_token_count"] = value + cached_tokens = prompt_tokens_details.get("cached_tokens") + if isinstance(cached_tokens, int) and not isinstance(cached_tokens, bool): + details["prompt/cached_tokens"] = cached_tokens + details["cache_read_input_token_count"] = cached_tokens return details or None # endregion diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py index 5862a7a526..cb4a8c9501 100644 --- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py +++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py @@ -241,6 +241,29 @@ async def test_get_response_includes_cached_input_tokens() -> None: assert response.usage_details["cache_read_input_token_count"] == 80 +@pytest.mark.parametrize("cached_tokens", ["80", 80.5, True, False]) +async def test_get_response_ignores_invalid_cached_input_tokens(cached_tokens: Any) -> None: + client, _ = make_client( + json_response( + make_response_payload( + content="hello", + usage={ + "prompt_tokens": 100, + "completion_tokens": 7, + "total_tokens": 107, + "prompt_tokens_details": {"cached_tokens": cached_tokens}, + }, + ) + ) + ) + + response = await client.get_response([Message("user", ["hi"])]) + + assert response.usage_details is not None + assert "prompt/cached_tokens" not in response.usage_details + assert "cache_read_input_token_count" not in response.usage_details + + @pytest.mark.parametrize( ("status_code", "expected_exception"), [ From 2d72515cd6d8f6ea5ba3508271876da6423252be Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Tue, 11 Aug 2026 16:31:15 -0700 Subject: [PATCH 3/3] fix(mistral): satisfy strict cached token typing Narrow prompt token details before reading cached_tokens so the Mistral package passes strict Pyright without changing runtime validation.\n\nAddresses https://github.com/microsoft/agent-framework/pull/7597#discussion_r3750712320 --- python/packages/mistral/agent_framework_mistral/_chat_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py index f3cd5ba420..bc55704392 100644 --- a/python/packages/mistral/agent_framework_mistral/_chat_client.py +++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py @@ -855,7 +855,7 @@ def _parse_usage(self, usage: Mapping[str, Any] | None) -> UsageDetails | None: details["total_token_count"] = value prompt_tokens_details = usage.get("prompt_tokens_details") if isinstance(prompt_tokens_details, Mapping): - cached_tokens = prompt_tokens_details.get("cached_tokens") + cached_tokens = cast("Mapping[str, Any]", prompt_tokens_details).get("cached_tokens") if isinstance(cached_tokens, int) and not isinstance(cached_tokens, bool): details["prompt/cached_tokens"] = cached_tokens details["cache_read_input_token_count"] = cached_tokens