fix: OpenAI Chat Completions tool calls (#1153)

* fix tool calls

* remove use of tracer

* linting

Author: dot-agi

agentops/instrumentation/providers/openai/stream_wrapper.py

@@ -15,7 +15,7 @@
 from agentops.logging import logger
 from agentops.instrumentation.common.wrappers import _with_tracer_wrapper
 from agentops.instrumentation.providers.openai.utils import is_metrics_enabled
-from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes
+from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes, _create_tool_span
 from agentops.semconv import SpanAttributes, LLMRequestTypeValues, MessageAttributes
 
 
@@ -192,30 +192,11 @@ def _finalize_stream(self) -> None:
         if self._finish_reason:
             self._span.set_attribute(MessageAttributes.COMPLETION_FINISH_REASON.format(i=0), self._finish_reason)
 
-        # Set tool calls
+        # Create tool spans for each tool call
         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"] is not None:
-                    self._span.set_attribute(
-                        MessageAttributes.COMPLETION_TOOL_CALL_ID.format(i=0, j=idx), tool_call["id"]
-                    )
-
-                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"] 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"] is not None:
-                    self._span.set_attribute(
-                        MessageAttributes.COMPLETION_TOOL_CALL_ARGUMENTS.format(i=0, j=idx),
-                        tool_call["function"]["arguments"],
-                    )
+                # Create a child span for this tool call
+                _create_tool_span(self._span, tool_call)
 
         # Set usage if available from the API
         if self._usage is not None:
@@ -374,7 +355,7 @@ def chat_completion_stream_wrapper(tracer, wrapped, instance, args, kwargs):
             return OpenaiStreamWrapper(response, span, kwargs)
         else:
             # Handle non-streaming response
-            response_attributes = handle_chat_attributes(kwargs=kwargs, return_value=response)
+            response_attributes = handle_chat_attributes(kwargs=kwargs, return_value=response, span=span)
 
             for key, value in response_attributes.items():
                 if key not in request_attributes:  # Avoid overwriting request attributes
@@ -439,7 +420,7 @@ async def async_chat_completion_stream_wrapper(tracer, wrapped, instance, args,
             return OpenAIAsyncStreamWrapper(response, span, kwargs)
         else:
             # Handle non-streaming response
-            response_attributes = handle_chat_attributes(kwargs=kwargs, return_value=response)
+            response_attributes = handle_chat_attributes(kwargs=kwargs, return_value=response, span=span)
 
             for key, value in response_attributes.items():
                 if key not in request_attributes:  # Avoid overwriting request attributes

agentops/instrumentation/providers/openai/wrappers/chat.py

@@ -8,28 +8,72 @@
 import logging
 from typing import Any, Dict, Optional, Tuple
 
+from opentelemetry.trace import Span
+
 from agentops.instrumentation.providers.openai.utils import is_openai_v1
 from agentops.instrumentation.providers.openai.wrappers.shared import (
     model_as_dict,
     should_send_prompts,
 )
 from agentops.instrumentation.common.attributes import AttributeMap
 from agentops.semconv import SpanAttributes, LLMRequestTypeValues
+from agentops.semconv.tool import ToolAttributes
+from agentops.semconv.span_kinds import AgentOpsSpanKindValues
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import SpanKind, Status, StatusCode, get_tracer
 
 logger = logging.getLogger(__name__)
 
 LLM_REQUEST_TYPE = LLMRequestTypeValues.CHAT
 
 
+def _create_tool_span(parent_span, tool_call_data):
+    """
+    Create a distinct span for each tool call.
+
+    Args:
+        parent_span: The parent LLM span
+        tool_call_data: The tool call data dictionary
+    """
+    # Get the tracer for this module
+    tracer = get_tracer(__name__)
+
+    # Create a child span for the tool call
+    with tracer.start_as_current_span(
+        name=f"tool_call.{tool_call_data['function']['name']}",
+        kind=SpanKind.INTERNAL,
+        context=context_api.set_value("current_span", parent_span),
+    ) as tool_span:
+        # Set the span kind to TOOL
+        tool_span.set_attribute("agentops.span.kind", AgentOpsSpanKindValues.TOOL)
+
+        # Set tool-specific attributes
+        tool_span.set_attribute(ToolAttributes.TOOL_NAME, tool_call_data["function"]["name"])
+        tool_span.set_attribute(ToolAttributes.TOOL_PARAMETERS, tool_call_data["function"]["arguments"])
+        tool_span.set_attribute("tool.call.id", tool_call_data["id"])
+        tool_span.set_attribute("tool.call.type", tool_call_data["type"])
+
+        # Set status to OK for successful tool call creation
+        tool_span.set_status(Status(StatusCode.OK))
+
+
 def handle_chat_attributes(
     args: Optional[Tuple] = None,
     kwargs: Optional[Dict] = None,
     return_value: Optional[Any] = None,
+    span: Optional[Span] = None,
 ) -> AttributeMap:
     """Extract attributes from chat completion calls.
 
     This function is designed to work with the common wrapper pattern,
     extracting attributes from the method arguments and return value.
+
+    Args:
+        args: Method arguments (not used in this implementation)
+        kwargs: Method keyword arguments
+        return_value: Method return value
+        span: The parent span for creating tool spans
     """
     attributes = {
         SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value,
@@ -191,12 +235,20 @@ def handle_chat_attributes(
                     # Tool calls
                     if "tool_calls" in message:
                         tool_calls = message["tool_calls"]
-                        if tool_calls:  # Check if tool_calls is not None
+                        if tool_calls and span is not None:
                             for i, tool_call in enumerate(tool_calls):
+                                # Convert tool_call to the format expected by _create_tool_span
                                 function = tool_call.get("function", {})
-                                attributes[f"{prefix}.tool_calls.{i}.id"] = tool_call.get("id")
-                                attributes[f"{prefix}.tool_calls.{i}.name"] = function.get("name")
-                                attributes[f"{prefix}.tool_calls.{i}.arguments"] = function.get("arguments")
+                                tool_call_data = {
+                                    "id": tool_call.get("id", ""),
+                                    "type": tool_call.get("type", "function"),
+                                    "function": {
+                                        "name": function.get("name", ""),
+                                        "arguments": function.get("arguments", ""),
+                                    },
+                                }
+                                # Create a child span for this tool call
+                                _create_tool_span(span, tool_call_data)
 
         # Prompt filter results
         if "prompt_filter_results" in response_dict: