monocle-apptrace 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/openai/_helper.py

@@ -3,33 +3,123 @@ This module provides utility functions for extracting system, user,
 and assistant messages from various input formats.
 """
 
+import json
 import logging
+from opentelemetry.context import get_value
 from monocle_apptrace.instrumentation.common.utils import (
     Option,
+    get_json_dumps,
     try_option,
     get_exception_message,
     get_parent_span,
     get_status_code,
 )
 from monocle_apptrace.instrumentation.common.span_handler import NonFrameworkSpanHandler, WORKFLOW_TYPE_MAP
+from monocle_apptrace.instrumentation.metamodel.finish_types import (
+    map_openai_finish_reason_to_finish_type,
+    OPENAI_FINISH_REASON_MAPPING
+)
+from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY, CHILD_ERROR_CODE, INFERENCE_AGENT_DELEGATION, INFERENCE_COMMUNICATION, INFERENCE_TOOL_CALL
 
 logger = logging.getLogger(__name__)
 
-
 def extract_messages(kwargs):
     """Extract system and user messages"""
     try:
         messages = []
         if 'instructions' in kwargs:
-            messages.append({'instructions': kwargs.get('instructions', {})})
+            messages.append({'system': kwargs.get('instructions', {})})
         if 'input' in kwargs:
-            messages.append({'input': kwargs.get('input', {})})
+            if isinstance(kwargs['input'], str):
+                messages.append({'user': kwargs.get('input', "")})
+            # [
+            #     {
+            #         "role": "developer",
+            #         "content": "Talk like a pirate."
+            #     },
+            #     {
+            #         "role": "user",
+            #         "content": "Are semicolons optional in JavaScript?"
+            #     }
+            # ]
+            if isinstance(kwargs['input'], list):
+                # kwargs['input']
+                # [
+                #     {
+                #         "content": "I need to book a flight from NYC to LAX and also book the Hilton hotel in Los Angeles. Also check the weather in Los Angeles.",
+                #         "role": "user"
+                #     },
+                #     {
+                #         "arguments": "{}",
+                #         "call_id": "call_dSljcToR2LWwqWibPt0qjeHD",
+                #         "name": "transfer_to_flight_agent",
+                #         "type": "function_call",
+                #         "id": "fc_689c30f96f708191aabb0ffd8098cdbd016ef325124ac05f",
+                #         "status": "completed"
+                #     },
+                #     {
+                #         "arguments": "{}",
+                #         "call_id": "call_z0MTZroziWDUd0fxVemGM5Pg",
+                #         "name": "transfer_to_hotel_agent",
+                #         "type": "function_call",
+                #         "id": "fc_689c30f99b808191a8743ff407fa8ee2016ef325124ac05f",
+                #         "status": "completed"
+                #     },
+                #     {
+                #         "arguments": "{\"city\":\"Los Angeles\"}",
+                #         "call_id": "call_rrdRSPv5vcB4pgl6P4W8U2bX",
+                #         "name": "get_weather_tool",
+                #         "type": "function_call",
+                #         "id": "fc_689c30f9b824819196d4ad9379d570f7016ef325124ac05f",
+                #         "status": "completed"
+                #     },
+                #     {
+                #         "call_id": "call_rrdRSPv5vcB4pgl6P4W8U2bX",
+                #         "output": "The weather in Los Angeles is sunny and 75.",
+                #         "type": "function_call_output"
+                #     },
+                #     {
+                #         "call_id": "call_z0MTZroziWDUd0fxVemGM5Pg",
+                #         "output": "Multiple handoffs detected, ignoring this one.",
+                #         "type": "function_call_output"
+                #     },
+                #     {
+                #         "call_id": "call_dSljcToR2LWwqWibPt0qjeHD",
+                #         "output": "{\"assistant\": \"Flight Agent\"}",
+                #         "type": "function_call_output"
+                #     }
+                # ]
+                for item in kwargs['input']:
+                    if isinstance(item, dict) and 'role' in item and 'content' in item:
+                        messages.append({item['role']: item['content']})
+                    elif isinstance(item, dict) and 'type' in item and item['type'] == 'function_call':
+                        messages.append({
+                            "tool_function": item.get("name", ""),
+                            "tool_arguments": item.get("arguments", ""),
+                            "call_id": item.get("call_id", "")
+                        })
+                    elif isinstance(item, dict) and 'type' in item and item['type'] == 'function_call_output':
+                        messages.append({
+                            "call_id": item.get("call_id", ""),
+                            "output": item.get("output", "")
+                        })
         if 'messages' in kwargs and len(kwargs['messages']) >0:
             for msg in kwargs['messages']:
                 if msg.get('content') and msg.get('role'):
                     messages.append({msg['role']: msg['content']})
+                elif msg.get('tool_calls') and msg.get('role'):
+                    try:
+                        tool_call_messages = []
+                        for tool_call in msg['tool_calls']:
+                            tool_call_messages.append(get_json_dumps({
+                                "tool_function": tool_call.function.name,
+                                "tool_arguments": tool_call.function.arguments,
+                            }))
+                        messages.append({msg['role']: tool_call_messages})
+                    except Exception as e:
+                        logger.warning("Warning: Error occurred while processing tool calls: %s", str(e))
 
-        return [str(message) for message in messages]
+        return [get_json_dumps(message) for message in messages]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
@@ -37,25 +127,62 @@ def extract_messages(kwargs):
 
 def extract_assistant_message(arguments):
     try:
+        messages = []
         status = get_status_code(arguments)
-        response: str = ""
-        if status == 'success':
+        if status == 'success' or status == 'completed':
             response = arguments["result"]
-            if hasattr(response,"output_text") and len(response.output_text):
-                return response.output_text
-            if response is not None and hasattr(response,"choices") and len(response.choices) >0:
-                if hasattr(response.choices[0],"message"):
-                    return response.choices[0].message.content
+            if hasattr(response, "tools") and isinstance(response.tools, list) and len(response.tools) > 0 and isinstance(response.tools[0], dict):
+                tools = []
+                for tool in response.tools:
+                    tools.append({
+                        "tool_id": tool.get("id", ""),
+                        "tool_name": tool.get("name", ""),
+                        "tool_arguments": tool.get("arguments", "")
+                    })
+                messages.append({"tools": tools})
+            if hasattr(response, "output") and isinstance(response.output, list) and len(response.output) > 0:
+                response_messages = []
+                role = "assistant"
+                for response_message in response.output:
+                    if(response_message.type == "function_call"):
+                        role = "tools"
+                        response_messages.append({
+                            "tool_id": response_message.call_id,
+                            "tool_name": response_message.name,
+                            "tool_arguments": response_message.arguments
+                        })
+                if len(response_messages) > 0:
+                    messages.append({role: response_messages})
+                    
+            if hasattr(response, "output_text") and len(response.output_text):
+                role = response.role if hasattr(response, "role") else "assistant"
+                messages.append({role: response.output_text})
+            if (
+                response is not None
+                and hasattr(response, "choices")
+                and len(response.choices) > 0
+            ):
+                if hasattr(response.choices[0], "message"):
+                    role = (
+                        response.choices[0].message.role
+                        if hasattr(response.choices[0].message, "role")
+                        else "assistant"
+                    )
+                    messages.append({role: response.choices[0].message.content})
+            return get_json_dumps(messages[0]) if messages else ""
         else:
             if arguments["exception"] is not None:
-                response = get_exception_message(arguments)
+                return get_exception_message(arguments)
             elif hasattr(arguments["result"], "error"):
-                response = arguments["result"].error
-        return response
+                return arguments["result"].error
+
     except (IndexError, AttributeError) as e:
-        logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
+        logger.warning(
+            "Warning: Error occurred in extract_assistant_message: %s", str(e)
+        )
         return None
 
+
 def extract_provider_name(instance):
     provider_url: Option[str] = try_option(getattr, instance._client.base_url, 'host')
     return provider_url.unwrap_or(None)
@@ -129,7 +256,7 @@ def get_inference_type(instance):
 
 class OpenAISpanHandler(NonFrameworkSpanHandler):
     def is_teams_span_in_progress(self) -> bool:
-        return self.is_framework_span_in_progess() and self.get_workflow_name_in_progress() == WORKFLOW_TYPE_MAP["teams.ai"]
+        return self.is_framework_span_in_progress() and self.get_workflow_name_in_progress() == WORKFLOW_TYPE_MAP["teams.ai"]
 
     # If openAI is being called by Teams AI SDK, then retain the metadata part of the span events
     def skip_processor(self, to_wrap, wrapped, instance, span, args, kwargs) -> list[str]:
@@ -144,3 +271,48 @@ class OpenAISpanHandler(NonFrameworkSpanHandler):
             return super().hydrate_events(to_wrap, wrapped, instance, args, kwargs, ret_result, span=parent_span, parent_span=None, ex=ex)
 
         return super().hydrate_events(to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=parent_span, ex=ex)
+
+    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span):
+        # TeamsAI doesn't capture the status and other metadata from underlying OpenAI SDK.
+        # Thus we save the OpenAI status code in the parent span and retrieve it here to preserve meaningful error codes.
+        if self.is_teams_span_in_progress() and ex is not None:
+            if len(span.events) > 1 and span.events[1].name == "data.output" and span.events[1].attributes.get("error_code") is not None:
+                parent_span.set_attribute(CHILD_ERROR_CODE, span.events[1].attributes.get("error_code"))
+        super().post_task_processing(to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span)
+
+def extract_finish_reason(arguments):
+    """Extract finish_reason from OpenAI response"""
+    try:
+        if arguments["exception"] is not None:
+            if hasattr(arguments["exception"], "code") and arguments["exception"].code in OPENAI_FINISH_REASON_MAPPING.keys():
+                return arguments["exception"].code
+        response = arguments["result"]
+
+        # Handle streaming responses
+        if hasattr(response, "finish_reason") and response.finish_reason:
+            return response.finish_reason
+
+        # Handle non-streaming responses
+        if response is not None and hasattr(response, "choices") and len(response.choices) > 0:
+            if hasattr(response.choices[0], "finish_reason"):
+                return response.choices[0].finish_reason
+    except (IndexError, AttributeError) as e:
+        logger.warning("Warning: Error occurred in extract_finish_reason: %s", str(e))
+        return None
+    return None
+
+def map_finish_reason_to_finish_type(finish_reason):
+    """Map OpenAI finish_reason to finish_type based on the possible errors mapping"""
+    return map_openai_finish_reason_to_finish_type(finish_reason)
+
+def agent_inference_type(arguments):
+    """Extract agent inference type from OpenAI response"""
+    message = json.loads(extract_assistant_message(arguments))
+    # message["tools"][0]["tool_name"]
+    if message and message.get("tools") and isinstance(message["tools"], list) and len(message["tools"]) > 0:
+        agent_prefix = get_value(AGENT_PREFIX_KEY)
+        tool_name = message["tools"][0].get("tool_name", "")
+        if tool_name and agent_prefix and tool_name.startswith(agent_prefix):
+            return INFERENCE_AGENT_DELEGATION
+        return INFERENCE_TOOL_CALL
+    return INFERENCE_COMMUNICATION

monocle_apptrace/instrumentation/metamodel/openai/entities/inference.py

@@ -6,117 +6,159 @@ from monocle_apptrace.instrumentation.metamodel.openai import (
     _helper,
 )
 from monocle_apptrace.instrumentation.common.utils import (
+    get_error_message,
     patch_instance_method,
     resolve_from_alias,
-    get_status,
-    get_exception_status_code,
-    get_status_code,
 )
 
 logger = logging.getLogger(__name__)
 
 
+def _process_stream_item(item, state):
+    """Process a single stream item and update state."""
+    try:
+        if (
+            hasattr(item, "type")
+            and isinstance(item.type, str)
+            and item.type.startswith("response.")
+        ):
+            if state["waiting_for_first_token"]:
+                state["waiting_for_first_token"] = False
+                state["first_token_time"] = time.time_ns()
+            if item.type == "response.output_text.delta":
+                state["accumulated_response"] += item.delta
+            if item.type == "response.completed":
+                state["stream_closed_time"] = time.time_ns()
+                if hasattr(item, "response") and hasattr(item.response, "usage"):
+                    state["token_usage"] = item.response.usage
+        elif (
+            hasattr(item, "choices")
+            and item.choices
+            and item.choices[0].delta
+            and item.choices[0].delta.content
+        ):
+            if hasattr(item.choices[0].delta, "role") and item.choices[0].delta.role:
+                state["role"] = item.choices[0].delta.role
+            if state["waiting_for_first_token"]:
+                state["waiting_for_first_token"] = False
+                state["first_token_time"] = time.time_ns()
+
+            state["accumulated_response"] += item.choices[0].delta.content
+        elif (
+            hasattr(item, "object")
+            and item.object == "chat.completion.chunk"
+            and item.usage
+        ):
+            # Handle the case where the response is a chunk
+            state["token_usage"] = item.usage
+            state["stream_closed_time"] = time.time_ns()
+            # Capture finish_reason from the chunk
+            if (
+                hasattr(item, "choices")
+                and item.choices
+                and len(item.choices) > 0
+                and hasattr(item.choices[0], "finish_reason")
+                and item.choices[0].finish_reason
+            ):
+                finish_reason = item.choices[0].finish_reason
+                state["finish_reason"] = finish_reason
+
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred while processing stream item: %s",
+            str(e),
+        )
+    finally:
+        state["accumulated_temp_list"].append(item)
+
+
+def _create_span_result(state, stream_start_time):
+    # extract tool calls from the accumulated_temp_list
+    # this can only be done when all the streaming is complete.
+    for item in state["accumulated_temp_list"]:
+        try:
+            if (
+                item.choices
+                and isinstance(item.choices, list)
+                and hasattr(item.choices[0], "delta")
+                and hasattr(item.choices[0].delta, "tool_calls")
+                and item.choices[0].delta.tool_calls
+                and item.choices[0].delta.tool_calls[0].id
+                and item.choices[0].delta.tool_calls[0].function
+            ):
+                state["tools"] = state.get("tools", [])
+                state["tools"].append(
+                    {
+                        "id": item.choices[0].delta.tool_calls[0].id,
+                        "name": item.choices[0].delta.tool_calls[0].function.name,
+                        "arguments": item.choices[0]
+                        .delta.tool_calls[0]
+                        .function.arguments,
+                    }
+                )
+            if (item.choices and item.choices[0].finish_reason):
+                state["finish_reason"] = item.choices[0].finish_reason
+        except Exception as e:
+            logger.warning(
+                "Warning: Error occurred while processing tool calls: %s",
+                str(e),
+            )
+
+    """Create the span result object."""
+    return SimpleNamespace(
+        type="stream",
+        timestamps={
+            "role": state["role"],
+            "data.input": int(stream_start_time),
+            "data.output": int(state["first_token_time"]),
+            "metadata": int(state["stream_closed_time"] or time.time_ns()),
+        },
+        output_text=state["accumulated_response"],
+        tools=state["tools"] if "tools" in state else None,
+        usage=state["token_usage"],
+        finish_reason=state["finish_reason"],
+    )
+
+
 def process_stream(to_wrap, response, span_processor):
-    waiting_for_first_token = True
     stream_start_time = time.time_ns()
-    first_token_time = stream_start_time
-    stream_closed_time = None
-    accumulated_response = ""
-    token_usage = None
-    accumulated_temp_list = []
+
+    # Shared state for both sync and async processing
+    state = {
+        "waiting_for_first_token": True,
+        "first_token_time": stream_start_time,
+        "stream_closed_time": None,
+        "accumulated_response": "",
+        "token_usage": None,
+        "accumulated_temp_list": [],
+        "finish_reason": None,
+        "role": "assistant",
+    }
 
     if to_wrap and hasattr(response, "__iter__"):
         original_iter = response.__iter__
 
         def new_iter(self):
-            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage
-
             for item in original_iter():
-                try:
-                    if (
-                        item.choices
-                        and item.choices[0].delta
-                        and item.choices[0].delta.content
-                    ):
-                        if waiting_for_first_token:
-                            waiting_for_first_token = False
-                            first_token_time = time.time_ns()
-
-                        accumulated_response += item.choices[0].delta.content
-                        # token_usage = item.usage
-                    elif item.object == "chat.completion.chunk" and item.usage:
-                        # Handle the case where the response is a chunk
-                        token_usage = item.usage
-                        stream_closed_time = time.time_ns()
-
-                except Exception as e:
-                    logger.warning(
-                        "Warning: Error occurred while processing item in new_iter: %s",
-                        str(e),
-                    )
-                finally:
-                    accumulated_temp_list.append(item)
-                    yield item
+                _process_stream_item(item, state)
+                yield item
 
             if span_processor:
-                ret_val = SimpleNamespace(
-                    type="stream",
-                    timestamps={
-                        "data.input": int(stream_start_time),
-                        "data.output": int(first_token_time),
-                        "metadata": int(stream_closed_time or time.time_ns()),
-                    },
-                    output_text=accumulated_response,
-                    usage=token_usage,
-                )
+                ret_val = _create_span_result(state, stream_start_time)
                 span_processor(ret_val)
 
         patch_instance_method(response, "__iter__", new_iter)
-        
+
     if to_wrap and hasattr(response, "__aiter__"):
         original_iter = response.__aiter__
 
         async def new_aiter(self):
-            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage
-
             async for item in original_iter():
-                try:
-                    if (
-                        item.choices
-                        and item.choices[0].delta
-                        and item.choices[0].delta.content
-                    ):
-                        if waiting_for_first_token:
-                            waiting_for_first_token = False
-                            first_token_time = time.time_ns()
-
-                        accumulated_response += item.choices[0].delta.content
-                        # token_usage = item.usage
-                    elif item.object == "chat.completion.chunk" and item.usage:
-                        # Handle the case where the response is a chunk
-                        token_usage = item.usage
-                        stream_closed_time = time.time_ns()
-
-                except Exception as e:
-                    logger.warning(
-                        "Warning: Error occurred while processing item in new_aiter: %s",
-                        str(e),
-                    )
-                finally:
-                    accumulated_temp_list.append(item)
-                    yield item
+                _process_stream_item(item, state)
+                yield item
 
             if span_processor:
-                ret_val = SimpleNamespace(
-                    type="stream",
-                    timestamps={
-                        "data.input": int(stream_start_time),
-                        "data.output": int(first_token_time),
-                        "metadata": int(stream_closed_time or time.time_ns()),
-                    },
-                    output_text=accumulated_response,
-                    usage=token_usage,
-                )
+                ret_val = _create_span_result(state, stream_start_time)
                 span_processor(ret_val)
 
         patch_instance_method(response, "__aiter__", new_aiter)
@@ -198,6 +240,10 @@ INFERENCE = {
         {
             "name": "data.output",
             "attributes": [
+                {
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments),
+                },
                 {
                     "_comment": "this is result from LLM",
                     "attribute": "response",
@@ -205,14 +251,6 @@ INFERENCE = {
                         arguments,
                     ),
                 },
-                {
-                    "attribute": "status",
-                    "accessor": lambda arguments: get_status(arguments)
-                },
-                {
-                    "attribute": "status_code",
-                    "accessor": lambda arguments: get_status_code(arguments)
-                }
             ],
         },
         {
@@ -223,6 +261,24 @@ INFERENCE = {
                     "accessor": lambda arguments: _helper.update_span_from_llm_response(
                         arguments["result"]
                     ),
+                },
+                {
+                    "_comment": "finish reason from OpenAI response",
+                    "attribute": "finish_reason",
+                    "accessor": lambda arguments: _helper.extract_finish_reason(
+                        arguments
+                    ),
+                },
+                {
+                    "_comment": "finish type mapped from finish reason",
+                    "attribute": "finish_type",
+                    "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(
+                        _helper.extract_finish_reason(arguments)
+                    ),
+                },
+                {
+                    "attribute": "inference_sub_type",
+                    "accessor": lambda arguments: _helper.agent_inference_type(arguments)
                 }
             ],
         },

monocle_apptrace/instrumentation/metamodel/openai/entities/retrieval.py

@@ -4,7 +4,7 @@ from monocle_apptrace.instrumentation.metamodel.openai import (
 from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
 
 RETRIEVAL = {
-    "type": "retrieval",
+    "type": "embedding",
     "attributes": [
         [
             {