From 7ba7531ca38df736237b55f1149e597446d76236 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 11 Sep 2026 11:32:17 -0700 Subject: [PATCH 1/4] [openai] Record cache-write and modality token usage Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../.changelog/603.added | 1 + .../genai/openai/chat_wrappers.py | 9 +- .../instrumentation/genai/openai/patch.py | 2 + .../instrumentation/genai/openai/utils.py | 28 ++ .../tests/test_chat_token_usage.py | 281 ++++++++++++++++++ 5 files changed, 320 insertions(+), 1 deletion(-) create mode 100644 instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added create mode 100644 instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added b/instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added new file mode 100644 index 000000000..54b8ed18b --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added @@ -0,0 +1 @@ +Record cache-write and modality token usage from OpenAI Chat Completions responses and streams. diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py index 0e9fbf053..b2705ed80 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py @@ -25,7 +25,11 @@ ) from .chat_buffers import ChoiceBuffer -from .utils import get_property_value, map_finish_reason +from .utils import ( + get_property_value, + map_finish_reason, + set_chat_usage_details, +) _logger = logging.getLogger(__name__) @@ -105,6 +109,9 @@ def _set_usage(self, chunk: ChatCompletionChunk) -> None: self._self_cached_prompt_tokens = get_property_value( prompt_tokens_details, "cached_tokens" ) + set_chat_usage_details( + invocation=self._self_invocation, usage=usage + ) def _process_chunk(self, chunk: ChatCompletionChunk) -> None: if not isinstance(chunk, ChatCompletionChunk): diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/patch.py b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/patch.py index cb976378d..ff449e529 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/patch.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/patch.py @@ -27,6 +27,7 @@ get_server_address_and_port, get_value, is_streaming, + set_chat_usage_details, ) _logger = logging.getLogger(__name__) @@ -208,6 +209,7 @@ def _set_response_properties( chat_invocation.cache_read_input_tokens = get_property_value( prompt_tokens_details, "cached_tokens" ) + set_chat_usage_details(invocation=chat_invocation, usage=result.usage) if getattr(result, "system_fingerprint", None): chat_invocation.attributes.update( diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/utils.py b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/utils.py index a291b88a3..db6d61621 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/utils.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/utils.py @@ -10,6 +10,7 @@ import openai from openai import NotGiven +from openai.types import CompletionUsage from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, @@ -65,6 +66,33 @@ def get_property_value(obj, property_name): return getattr(obj, property_name, None) +def set_chat_usage_details( + invocation: InferenceInvocation, usage: CompletionUsage +) -> None: + prompt_details: object = get_property_value(usage, "prompt_tokens_details") + completion_details: object = get_property_value( + usage, "completion_tokens_details" + ) + invocation.cache_write_input_tokens = get_property_value( + prompt_details, "cache_write_tokens" + ) + invocation.text_input_tokens = get_property_value( + prompt_details, "text_tokens" + ) + invocation.image_input_tokens = get_property_value( + prompt_details, "image_tokens" + ) + invocation.audio_input_tokens = get_property_value( + prompt_details, "audio_tokens" + ) + invocation.text_output_tokens = get_property_value( + completion_details, "text_tokens" + ) + invocation.audio_output_tokens = get_property_value( + completion_details, "audio_tokens" + ) + + def get_server_address_and_port( client_instance, ) -> tuple[str | None, int | None]: diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py new file mode 100644 index 000000000..a51ae6f0a --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py @@ -0,0 +1,281 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from collections.abc import Iterator +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread + +import pytest +from openai import AsyncOpenAI, OpenAI + +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, +) + + +@pytest.fixture(autouse=True) +def fixture_vcr() -> Iterator[None]: + yield + + +@pytest.fixture( + params=["instrument_no_content", "instrument_with_content"], + autouse=True, +) +def instrumentation( + request: pytest.FixtureRequest, content_mode: tuple[bool, str] +) -> None: + request.getfixturevalue(request.param) + + +@pytest.fixture( + params=[ + pytest.param( + { + "prompt_tokens_details": { + "cached_tokens": 5, + "cache_write_tokens": 10, + "text_tokens": 70, + "image_tokens": 20, + "audio_tokens": 10, + }, + "completion_tokens_details": { + "text_tokens": 18, + "audio_tokens": 2, + }, + }, + id="all-details", + ), + pytest.param( + {"prompt_tokens_details": {"audio_tokens": 10}}, + id="partial-details", + ), + pytest.param({}, id="absent-details"), + pytest.param( + { + "prompt_tokens_details": None, + "completion_tokens_details": None, + }, + id="null-details", + ), + pytest.param( + { + "prompt_tokens_details": { + "cached_tokens": 0, + "cache_write_tokens": 0, + "text_tokens": 0, + "image_tokens": 0, + "audio_tokens": 0, + }, + "completion_tokens_details": { + "text_tokens": 0, + "audio_tokens": 0, + }, + }, + id="zero-details", + ), + ], +) +def usage(request: pytest.FixtureRequest) -> dict[str, object]: + return { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + **request.param, + } + + +@pytest.fixture +def api_url(usage: dict[str, object]) -> Iterator[str]: + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + body = json.loads( + self.rfile.read(int(self.headers["Content-Length"])) + ) + response: dict[str, object] = { + "id": "chatcmpl-usage", + "created": 1, + "model": "gpt-4", + } + if body["stream"]: + response["object"] = "chat.completion.chunk" + chunks = [ + { + **response, + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "content": "hello", + }, + "finish_reason": None, + } + ], + }, + { + **response, + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop", + } + ], + }, + {**response, "choices": [], "usage": usage}, + ] + payload = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" for chunk in chunks + ) + + "data: [DONE]\n\n" + ).encode() + content_type = "text/event-stream" + else: + response.update( + object="chat.completion", + usage=usage, + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": "hello", + }, + "finish_reason": "stop", + } + ], + ) + payload = json.dumps(response).encode() + content_type = "application/json" + + self.send_response(HTTPStatus.OK) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + thread = Thread( + target=server.serve_forever, + kwargs={"poll_interval": 0.01}, + daemon=True, + ) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/v1" + finally: + server.shutdown() + thread.join(timeout=5) + + +def assert_usage( + span_exporter: InMemorySpanExporter, usage: dict[str, object] +) -> None: + (span,) = span_exporter.get_finished_spans() + assert span.attributes is not None + assert span.attributes[GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS] == 100 + assert span.attributes[GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span.attributes[GenAIAttributes.GEN_AI_RESPONSE_FINISH_REASONS] == ( + "stop", + ) + expected: dict[str, int] = { + GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS: 100, + GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS: 20, + } + for field, suffix, mapping in ( + ( + "prompt_tokens_details", + "input_tokens", + { + "cached_tokens": "cache_read", + "cache_write_tokens": "cache_write", + "text_tokens": "text", + "image_tokens": "image", + "audio_tokens": "audio", + }, + ), + ( + "completion_tokens_details", + "output_tokens", + {"text_tokens": "text", "audio_tokens": "audio"}, + ), + ): + details = usage.get(field) + if isinstance(details, dict): + for name, attribute in mapping.items(): + if value := details.get(name): + expected[f"gen_ai.usage.{attribute}.{suffix}"] = value + actual = { + key: value + for key, value in span.attributes.items() + if key.startswith("gen_ai.usage.") + } + assert actual == expected + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_chat_detailed_token_usage( + api_url: str, + streaming: bool, + usage: dict[str, object], + span_exporter: InMemorySpanExporter, +) -> None: + with OpenAI(base_url=api_url, api_key="test", max_retries=0) as client: + response = client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hello"}], + stream=streaming, + **( + {"stream_options": {"include_usage": True}} + if streaming + else {} + ), + ) + if streaming: + with response: + chunks = list(response) + assert chunks[-1].usage.prompt_tokens == 100 + else: + assert response.usage.prompt_tokens == 100 + + assert_usage(span_exporter=span_exporter, usage=usage) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_async_chat_detailed_token_usage( + api_url: str, + streaming: bool, + usage: dict[str, object], + span_exporter: InMemorySpanExporter, +) -> None: + async with AsyncOpenAI( + base_url=api_url, api_key="test", max_retries=0 + ) as client: + response = await client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hello"}], + stream=streaming, + **( + {"stream_options": {"include_usage": True}} + if streaming + else {} + ), + ) + if streaming: + async with response: + chunks = [chunk async for chunk in response] + assert chunks[-1].usage.prompt_tokens == 100 + else: + assert response.usage.prompt_tokens == 100 + + assert_usage(span_exporter=span_exporter, usage=usage) From 6fdd9f71862e158b907b95ed909d273c1eb7bb5d Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 11 Sep 2026 11:35:34 -0700 Subject: [PATCH 2/4] Associate the changelog entry with its pull request Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../.changelog/{603.added => 681.added} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/{603.added => 681.added} (100%) diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added b/instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/681.added similarity index 100% rename from instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/603.added rename to instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/681.added From dd9d211f72cb41bcdc810c825dd8736465ca3483 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 11 Sep 2026 11:39:02 -0700 Subject: [PATCH 3/4] Show exported usage in the HTTP regression output Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../tests/test_chat_token_usage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py index a51ae6f0a..cb3495481 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py @@ -219,6 +219,7 @@ def assert_usage( for key, value in span.attributes.items() if key.startswith("gen_ai.usage.") } + print(json.dumps({"exported_usage": actual}, sort_keys=True)) assert actual == expected From 7ff833ef2bd8f3674aa382cc8e42acc93d68b526 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 11 Sep 2026 16:53:07 -0700 Subject: [PATCH 4/4] Buffer Chat Completions usage until stream cleanup Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../genai/openai/chat_wrappers.py | 54 +++- .../tests/test_chat_token_usage.py | 251 +++++++++--------- 2 files changed, 178 insertions(+), 127 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py index eda5e3d87..f750ddc2b 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/chat_wrappers.py @@ -28,7 +28,6 @@ from .utils import ( get_property_value, map_finish_reason, - set_chat_usage_details, ) _logger = logging.getLogger(__name__) @@ -45,6 +44,12 @@ class _ChatStreamMixin: _self_prompt_tokens: int | None _self_completion_tokens: int | None _self_cached_prompt_tokens: int | None + _self_cache_write_prompt_tokens: int | None + _self_text_prompt_tokens: int | None + _self_image_prompt_tokens: int | None + _self_audio_prompt_tokens: int | None + _self_text_completion_tokens: int | None + _self_audio_completion_tokens: int | None _self_reasoning_tokens: int | None def _set_response_model(self, chunk: ChatCompletionChunk) -> None: @@ -110,12 +115,27 @@ def _set_usage(self, chunk: ChatCompletionChunk) -> None: self._self_cached_prompt_tokens = get_property_value( prompt_tokens_details, "cached_tokens" ) - set_chat_usage_details( - invocation=self._self_invocation, usage=usage + self._self_cache_write_prompt_tokens = get_property_value( + prompt_tokens_details, "cache_write_tokens" + ) + self._self_text_prompt_tokens = get_property_value( + prompt_tokens_details, "text_tokens" + ) + self._self_image_prompt_tokens = get_property_value( + prompt_tokens_details, "image_tokens" + ) + self._self_audio_prompt_tokens = get_property_value( + prompt_tokens_details, "audio_tokens" ) completion_tokens_details = getattr( usage, "completion_tokens_details", None ) + self._self_text_completion_tokens = get_property_value( + completion_tokens_details, "text_tokens" + ) + self._self_audio_completion_tokens = get_property_value( + completion_tokens_details, "audio_tokens" + ) if completion_tokens_details is not None: self._self_reasoning_tokens = get_property_value( completion_tokens_details, "reasoning_tokens" @@ -184,6 +204,22 @@ def _cleanup(self, error: BaseException | None = None) -> None: self._self_invocation.cache_read_input_tokens = ( self._self_cached_prompt_tokens ) + self._self_invocation.cache_write_input_tokens = ( + self._self_cache_write_prompt_tokens + ) + self._self_invocation.text_input_tokens = self._self_text_prompt_tokens + self._self_invocation.image_input_tokens = ( + self._self_image_prompt_tokens + ) + self._self_invocation.audio_input_tokens = ( + self._self_audio_prompt_tokens + ) + self._self_invocation.text_output_tokens = ( + self._self_text_completion_tokens + ) + self._self_invocation.audio_output_tokens = ( + self._self_audio_completion_tokens + ) self._self_invocation.thinking_tokens = self._self_reasoning_tokens finish_reasons = [ choice.finish_reason @@ -226,6 +262,12 @@ def __init__( self._self_prompt_tokens = None self._self_completion_tokens = None self._self_cached_prompt_tokens = None + self._self_cache_write_prompt_tokens = None + self._self_text_prompt_tokens = None + self._self_image_prompt_tokens = None + self._self_audio_prompt_tokens = None + self._self_text_completion_tokens = None + self._self_audio_completion_tokens = None self._self_reasoning_tokens = None @@ -248,6 +290,12 @@ def __init__( self._self_prompt_tokens = None self._self_completion_tokens = None self._self_cached_prompt_tokens = None + self._self_cache_write_prompt_tokens = None + self._self_text_prompt_tokens = None + self._self_image_prompt_tokens = None + self._self_audio_prompt_tokens = None + self._self_text_completion_tokens = None + self._self_audio_completion_tokens = None self._self_reasoning_tokens = None diff --git a/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py index cb3495481..5a864a4f5 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py +++ b/instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_chat_token_usage.py @@ -5,19 +5,25 @@ import json from collections.abc import Iterator -from http import HTTPStatus -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from threading import Thread +from typing import Any +from unittest.mock import AsyncMock, MagicMock, Mock import pytest -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream +from openai.types.chat import ChatCompletion, ChatCompletionChunk +from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta +from opentelemetry.instrumentation.genai.openai.chat_wrappers import ( + AsyncChatStreamWrapper, + ChatStreamWrapper, +) from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( InMemorySpanExporter, ) from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, ) +from opentelemetry.util.genai.invocation import InferenceInvocation @pytest.fixture(autouse=True) @@ -83,7 +89,7 @@ def instrumentation( ), ], ) -def usage(request: pytest.FixtureRequest) -> dict[str, object]: +def usage(request: pytest.FixtureRequest) -> dict[str, Any]: return { "prompt_tokens": 100, "completion_tokens": 20, @@ -93,92 +99,68 @@ def usage(request: pytest.FixtureRequest) -> dict[str, object]: @pytest.fixture -def api_url(usage: dict[str, object]) -> Iterator[str]: - class Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: - body = json.loads( - self.rfile.read(int(self.headers["Content-Length"])) - ) - response: dict[str, object] = { - "id": "chatcmpl-usage", - "created": 1, - "model": "gpt-4", +def chat_completion(usage: dict[str, Any]) -> ChatCompletion: + response: dict[str, Any] = { + "id": "chatcmpl-usage", + "created": 1, + "model": "gpt-4", + "object": "chat.completion", + "usage": usage, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", } - if body["stream"]: - response["object"] = "chat.completion.chunk" - chunks = [ - { - **response, - "choices": [ - { - "index": 0, - "delta": { - "role": "assistant", - "content": "hello", - }, - "finish_reason": None, - } - ], - }, - { - **response, - "choices": [ - { - "index": 0, - "delta": {}, - "finish_reason": "stop", - } - ], - }, - {**response, "choices": [], "usage": usage}, - ] - payload = ( - "".join( - f"data: {json.dumps(chunk)}\n\n" for chunk in chunks - ) - + "data: [DONE]\n\n" - ).encode() - content_type = "text/event-stream" - else: - response.update( - object="chat.completion", - usage=usage, - choices=[ - { - "index": 0, - "message": { - "role": "assistant", - "content": "hello", - }, - "finish_reason": "stop", - } - ], + ], + } + return ChatCompletion(**response) + + +@pytest.fixture +def chat_chunks(chat_completion: ChatCompletion) -> list[ChatCompletionChunk]: + common: dict[str, Any] = { + "id": chat_completion.id, + "created": chat_completion.created, + "model": chat_completion.model, + "object": "chat.completion.chunk", + } + return [ + ChatCompletionChunk( + **common, + choices=[ + Choice( + index=0, + delta=ChoiceDelta(role="assistant", content="hello"), + finish_reason=None, ) - payload = json.dumps(response).encode() - content_type = "application/json" + ], + ), + ChatCompletionChunk( + **common, + choices=[ + Choice( + index=0, + delta=ChoiceDelta(), + finish_reason="stop", + ) + ], + ), + ChatCompletionChunk(**common, choices=[], usage=chat_completion.usage), + ] - self.send_response(HTTPStatus.OK) - self.send_header("Content-Type", content_type) - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: - thread = Thread( - target=server.serve_forever, - kwargs={"poll_interval": 0.01}, - daemon=True, - ) - thread.start() - try: - yield f"http://127.0.0.1:{server.server_port}/v1" - finally: - server.shutdown() - thread.join(timeout=5) +def assert_usage_is_buffered(invocation: InferenceInvocation) -> None: + assert invocation.cache_write_input_tokens is None + assert invocation.text_input_tokens is None + assert invocation.image_input_tokens is None + assert invocation.audio_input_tokens is None + assert invocation.text_output_tokens is None + assert invocation.audio_output_tokens is None def assert_usage( - span_exporter: InMemorySpanExporter, usage: dict[str, object] + span_exporter: InMemorySpanExporter, usage: dict[str, Any] ) -> None: (span,) = span_exporter.get_finished_spans() assert span.attributes is not None @@ -221,32 +203,44 @@ def assert_usage( } print(json.dumps({"exported_usage": actual}, sort_keys=True)) assert actual == expected + assert all(type(value) is int for value in actual.values()) @pytest.mark.parametrize("streaming", [False, True]) def test_chat_detailed_token_usage( - api_url: str, + openai_client: OpenAI, + monkeypatch: pytest.MonkeyPatch, + chat_completion: ChatCompletion, + chat_chunks: list[ChatCompletionChunk], streaming: bool, - usage: dict[str, object], + usage: dict[str, Any], span_exporter: InMemorySpanExporter, ) -> None: - with OpenAI(base_url=api_url, api_key="test", max_retries=0) as client: - response = client.chat.completions.create( - model="gpt-4", - messages=[{"role": "user", "content": "hello"}], - stream=streaming, - **( - {"stream_options": {"include_usage": True}} - if streaming - else {} - ), - ) - if streaming: - with response: - chunks = list(response) - assert chunks[-1].usage.prompt_tokens == 100 - else: - assert response.usage.prompt_tokens == 100 + result: ChatCompletion | MagicMock = chat_completion + if streaming: + result = MagicMock(spec=Stream) + result.__iter__.return_value = chat_chunks + monkeypatch.setattr( + target=openai_client, + name="request", + value=Mock(return_value=result), + ) + response = openai_client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hello"}], + stream=streaming, + **({"stream_options": {"include_usage": True}} if streaming else {}), + ) + if streaming: + assert isinstance(response, ChatStreamWrapper) + received: list[ChatCompletionChunk] = [] + with response: + for chunk in response: + assert_usage_is_buffered(invocation=response._self_invocation) + received.append(chunk) + assert received == chat_chunks + else: + assert response == chat_completion assert_usage(span_exporter=span_exporter, usage=usage) @@ -254,29 +248,38 @@ def test_chat_detailed_token_usage( @pytest.mark.asyncio @pytest.mark.parametrize("streaming", [False, True]) async def test_async_chat_detailed_token_usage( - api_url: str, + async_openai_client: AsyncOpenAI, + monkeypatch: pytest.MonkeyPatch, + chat_completion: ChatCompletion, + chat_chunks: list[ChatCompletionChunk], streaming: bool, - usage: dict[str, object], + usage: dict[str, Any], span_exporter: InMemorySpanExporter, ) -> None: - async with AsyncOpenAI( - base_url=api_url, api_key="test", max_retries=0 - ) as client: - response = await client.chat.completions.create( - model="gpt-4", - messages=[{"role": "user", "content": "hello"}], - stream=streaming, - **( - {"stream_options": {"include_usage": True}} - if streaming - else {} - ), - ) - if streaming: - async with response: - chunks = [chunk async for chunk in response] - assert chunks[-1].usage.prompt_tokens == 100 - else: - assert response.usage.prompt_tokens == 100 + result: ChatCompletion | MagicMock = chat_completion + if streaming: + result = MagicMock(spec=AsyncStream) + result.__aiter__.return_value = chat_chunks + monkeypatch.setattr( + target=async_openai_client, + name="request", + value=AsyncMock(return_value=result), + ) + response = await async_openai_client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hello"}], + stream=streaming, + **({"stream_options": {"include_usage": True}} if streaming else {}), + ) + if streaming: + assert isinstance(response, AsyncChatStreamWrapper) + received: list[ChatCompletionChunk] = [] + async with response: + async for chunk in response: + assert_usage_is_buffered(invocation=response._self_invocation) + received.append(chunk) + assert received == chat_chunks + else: + assert response == chat_completion assert_usage(span_exporter=span_exporter, usage=usage)