[object Object]

← back to Exo

Complete responses api usage response field (#1885)

77ffe039b3a1edd1390fbba893e30a6ce9094988 · 2026-04-13 18:38:33 +0100 · ciaranbor

## Motivation

The Responses API usage response was missing `input_tokens_details` and
`output_tokens_details`. The chat completions API already reports these.

## Changes

- Added `InputTokensDetails` (`cached_tokens`) and `OutputTokensDetails`
(`reasoning_tokens`) to `ResponseUsage`
- Extracted shared `_build_response_usage()` helper for both streaming
and non-streaming paths

## Test Plan

### Manual Testing

4-node cluster, `Qwen3-30B-A3B-4bit` — verified both detail objects
present with correct values in streaming and non-streaming responses.

### Automated Testing

13 tests in `test_openai_responses_api.py`.

Files touched

Diff

commit 77ffe039b3a1edd1390fbba893e30a6ce9094988
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date:   Mon Apr 13 18:38:33 2026 +0100

    Complete responses api usage response field (#1885)
    
    ## Motivation
    
    The Responses API usage response was missing `input_tokens_details` and
    `output_tokens_details`. The chat completions API already reports these.
    
    ## Changes
    
    - Added `InputTokensDetails` (`cached_tokens`) and `OutputTokensDetails`
    (`reasoning_tokens`) to `ResponseUsage`
    - Extracted shared `_build_response_usage()` helper for both streaming
    and non-streaming paths
    
    ## Test Plan
    
    ### Manual Testing
    
    4-node cluster, `Qwen3-30B-A3B-4bit` — verified both detail objects
    present with correct values in streaming and non-streaming responses.
    
    ### Automated Testing
    
    13 tests in `test_openai_responses_api.py`.
---
 src/exo/api/adapters/responses.py              |  33 ++--
 src/exo/api/tests/test_openai_responses_api.py | 225 ++++++++++++++++++++++++-
 src/exo/api/types/openai_responses.py          |  14 ++
 3 files changed, 257 insertions(+), 15 deletions(-)

diff --git a/src/exo/api/adapters/responses.py b/src/exo/api/adapters/responses.py
index a3c248c0..ad11b059 100644
--- a/src/exo/api/adapters/responses.py
+++ b/src/exo/api/adapters/responses.py
@@ -23,6 +23,7 @@ from exo.api.types.openai_responses import (
     FunctionCallInputItem,
     FunctionCallOutputInputItem,
     ImageGenerationCallInputItem,
+    InputTokensDetails,
     ItemReferenceInputItem,
     LocalShellCallInputItem,
     LocalShellCallOutputInputItem,
@@ -30,6 +31,7 @@ from exo.api.types.openai_responses import (
     McpApprovalResponseInputItem,
     McpCallInputItem,
     McpListToolsInputItem,
+    OutputTokensDetails,
     ReasoningInputItem,
     ResponseCompletedEvent,
     ResponseContentPart,
@@ -82,6 +84,21 @@ from exo.shared.types.text_generation import (
 )
 
 
+def _build_response_usage(usage: Usage) -> ResponseUsage:
+    """Build a ResponseUsage from the internal Usage type."""
+    return ResponseUsage(
+        input_tokens=usage.prompt_tokens,
+        input_tokens_details=InputTokensDetails(
+            cached_tokens=usage.prompt_tokens_details.cached_tokens,
+        ),
+        output_tokens=usage.completion_tokens,
+        output_tokens_details=OutputTokensDetails(
+            reasoning_tokens=usage.completion_tokens_details.reasoning_tokens,
+        ),
+        total_tokens=usage.total_tokens,
+    )
+
+
 def _format_sse(event: ResponsesStreamEvent) -> str:
     """Format a streaming event as an SSE message."""
     return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n"
@@ -436,13 +453,7 @@ async def collect_responses_response(
         raise ValueError(error_message)
 
     # Create usage from usage data if available
-    usage = None
-    if last_usage is not None:
-        usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
-        )
+    usage = _build_response_usage(last_usage) if last_usage is not None else None
 
     output: list[ResponseItem] = []
     if thinking_parts:
@@ -791,13 +802,7 @@ async def generate_responses_stream(
     yield _format_sse(item_done)
 
     # Create usage from usage data if available
-    usage = None
-    if last_usage is not None:
-        usage = ResponseUsage(
-            input_tokens=last_usage.prompt_tokens,
-            output_tokens=last_usage.completion_tokens,
-            total_tokens=last_usage.total_tokens,
-        )
+    usage = _build_response_usage(last_usage) if last_usage is not None else None
 
     # response.completed
     output: list[ResponseItem] = []
diff --git a/src/exo/api/tests/test_openai_responses_api.py b/src/exo/api/tests/test_openai_responses_api.py
index 7165df8e..3a319581 100644
--- a/src/exo/api/tests/test_openai_responses_api.py
+++ b/src/exo/api/tests/test_openai_responses_api.py
@@ -4,14 +4,30 @@ ResponsesRequest is the API wire type for the Responses endpoint.
 The responses adapter converts it to TextGenerationTaskParams for the pipeline.
 """
 
+import json
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
 import pydantic
 import pytest
 
+from exo.api.adapters.responses import (
+    collect_responses_response,
+    generate_responses_stream,
+)
+from exo.api.types import CompletionTokensDetails, PromptTokensDetails, Usage
 from exo.api.types.openai_responses import (
+    InputTokensDetails,
+    OutputTokensDetails,
     ResponseInputMessage,
     ResponsesRequest,
+    ResponsesResponse,
+    ResponseUsage,
 )
-from exo.shared.types.common import ModelId
+from exo.shared.types.chunks import TokenChunk
+from exo.shared.types.common import CommandId, ModelId
+
+_TEST_MODEL = ModelId("test-model")
 
 
 class TestResponsesRequestValidation:
@@ -46,3 +62,210 @@ class TestResponsesRequestValidation:
             input=[ResponseInputMessage(role="user", content="Hello")],
         )
         assert len(request.input) == 1
+
+
+class TestResponseUsage:
+    """Tests for ResponseUsage with input_tokens_details and output_tokens_details."""
+
+    def test_usage_defaults_to_zero_details(self):
+        usage = ResponseUsage(
+            input_tokens=10,
+            input_tokens_details=InputTokensDetails(),
+            output_tokens=20,
+            output_tokens_details=OutputTokensDetails(),
+            total_tokens=30,
+        )
+        assert usage.input_tokens_details.cached_tokens == 0
+        assert usage.output_tokens_details.reasoning_tokens == 0
+
+    def test_usage_with_reasoning_tokens(self):
+        usage = ResponseUsage(
+            input_tokens=10,
+            input_tokens_details=InputTokensDetails(),
+            output_tokens=20,
+            output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
+            total_tokens=30,
+        )
+        assert usage.output_tokens_details.reasoning_tokens == 5
+
+    def test_usage_with_cached_tokens(self):
+        usage = ResponseUsage(
+            input_tokens=10,
+            input_tokens_details=InputTokensDetails(cached_tokens=7),
+            output_tokens=20,
+            output_tokens_details=OutputTokensDetails(),
+            total_tokens=30,
+        )
+        assert usage.input_tokens_details.cached_tokens == 7
+
+    def test_usage_serialization(self):
+        usage = ResponseUsage(
+            input_tokens=10,
+            input_tokens_details=InputTokensDetails(cached_tokens=3),
+            output_tokens=20,
+            output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
+            total_tokens=30,
+        )
+        data = usage.model_dump()
+        assert data["input_tokens_details"] == {"cached_tokens": 3}
+        assert data["output_tokens_details"] == {"reasoning_tokens": 5}
+
+    def test_usage_serialization_zero_details(self):
+        usage = ResponseUsage(
+            input_tokens=10,
+            input_tokens_details=InputTokensDetails(),
+            output_tokens=20,
+            output_tokens_details=OutputTokensDetails(),
+            total_tokens=30,
+        )
+        data = usage.model_dump()
+        assert data["input_tokens_details"] == {"cached_tokens": 0}
+        assert data["output_tokens_details"] == {"reasoning_tokens": 0}
+
+
+def _make_usage(
+    prompt_tokens: int, completion_tokens: int, reasoning_tokens: int = 0
+) -> Usage:
+    """Create a Usage object for testing."""
+    return Usage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        prompt_tokens_details=PromptTokensDetails(),
+        completion_tokens_details=CompletionTokensDetails(
+            reasoning_tokens=reasoning_tokens
+        ),
+    )
+
+
+async def _token_chunks(
+    chunks: list[TokenChunk],
+) -> AsyncGenerator[TokenChunk, None]:
+    for chunk in chunks:
+        yield chunk
+
+
+class TestCollectResponsesResponseReasoningTokens:
+    """Tests for reasoning_tokens in collect_responses_response."""
+
+    async def test_non_streaming_includes_reasoning_tokens(self):
+        usage = _make_usage(10, 25, reasoning_tokens=8)
+        chunks = [
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=0,
+                text="thinking...",
+                is_thinking=True,
+                usage=None,
+            ),
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=1,
+                text="Hello world",
+                is_thinking=False,
+                usage=usage,
+            ),
+        ]
+        command_id = CommandId("test-cmd-001")
+        result_parts: list[str] = []
+        async for part in collect_responses_response(
+            command_id, "test-model", _token_chunks(chunks)
+        ):
+            result_parts.append(part)
+        assert len(result_parts) == 1
+        response = ResponsesResponse.model_validate_json(result_parts[0])
+        assert response.usage is not None
+        assert response.usage.input_tokens_details.cached_tokens == 0
+        assert response.usage.output_tokens_details.reasoning_tokens == 8
+
+    async def test_non_streaming_zero_reasoning_tokens(self):
+        usage = _make_usage(10, 20, reasoning_tokens=0)
+        chunks = [
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=0,
+                text="Hello world",
+                is_thinking=False,
+                usage=usage,
+            ),
+        ]
+        command_id = CommandId("test-cmd-002")
+        result_parts: list[str] = []
+        async for part in collect_responses_response(
+            command_id, "test-model", _token_chunks(chunks)
+        ):
+            result_parts.append(part)
+        assert len(result_parts) == 1
+        response = ResponsesResponse.model_validate_json(result_parts[0])
+        assert response.usage is not None
+        assert response.usage.output_tokens_details.reasoning_tokens == 0
+        assert response.usage.input_tokens_details.cached_tokens == 0
+
+
+class TestGenerateResponsesStreamReasoningTokens:
+    """Tests for reasoning_tokens in generate_responses_stream."""
+
+    async def test_streaming_includes_reasoning_tokens(self):
+        usage = _make_usage(10, 25, reasoning_tokens=8)
+        chunks = [
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=0,
+                text="thinking...",
+                is_thinking=True,
+                usage=None,
+            ),
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=1,
+                text="Hello world",
+                is_thinking=False,
+                usage=usage,
+            ),
+        ]
+        command_id = CommandId("test-cmd-003")
+        events: list[str] = []
+        async for event in generate_responses_stream(
+            command_id, "test-model", _token_chunks(chunks)
+        ):
+            events.append(event)
+
+        # The last event should be response.completed
+        last_event = events[-1]
+        # Parse the SSE data
+        data_line = [
+            line for line in last_event.strip().split("\n") if line.startswith("data: ")
+        ][0]
+        data = cast(dict[str, Any], json.loads(data_line.removeprefix("data: ")))
+        assert data["type"] == "response.completed"
+        response_usage = cast(dict[str, Any], data["response"]["usage"])
+        assert response_usage["input_tokens_details"] == {"cached_tokens": 0}
+        assert response_usage["output_tokens_details"] == {"reasoning_tokens": 8}
+
+    async def test_streaming_zero_reasoning_tokens(self):
+        usage = _make_usage(10, 20, reasoning_tokens=0)
+        chunks = [
+            TokenChunk(
+                model=_TEST_MODEL,
+                token_id=0,
+                text="Hello world",
+                is_thinking=False,
+                usage=usage,
+            ),
+        ]
+        command_id = CommandId("test-cmd-004")
+        events: list[str] = []
+        async for event in generate_responses_stream(
+            command_id, "test-model", _token_chunks(chunks)
+        ):
+            events.append(event)
+
+        last_event = events[-1]
+        data_line = [
+            line for line in last_event.strip().split("\n") if line.startswith("data: ")
+        ][0]
+        data = cast(dict[str, Any], json.loads(data_line.removeprefix("data: ")))
+        assert data["type"] == "response.completed"
+        response_usage = cast(dict[str, Any], data["response"]["usage"])
+        assert response_usage["input_tokens_details"] == {"cached_tokens": 0}
+        assert response_usage["output_tokens_details"] == {"reasoning_tokens": 0}
diff --git a/src/exo/api/types/openai_responses.py b/src/exo/api/types/openai_responses.py
index 0b747648..753b57d0 100644
--- a/src/exo/api/types/openai_responses.py
+++ b/src/exo/api/types/openai_responses.py
@@ -407,11 +407,25 @@ class ResponseReasoningItem(BaseModel, frozen=True):
 ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem
 
 
+class InputTokensDetails(BaseModel, frozen=True):
+    """Breakdown of input token counts in Responses API response."""
+
+    cached_tokens: int = 0
+
+
+class OutputTokensDetails(BaseModel, frozen=True):
+    """Breakdown of output token counts in Responses API response."""
+
+    reasoning_tokens: int = 0
+
+
 class ResponseUsage(BaseModel, frozen=True):
     """Token usage in Responses API response."""
 
     input_tokens: int
+    input_tokens_details: InputTokensDetails
     output_tokens: int
+    output_tokens_details: OutputTokensDetails
     total_tokens: int
 
 

← 3f0df404 Reduce memory consumption by adding Flash Attention to Qwen3  ·  back to Exo  ·  Add prefix cache flag to exo bench (#1888) f2709dcd →