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

monocle_apptrace/instrumentation/metamodel/llamaindex/methods.py

@@ -2,7 +2,7 @@ from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.inference import (
     INFERENCE,
 )
-from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import AGENT
+from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import AGENT, TOOLS, AGENT_REQUEST
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.retrieval import (
     RETRIEVAL,
 )
@@ -78,6 +78,14 @@ LLAMAINDEX_METHODS = [
         "wrapper_method": atask_wrapper,
         "output_processor": INFERENCE
     },
+    {
+        "package": "llama_index.core.agent.workflow.multi_agent_workflow",
+        "object": "AgentWorkflow",
+        "method": "run",
+        "span_handler": "llamaindex_agent_handler",
+        "wrapper_method": task_wrapper,
+        "output_processor": AGENT_REQUEST
+    },
     {
         "package": "llama_index.core.agent",
         "object": "ReActAgent",
@@ -85,6 +93,51 @@ LLAMAINDEX_METHODS = [
         "wrapper_method": task_wrapper,
         "output_processor": AGENT
     },
+    {
+        "package": "llama_index.core.agent",
+        "object": "ReActAgent",
+        "method": "achat",
+        "wrapper_method": atask_wrapper,
+        "output_processor": AGENT
+    },
+    {
+        "package": "llama_index.core.agent.workflow.function_agent",
+        "object": "FunctionAgent",
+        "method": "finalize",
+        "wrapper_method": atask_wrapper,
+        "output_processor": AGENT
+    },
+    {
+        "package": "llama_index.core.agent.workflow.function_agent",
+        "object": "FunctionAgent",
+        "method": "take_step",
+        "span_handler": "llamaindex_agent_handler",
+        "wrapper_method": atask_wrapper
+    },
+    {
+        "package": "llama_index.core.tools.function_tool",
+        "object": "FunctionTool",
+        "method": "call",
+        "span_handler": "llamaindex_single_agent_tool_handler",
+        "wrapper_method": task_wrapper,
+        "output_processor": TOOLS
+    },
+    {
+        "package": "llama_index.core.tools.function_tool",
+        "object": "FunctionTool",
+        "method": "acall",
+        "span_handler": "llamaindex_single_agent_tool_handler",
+        "wrapper_method": atask_wrapper,
+        "output_processor": TOOLS
+    },
+    {
+        "package": "llama_index.core.agent.workflow.multi_agent_workflow",
+        "object": "AgentWorkflow",
+        "method": "_call_tool",
+        "span_handler": "llamaindex_tool_handler",
+        "wrapper_method": atask_wrapper,
+        "output_processor": TOOLS
+    },
     {
         "package": "llama_index.llms.anthropic",
         "object": "Anthropic",
@@ -98,5 +151,19 @@ LLAMAINDEX_METHODS = [
         "method": "achat",
         "wrapper_method": atask_wrapper,
         "output_processor": INFERENCE
+    },
+    {
+        "package": "llama_index.llms.gemini",
+        "object": "Gemini",
+        "method": "chat",
+        "wrapper_method": task_wrapper,
+        "output_processor": INFERENCE
+    },
+    {
+        "package": "llama_index.llms.gemini",
+        "object": "Gemini",
+        "method": "achat",
+        "wrapper_method": atask_wrapper,
+        "output_processor": INFERENCE
     }
 ]

monocle_apptrace/instrumentation/metamodel/mcp/_helper.py

@@ -0,0 +1,118 @@
+from monocle_apptrace.instrumentation.common.utils import with_tracer_wrapper
+from opentelemetry.context import attach, set_value, get_value, detach
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+import logging
+import json
+
+logger = logging.getLogger(__name__)
+
+
+def log(arguments):
+    print(f"Arguments: {arguments}")
+
+
+def get_output_text(arguments):
+    # arguments["result"].content[0].text
+    if (
+        "result" in arguments
+        and hasattr(arguments["result"], "tools")
+        and isinstance(arguments["result"].tools, list)
+    ):
+        tools = []
+        for tool in arguments["result"].tools:
+            if hasattr(tool, "name"):
+                tools.append(tool.name)
+        return tools
+    if (
+        "result" in arguments
+        and hasattr(arguments["result"], "content")
+        and isinstance(arguments["result"].content, list)
+    ):
+        ret_val = []
+        for content in arguments["result"].content:
+            if hasattr(content, "text"):
+                ret_val.append(content.text)
+        return ret_val
+
+
+def get_name(arguments):
+    """Get the name of the tool from the instance."""
+
+    args = arguments["args"]
+    if (
+        args
+        and hasattr(args[0], "root")
+        and hasattr(args[0].root, "params")
+        and hasattr(args[0].root.params, "name")
+    ):
+        # If the first argument has a root with params and name, return that name
+        return args[0].root.params.name
+
+
+def get_type(arguments):
+    """Get the type of the tool from the instance."""
+    args = arguments["args"]
+    if args and hasattr(args[0], "root") and hasattr(args[0].root, "method"):
+        # If the first argument has a root with a method, return that method's name
+        return args[0].root.method
+
+
+def get_params_arguments(arguments):
+    """Get the params of the tool from the instance."""
+
+    args = arguments["args"]
+    if (
+        args
+        and hasattr(args[0], "root")
+        and hasattr(args[0].root, "params")
+        and hasattr(args[0].root.params, "arguments")
+    ):
+        # If the first argument has a root with params and arguments, return those arguments
+        try:
+            return json.dumps(args[0].root.params.arguments)
+        except (TypeError, ValueError) as e:
+            logger.error(f"Error serializing arguments: {e}")
+            return str(args[0].root.params.arguments)
+
+
+def get_url(arguments):
+    """Get the URL of the tool from the instance."""
+    url = get_value("mcp.url", None)
+
+    return url
+
+# this extracts the url from the langchain mcp adapter tools and attaches it to the context.
+@with_tracer_wrapper
+def langchain_mcp_wrapper(
+    tracer: any, handler: any, to_wrap, wrapped, instance, source_path, args, kwargs
+):
+    return_value = None
+    try:
+        return_value = wrapped(*args, **kwargs)
+        return return_value
+    finally:
+        if (
+            return_value
+            and hasattr(return_value, "coroutine")
+            and kwargs.get("connection", None)
+        ):
+            try:
+                # extract the URL from the connection and attach it to the context
+                url = kwargs.get("connection").get("url", None)
+                if url:
+                    # wrap coroutine methods and attach the URL to the context
+
+                    original_coroutine = return_value.coroutine
+
+                    async def wrapped_coroutine(*args1, **kwargs1):
+                        token = None
+                        try:
+                            token = attach(set_value("mcp.url", url))
+                            return await original_coroutine(*args1, **kwargs1)
+                        finally:
+                            detach(token)
+
+                    return_value.coroutine = wrapped_coroutine
+
+            except Exception as e:
+                pass

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

@@ -0,0 +1,48 @@
+from monocle_apptrace.instrumentation.metamodel.mcp import _helper
+
+TOOLS = {
+    "type": "agentic.mcp.invocation",
+    "attributes": [
+        [
+            {
+                "_comment": "name of the tool",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_name(arguments),
+            },
+            {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments: "mcp.server",
+            },
+            {
+                "_comment": "tool url",
+                "attribute": "url",
+                "accessor": lambda arguments: _helper.get_url(arguments),
+            },
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is Tool input",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.get_params_arguments(
+                        arguments
+                    ),
+                },
+            ],
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is Tool output",
+                    "attribute": "output",
+                    "accessor": lambda arguments: _helper.get_output_text(arguments)
+                },
+            ],
+        },
+    ],
+}

monocle_apptrace/instrumentation/metamodel/mcp/mcp_processor.py

@@ -0,0 +1,8 @@
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.metamodel.mcp._helper import (
+   get_name
+)
+
+class MCPAgentHandler(SpanHandler):
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        return get_name({"args": args, "kwargs": kwargs}) is None or args[0].root.method == "tools/list"

monocle_apptrace/instrumentation/metamodel/mcp/methods.py

@@ -0,0 +1,21 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper
+from monocle_apptrace.instrumentation.metamodel.mcp import _helper
+from monocle_apptrace.instrumentation.metamodel.mcp.entities.inference import TOOLS
+
+
+MCP_METHODS = [
+    {
+        "package": "mcp.shared.session",
+        "object": "BaseSession",
+        "method": "send_request",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "mcp_agent_handler",
+        "output_processor": TOOLS,
+    },
+    {
+        "package": "langchain_mcp_adapters.tools",
+        "object": "",
+        "method": "convert_mcp_tool_to_langchain_tool",
+        "wrapper_method": _helper.langchain_mcp_wrapper,
+    },
+]

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