From 5435347e24349155a4794fddb4c0e2557ede9488 Mon Sep 17 00:00:00 2001
From: Tanisha-fil <tanisha@fil.org>
Date: Fri, 14 Aug 2026 14:27:25 +0400
Subject: [PATCH] Account for OpenAI Responses cached tokens

---
 src/art/api_costs.py              |  8 +++-
 tests/unit/test_track_api_cost.py | 66 +++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/art/api_costs.py b/src/art/api_costs.py
index 50c7cd754..8d7b63802 100644
--- a/src/art/api_costs.py
+++ b/src/art/api_costs.py
@@ -206,9 +206,15 @@ def _extract_direct_response_cost(response: Any) -> float | None:
 def _extract_openai_token_counts(response: Any) -> _OpenAITokenUsage | None:
     usage = _response_usage(response)
     prompt_tokens = _read_usage_field(usage, "prompt_tokens")
+    if prompt_tokens is None:
+        prompt_tokens = _read_usage_field(usage, "input_tokens")
     completion_tokens = _read_usage_field(usage, "completion_tokens")
+    if completion_tokens is None:
+        completion_tokens = _read_usage_field(usage, "output_tokens")
     cached_prompt_tokens = (
-        _read_usage_nested_field(usage, "prompt_tokens_details", "cached_tokens") or 0.0
+        _read_usage_nested_field(usage, "prompt_tokens_details", "cached_tokens")
+        or _read_usage_nested_field(usage, "input_tokens_details", "cached_tokens")
+        or 0.0
     )
     if (
         prompt_tokens is None
diff --git a/tests/unit/test_track_api_cost.py b/tests/unit/test_track_api_cost.py
index fbd938dbc..975364e00 100644
--- a/tests/unit/test_track_api_cost.py
+++ b/tests/unit/test_track_api_cost.py
@@ -60,6 +60,40 @@ def __init__(
         self.model = model
 
 
+class _OpenAIResponsesUsage:
+    def __init__(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        *,
+        cached_tokens: int = 0,
+    ) -> None:
+        self.input_tokens = input_tokens
+        self.output_tokens = output_tokens
+        self.input_tokens_details = type(
+            "InputTokensDetails",
+            (),
+            {"cached_tokens": cached_tokens},
+        )()
+
+
+class _OpenAIResponsesResponse:
+    def __init__(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        *,
+        cached_tokens: int = 0,
+        model: str | None = None,
+    ) -> None:
+        self.usage = _OpenAIResponsesUsage(
+            input_tokens,
+            output_tokens,
+            cached_tokens=cached_tokens,
+        )
+        self.model = model
+
+
 class _AnthropicUsage:
     def __init__(
         self,
@@ -162,6 +196,38 @@ async def _judge() -> _OpenAIResponse:
         metrics = await builder.flush()
         assert metrics["costs/train/llm_judge/cached_openai"] == pytest.approx(0.00255)
 
+    @pytest.mark.asyncio
+    async def test_openai_cost_extraction_accounts_for_responses_cached_tokens(
+        self,
+    ) -> None:
+        builder = MetricsBuilder(cost_context="train")
+
+        @track_api_cost(
+            source="llm_judge/cached_openai_responses",
+            provider="openai",
+            model_name="openai/gpt-4.1",
+            prompt_price_per_million=2.0,
+            completion_price_per_million=8.0,
+            cached_prompt_price_per_million=0.5,
+        )
+        async def _judge() -> _OpenAIResponsesResponse:
+            return _OpenAIResponsesResponse(
+                input_tokens=2_000,
+                output_tokens=100,
+                cached_tokens=1_500,
+            )
+
+        token = builder.activate()
+        try:
+            await _judge()
+        finally:
+            token.var.reset(token)
+
+        metrics = await builder.flush()
+        assert metrics["costs/train/llm_judge/cached_openai_responses"] == pytest.approx(
+            0.00255
+        )
+
     @pytest.mark.asyncio
     async def test_anthropic_cost_extraction_uses_registered_model_pricing(
         self,