Enhance OpenAI Stream Wrapper to Include Token Count Metrics

Author: fenilfaldu

agentops/instrumentation/openai/stream_wrapper.py

@@ -14,6 +14,7 @@
 
 from agentops.logging import logger
 from agentops.instrumentation.common.wrappers import _with_tracer_wrapper
+from agentops.instrumentation.openai.utils import is_metrics_enabled
 from agentops.instrumentation.openai.wrappers.chat import handle_chat_attributes
 from agentops.semconv import SpanAttributes, LLMRequestTypeValues, MessageAttributes
 
@@ -195,22 +196,22 @@ def _finalize_stream(self) -> None:
         if len(self._tool_calls) > 0:
             for idx, tool_call in self._tool_calls.items():
                 # Only set attributes if values are not None
-                if tool_call["id"]:
+                if tool_call["id"] is not None:
                     self._span.set_attribute(
                         MessageAttributes.COMPLETION_TOOL_CALL_ID.format(i=0, j=idx), tool_call["id"]
                     )
 
-                if tool_call["type"]:
+                if tool_call["type"] is not None:
                     self._span.set_attribute(
                         MessageAttributes.COMPLETION_TOOL_CALL_TYPE.format(i=0, j=idx), tool_call["type"]
                     )
 
-                if tool_call["function"]["name"]:
+                if tool_call["function"]["name"] is not None:
                     self._span.set_attribute(
                         MessageAttributes.COMPLETION_TOOL_CALL_NAME.format(i=0, j=idx), tool_call["function"]["name"]
                     )
 
-                if tool_call["function"]["arguments"]:
+                if tool_call["function"]["arguments"] is not None:
                     self._span.set_attribute(
                         MessageAttributes.COMPLETION_TOOL_CALL_ARGUMENTS.format(i=0, j=idx),
                         tool_call["function"]["arguments"],
@@ -339,6 +340,19 @@ def chat_completion_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         for key, value in request_attributes.items():
             span.set_attribute(key, value)
 
+        # Add include_usage to get token counts for streaming responses
+        if is_streaming and is_metrics_enabled():
+            # Add stream_options if it doesn't exist
+            if "stream_options" not in kwargs:
+                kwargs["stream_options"] = {"include_usage": True}
+                logger.debug("[OPENAI WRAPPER] Adding stream_options.include_usage=True to get token counts")
+            # If stream_options exists but doesn't have include_usage, add it
+            elif isinstance(kwargs["stream_options"], dict) and "include_usage" not in kwargs["stream_options"]:
+                kwargs["stream_options"]["include_usage"] = True
+                logger.debug(
+                    "[OPENAI WRAPPER] Adding include_usage=True to existing stream_options to get token counts"
+                )
+
         # Call the original method
         response = wrapped(*args, **kwargs)
 
@@ -395,6 +409,15 @@ async def async_chat_completion_stream_wrapper(tracer, wrapped, instance, args,
         for key, value in request_attributes.items():
             span.set_attribute(key, value)
 
+        # Add include_usage to get token counts for streaming responses
+        if is_streaming and is_metrics_enabled():
+            # Add stream_options if it doesn't exist
+            if "stream_options" not in kwargs:
+                kwargs["stream_options"] = {"include_usage": True}
+            # If stream_options exists but doesn't have include_usage, add it
+            elif isinstance(kwargs["stream_options"], dict) and "include_usage" not in kwargs["stream_options"]:
+                kwargs["stream_options"]["include_usage"] = True
+
         # Call the original method
         response = await wrapped(*args, **kwargs)
 

agentops/instrumentation/openai/wrappers/chat.py

@@ -137,7 +137,7 @@ def handle_chat_attributes(
             attributes[SpanAttributes.LLM_RESPONSE_ID] = response_dict["id"]
         if "model" in response_dict:
             attributes[SpanAttributes.LLM_RESPONSE_MODEL] = response_dict["model"]
-        if "system_fingerprint" in response_dict:
+        if "system_fingerprint" in response_dict and response_dict["system_fingerprint"] is not None:
             attributes[SpanAttributes.LLM_OPENAI_RESPONSE_SYSTEM_FINGERPRINT] = response_dict["system_fingerprint"]
 
         # Usage
@@ -176,9 +176,9 @@ def handle_chat_attributes(
                 if message:
                     if "role" in message:
                         attributes[f"{prefix}.role"] = message["role"]
-                    if "content" in message:
+                    if "content" in message and message["content"] is not None:
                         attributes[f"{prefix}.content"] = message["content"]
-                    if "refusal" in message:
+                    if "refusal" in message and message["refusal"] is not None:
                         attributes[f"{prefix}.refusal"] = message["refusal"]
 
                     # Function call