Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 6 additions & 25 deletions agentops/instrumentation/providers/openai/stream_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
60 changes: 56 additions & 4 deletions agentops/instrumentation/providers/openai/wrappers/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
Loading