fix tool calls

Author: dot-agi

agentops/instrumentation/providers/openai/stream_wrapper.py

@@ -9,13 +9,13 @@
 from typing import Any, AsyncIterator, Iterator
 
 from opentelemetry import context as context_api
-from opentelemetry.trace import Span, SpanKind, Status, StatusCode, set_span_in_context
+from opentelemetry.trace import Span, SpanKind, Status, StatusCode, set_span_in_context, get_tracer
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 
 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
 
 
@@ -30,17 +30,19 @@ class OpenaiStreamWrapper:
     - Chunk statistics
     """
 
-    def __init__(self, stream: Any, span: Span, request_kwargs: dict):
+    def __init__(self, stream: Any, span: Span, request_kwargs: dict, tracer=None):
         """Initialize the stream wrapper.
 
         Args:
             stream: The original OpenAI stream object
             span: The OpenTelemetry span for tracking
             request_kwargs: Original request parameters for context
+            tracer: The OpenTelemetry tracer for creating child spans
         """
         self._stream = stream
         self._span = span
         self._request_kwargs = request_kwargs
+        self._tracer = tracer
         self._start_time = time.time()
         self._first_token_time = None
         self._chunk_count = 0
@@ -192,30 +194,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
-        if len(self._tool_calls) > 0:
+        # Create tool spans for each tool call
+        if len(self._tool_calls) > 0 and self._tracer is not None:
             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, self._tracer)
 
         # Set usage if available from the API
         if self._usage is not None:
@@ -254,17 +237,19 @@ def _finalize_stream(self) -> None:
 class OpenAIAsyncStreamWrapper:
     """Async wrapper for OpenAI Chat Completions streaming responses."""
 
-    def __init__(self, stream: Any, span: Span, request_kwargs: dict):
+    def __init__(self, stream: Any, span: Span, request_kwargs: dict, tracer=None):
         """Initialize the async stream wrapper.
 
         Args:
             stream: The original OpenAI async stream object
             span: The OpenTelemetry span for tracking
             request_kwargs: Original request parameters for context
+            tracer: The OpenTelemetry tracer for creating child spans
         """
         self._stream = stream
         self._span = span
         self._request_kwargs = request_kwargs
+        self._tracer = tracer
         self._start_time = time.time()
         self._first_token_time = None
         self._chunk_count = 0
@@ -371,10 +356,10 @@ def chat_completion_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         if is_streaming:
             # Wrap the stream
             context_api.detach(token)
-            return OpenaiStreamWrapper(response, span, kwargs)
+            return OpenaiStreamWrapper(response, span, kwargs, tracer)
         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, tracer=tracer)
 
             for key, value in response_attributes.items():
                 if key not in request_attributes:  # Avoid overwriting request attributes
@@ -436,10 +421,10 @@ async def async_chat_completion_stream_wrapper(tracer, wrapped, instance, args,
         if is_streaming:
             # Wrap the stream
             context_api.detach(token)
-            return OpenAIAsyncStreamWrapper(response, span, kwargs)
+            return OpenAIAsyncStreamWrapper(response, span, kwargs, tracer)
         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, tracer=tracer)
 
             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
 
 logger = logging.getLogger(__name__)
 
 LLM_REQUEST_TYPE = LLMRequestTypeValues.CHAT
 
 
+def _create_tool_span(parent_span, tool_call_data, tracer):
+    """
+    Create a distinct span for each tool call.
+    
+    Args:
+        parent_span: The parent LLM span
+        tool_call_data: The tool call data dictionary
+        tracer: The OpenTelemetry tracer instance
+    """
+    # 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,
+    tracer: Optional[Any] = 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
+        tracer: The OpenTelemetry tracer for creating child 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 and tracer 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, tracer)
 
         # Prompt filter results
         if "prompt_filter_results" in response_dict: