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

monocle_apptrace/instrumentation/metamodel/llamaindex/entities/agent.py

@@ -3,21 +3,23 @@ from monocle_apptrace.instrumentation.metamodel.llamaindex import (
 )
 
 AGENT = {
-    "type": "agent",
+    "type": "agentic.invocation",
     "attributes": [
         [
             {
-                "_comment": "Agent name, type and Tools.",
-                "attribute": "name",
-                "accessor": lambda arguments: arguments['instance'].__class__.__name__
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.llamaindex'
             },
             {
-                "attribute": "type",
-                "accessor": lambda arguments: 'Agent.oai'
+                "_comment": "Agent name, type and Tools.",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_agent_name(arguments['instance'])
             },
             {
-                "attribute": "tools",
-                "accessor": lambda arguments: _helper.extract_tools(arguments['instance'])
+                "_comment": "agent description",
+                "attribute": "description",
+                "accessor": lambda arguments: _helper.get_agent_description(arguments['instance'])
             }
         ]
 
@@ -29,7 +31,7 @@ AGENT = {
              {
                  "_comment": "this is instruction and user query to LLM",
                  "attribute": "input",
-                 "accessor": lambda arguments: _helper.extract_messages(arguments['args'])
+                 "accessor": lambda arguments: _helper.extract_agent_input(arguments['args'])
              }
          ]
          },
@@ -39,9 +41,124 @@ AGENT = {
                 {
                     "_comment": "this is response from LLM",
                     "attribute": "response",
-                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                    "accessor": lambda arguments: _helper.extract_agent_response(arguments)
+                }
+            ]
+        }
+    ]
+} 
+
+AGENT_REQUEST = {
+    "type": "agentic.request",
+    "attributes": [
+        [
+              {
+                "_comment": "agent request type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.llamaindex'
+              }
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is request to LLM",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.extract_agent_request_input(arguments['kwargs'])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_agent_request_output(arguments)
                 }
             ]
         }
     ]
 }
+
+TOOLS = {
+      "type": "agentic.tool.invocation",
+      "attributes": [
+        [
+              {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments:'tool.llamaindex'
+              },
+              {
+                "_comment": "name of the tool",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_tool_name(arguments['args'], arguments['instance'])
+              },
+              {
+                  "_comment": "tool description",
+                  "attribute": "description",
+                  "accessor": lambda arguments: _helper.get_tool_description(arguments)
+              }
+        ],
+        [             
+              {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+              },          
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.llamaindex'
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Tool input",
+                "attribute": "Inputs",
+                "accessor": lambda arguments: _helper.extract_tool_args(arguments) 
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                "_comment": "this is response from Tool",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_tool_response(arguments['result'])
+             }
+          ]
+        }
+      ]
+}
+
+AGENT_DELEGATION = {
+    "type": "agentic.delegation",
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.llamaindex'
+              },
+              {
+                "_comment": "name of the agent",
+                "attribute": "from_agent",
+                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+              },
+              {
+                "_comment": "name of the agent called",
+                "attribute": "to_agent",
+                "accessor": lambda arguments: _helper.get_target_agent(arguments['result'])
+              }
+        ]
+      ]
+}

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

@@ -1,7 +1,7 @@
 from monocle_apptrace.instrumentation.metamodel.llamaindex import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import resolve_from_alias, get_llm_type, get_status, get_status_code
+from monocle_apptrace.instrumentation.common.utils import get_error_message, resolve_from_alias, get_llm_type, get_status, get_status_code
 
 INFERENCE = {
     "type": "inference.framework",
@@ -54,13 +54,8 @@ INFERENCE = {
             "name": "data.output",
             "attributes": [
                 {
-                    "_comment": "this is response from LLM",
-                    "attribute": "status",
-                    "accessor": lambda arguments: get_status(arguments)
-                },
-{
-                    "attribute": "status_code",
-                    "accessor": lambda arguments: get_status_code(arguments)
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
                 },
                 {
                     "attribute": "response",
@@ -74,6 +69,16 @@ INFERENCE = {
                 {
                     "_comment": "this is metadata usage from LLM",
                     "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'],arguments['instance'])
+                },
+                {
+                    "attribute": "finish_reason",
+                    "accessor": lambda arguments: _helper.extract_finish_reason(arguments)
+                },
+                {
+                    "attribute": "finish_type",
+                    "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(
+                        _helper.extract_finish_reason(arguments)
+                    )
                 }
             ]
         }

monocle_apptrace/instrumentation/metamodel/llamaindex/llamaindex_processor.py

@@ -0,0 +1,51 @@
+from opentelemetry.context import attach, detach, get_current, get_value, set_value, Context
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.metamodel.llamaindex._helper import (
+    is_delegation_tool, LLAMAINDEX_AGENT_NAME_KEY, get_agent_name
+)
+from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import (
+    AGENT_DELEGATION
+)
+
+TOOL_INVOCATION_STARTED:str = "llamaindex.tool_invocation_started"
+
+class DelegationHandler(SpanHandler):
+    # LlamaIndex uses an internal tool to initate delegation to other agents. The method is tool invoke() with tool name as `transfer_to_<agent_name>`.
+    # Hence we usea different output processor for tool invoke() to format the span as agentic.delegation.
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
+        if is_delegation_tool(args, instance):
+            agent_request_wrapper = to_wrap.copy()
+            agent_request_wrapper["output_processor"] = AGENT_DELEGATION
+        else:
+            agent_request_wrapper = to_wrap
+
+        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex)
+
+
+# There are two different APIs for tool calling FunctionTool.call() and AgentWorkflow.tool_call(). In case of single agent calling tool, only the FunctionTool.call() is used. In case of multi agent case,
+# the AgentWorkflow.tool_call() is used which inturn calls FunctionTool.call(). We can't entirely rely on the FunctionTool.call() to extract tool span details, especially the agent delegation details are not available there.
+# Hence we want to distinguish between single agent tool call and multi agent tool call. In case of multi agent tool call, we suppress the FunctionTool.call() span and use AgentWorkflow.tool_call() span to capture the tool call details.
+class LlamaIndexToolHandler(DelegationHandler):
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        return attach(set_value(TOOL_INVOCATION_STARTED, True))
+    
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
+        if token:
+            detach(token)
+
+class LlamaIndexSingleAgenttToolHandlerWrapper(DelegationHandler):
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs):
+        if get_value(TOOL_INVOCATION_STARTED) == True:
+            return True
+        return super().skip_span(to_wrap, wrapped, instance, args, kwargs)
+
+class LlamaIndexAgentHandler(SpanHandler):
+    # LlamaIndex uses direct OpenAI call for agent inferences. Given that the workflow type is set to llamaindex, the openAI inference does not record the input/output events.
+    # To avoid this, we set the workflow type to generic for agent inference spans so we can capture the prompts and responses.
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
+        retval = super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex)
+        if SpanHandler.is_root_span(parent_span):
+            span.set_attribute(LLAMAINDEX_AGENT_NAME_KEY, "")
+        else:
+            parent_span.set_attribute(LLAMAINDEX_AGENT_NAME_KEY, get_agent_name(instance))
+        return retval

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": atask_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,13 @@
+from opentelemetry.context import set_value, attach, detach
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.metamodel.mcp._helper import (
+   get_name
+
+)
+from monocle_apptrace.instrumentation.metamodel.langgraph.entities.inference import (
+    AGENT_DELEGATION, AGENT_REQUEST
+)
+
+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,
+    },
+]