diff --git a/agentops/instrumentation/providers/openai/wrappers/chat.py b/agentops/instrumentation/providers/openai/wrappers/chat.py index df730ee9e..a57ec1053 100644 --- a/agentops/instrumentation/providers/openai/wrappers/chat.py +++ b/agentops/instrumentation/providers/openai/wrappers/chat.py @@ -197,9 +197,9 @@ def handle_chat_attributes( attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = usage["completion_tokens"] # Reasoning tokens - output_details = usage.get("output_tokens_details", {}) - if isinstance(output_details, dict) and "reasoning_tokens" in output_details: - attributes[SpanAttributes.LLM_USAGE_REASONING_TOKENS] = output_details["reasoning_tokens"] + completion_details = usage.get("completion_tokens_details") + if isinstance(completion_details, dict) and completion_details.get("reasoning_tokens") is not None: + attributes[SpanAttributes.LLM_USAGE_REASONING_TOKENS] = completion_details["reasoning_tokens"] # Choices if should_send_prompts() and "choices" in response_dict: diff --git a/tests/unit/instrumentation/openai_core/test_chat_attributes.py b/tests/unit/instrumentation/openai_core/test_chat_attributes.py new file mode 100644 index 000000000..4c2d90981 --- /dev/null +++ b/tests/unit/instrumentation/openai_core/test_chat_attributes.py @@ -0,0 +1,41 @@ +from openai.types.chat import ChatCompletion + +from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes +from agentops.semconv import SpanAttributes + + +def _chat_completion(usage): + return ChatCompletion.model_validate( + { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 0, + "model": "o3-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": usage, + } + ) + + +class TestChatAttributes: + def test_reasoning_tokens_from_completion_tokens_details(self): + response = _chat_completion( + { + "prompt_tokens": 10, + "completion_tokens": 200, + "total_tokens": 210, + "completion_tokens_details": {"reasoning_tokens": 192}, + } + ) + + attributes = handle_chat_attributes(kwargs={"model": "o3-mini"}, return_value=response) + + assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 200 + assert attributes[SpanAttributes.LLM_USAGE_REASONING_TOKENS] == 192 + + def test_no_reasoning_tokens_without_completion_tokens_details(self): + response = _chat_completion({"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}) + + attributes = handle_chat_attributes(kwargs={"model": "gpt-4o"}, return_value=response) + + assert SpanAttributes.LLM_USAGE_REASONING_TOKENS not in attributes