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

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

@@ -0,0 +1,111 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_SUBTYPES, SPAN_TYPES
+from monocle_apptrace.instrumentation.metamodel.adk import _helper
+AGENT = {
+      "type": SPAN_TYPES.AGENTIC_INVOCATION,
+      "subtype": SPAN_SUBTYPES.ROUTING,
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.adk'
+              },
+              {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_agent_name(arguments['instance'])
+              },
+              {
+                "_comment": "agent description",
+                "attribute": "description",
+                "accessor": lambda arguments: _helper.get_agent_description(arguments['instance'])
+              },
+              {
+                "_comment": "delegating agent name",
+                "attribute": "from_agent",
+                "accessor": lambda arguments: _helper.get_delegating_agent(arguments)
+            }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Agent input",
+                "attribute": "query",
+                "accessor": lambda arguments: _helper.extract_agent_input(arguments)
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                "_comment": "this is response from LLM",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_agent_response(arguments['result'])
+            }
+          ]
+        }
+      ]
+    }
+
+REQUEST = {
+      "type": "agentic.request",
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.adk'
+              }
+        ],
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Agent input",
+                "attribute": "input",
+                "accessor": lambda arguments: _helper.extract_agent_request_input(arguments)
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                "_comment": "this is response from LLM",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_agent_response(arguments['result'])
+            }
+          ]
+        }
+      ]
+}
+
+DELEGATION = {
+      "type": "agentic.delegation",
+      "should_skip": lambda arguments: _helper.should_skip_delegation(arguments),
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.adk'
+              },
+              {
+                "_comment": "name of the agent",
+                "attribute": "from_agent",
+                "accessor": lambda arguments: _helper.get_delegating_agent(arguments)
+              },
+              {
+                "_comment": "name of the agent called",
+                "attribute": "to_agent",
+                "accessor": lambda arguments: _helper.get_agent_name(arguments['instance'])
+              }
+        ]
+      ]
+}

monocle_apptrace/instrumentation/metamodel/adk/entities/tool.py

@@ -0,0 +1,59 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_SUBTYPES, SPAN_TYPES
+from monocle_apptrace.instrumentation.metamodel.adk import _helper
+TOOL = {
+      "type": SPAN_TYPES.AGENTIC_TOOL_INVOCATION,
+      "subtype": SPAN_SUBTYPES.ROUTING,
+      "attributes": [
+        [
+              {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments:'tool.adk'
+              },
+              {
+                "_comment": "name of the tool",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_tool_name(arguments['instance'])
+              },
+              {
+                  "_comment": "tool description",
+                  "attribute": "description",
+                  "accessor": lambda arguments: _helper.get_tool_description(arguments['instance'])
+              }
+        ],
+        [
+              {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_source_agent(arguments)
+              },
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.adk'
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Tool input",
+                "attribute": "Inputs",
+                "accessor": lambda arguments: _helper.extract_tool_input(arguments) 
+            },
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                "_comment": "this is response from Tool",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_tool_response(arguments['result'])
+            }
+          ]
+        }
+      ]
+}

monocle_apptrace/instrumentation/metamodel/adk/methods.py

@@ -0,0 +1,31 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, atask_wrapper, atask_iter_wrapper
+from monocle_apptrace.instrumentation.metamodel.adk.entities.agent import (
+    AGENT, REQUEST, DELEGATION
+)
+from monocle_apptrace.instrumentation.metamodel.adk.entities.tool import (
+    TOOL
+)
+
+ADK_METHODS = [
+    {
+      "package": "google.adk.agents.base_agent",
+      "object": "BaseAgent",
+      "method": "run_async",
+      "wrapper_method": atask_iter_wrapper,
+      "output_processor_list": [DELEGATION, AGENT]
+    },
+    {
+      "package": "google.adk.tools.function_tool",
+      "object": "FunctionTool",
+      "method": "run_async",
+      "wrapper_method": atask_wrapper,
+      "output_processor": TOOL,
+    },
+    {
+      "package": "google.adk.runners",
+      "object": "Runner",
+      "method": "run_async",
+      "wrapper_method": atask_iter_wrapper,
+      "output_processor": REQUEST,
+    }
+]

monocle_apptrace/instrumentation/metamodel/agents/_helper.py

@@ -0,0 +1,225 @@
+from opentelemetry.context import get_value
+from monocle_apptrace.instrumentation.common.utils import (
+    resolve_from_alias,
+    get_json_dumps,
+)
+import logging
+
+logger = logging.getLogger(__name__)
+
+DELEGATION_NAME_PREFIX = "transfer_to_"
+ROOT_AGENT_NAME = "AgentsSDK"
+AGENTS_AGENT_NAME_KEY = "agent.openai_agents"
+
+
+def extract_agent_response(response):
+    """Extract the final output from an Agents SDK RunResult."""
+    try:
+        if response is not None and hasattr(response, "final_output"):
+            return str(response.final_output)
+        elif (
+            response is not None
+            and hasattr(response, "new_items")
+            and response.new_items
+        ):
+            # Try to extract from new_items if final_output is not available
+            last_item = response.new_items[-1]
+            if hasattr(last_item, "content"):
+                return str(last_item.content)
+            elif hasattr(last_item, "text"):
+                return str(last_item.text)
+        elif (
+            response is not None
+            and hasattr(response, "next_step")
+            and response.next_step
+            and hasattr(response.next_step, "output")
+            and response.next_step.output
+        ):
+            return str(response.next_step.output)
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_agent_response: %s", str(e))
+    return ""
+
+
+def extract_agent_input(arguments):
+    """Extract the input provided to the agent."""
+    try:
+        # For Runner.run and Runner.run_sync, the structure is:
+        # Runner.run(starting_agent, input, **kwargs)
+        # So args[0] = starting_agent, args[1] = input
+        if len(arguments["args"]) > 1:
+            input_data = arguments["args"][1]
+            if isinstance(input_data, str):
+                return input_data
+            elif isinstance(input_data, list):
+                # Handle list of input items
+                return get_json_dumps(input_data)
+
+        # Fallback to kwargs
+        if "original_input" in arguments["kwargs"]:
+            input_data = arguments["kwargs"]["original_input"]
+            if isinstance(input_data, str):
+                return input_data
+            elif isinstance(input_data, list):
+                return get_json_dumps(input_data)
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_agent_input: %s", str(e))
+    return None
+
+
+def get_agent_name(arguments) -> str:
+    """Get the name of an agent."""
+    instance = None
+    # For Runner methods, the agent is passed as an argument, not instance
+    if arguments["args"] and len(arguments["args"]) > 0:
+        instance = arguments["args"][0]
+    else:
+        instance = arguments["kwargs"].get("agent")
+    if instance is None:
+        return "Unknown Agent"
+    return get_name(instance)
+
+
+def get_agent_description(arguments) -> str:
+    """Get the description of an agent."""
+    instance = None
+    # For Runner methods, the agent is passed as an argument, not instance
+    if arguments["args"] and len(arguments["args"]) > 0:
+        instance = arguments["args"][0]
+    else:
+        instance = arguments["kwargs"].get("agent")
+    if instance is None:
+        return ""
+    return get_description(instance)
+
+
+def get_agent_instructions(arguments) -> str:
+    """Get the instructions of an agent."""
+    instance = None
+    # For Runner methods, the agent is passed as an argument, not instance
+    if arguments["args"] and len(arguments["args"]) > 0:
+        instance = arguments["args"][0]
+    else:
+        instance = arguments["kwargs"].get("agent")
+    if instance is None:
+        return ""
+    if hasattr(instance, "instructions"):
+        instructions = instance.instructions
+        if isinstance(instructions, str):
+            return instructions
+        elif callable(instructions):
+            # For dynamic instructions, we can't easily extract them without context
+            return "Dynamic instructions (function)"
+    return ""
+
+
+def extract_tool_response(result):
+    """Extract response from a tool execution."""
+    if result is not None:
+        if hasattr(result, "output"):
+            return str(result.output)
+        elif hasattr(result, "content"):
+            return str(result.content)
+        else:
+            return str(result)
+    return None
+
+
+def extract_tool_input(arguments):
+    """Extract input arguments passed to a tool."""
+    try:
+        # For function tools, the input is typically in args
+        if len(arguments["args"]) > 1:
+            tool_input = arguments["args"][
+                1
+            ]  # Second argument is usually the JSON string with params
+            if isinstance(tool_input, str):
+                return [tool_input]
+            elif isinstance(tool_input, dict):
+                return [get_json_dumps(tool_input)]
+
+        # Fallback to all args
+        return [str(arg) for arg in arguments["args"]]
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_tool_input: %s", str(e))
+    return []
+
+
+def get_name(instance):
+    """Get the name of an agent or tool instance."""
+    if hasattr(instance, "name"):
+        return str(instance.name)
+    elif hasattr(instance, "__name__"):
+        return str(instance.__name__)
+    return ""
+
+
+def get_runner_agent_name(instance) -> str:
+    """Get the name of an agent."""
+    return get_name(instance)
+
+def get_tool_type(span):
+    if (span.attributes.get("is_mcp", False)):
+        return "tool.mcp"
+    else:
+        return "tool.openai_agents"
+
+def get_tool_name(instance) -> str:
+    """Get the name of a tool."""
+    return get_name(instance)
+
+
+def is_root_agent_name(instance) -> bool:
+    """Check if this is the root agent."""
+    return get_name(instance) == ROOT_AGENT_NAME
+
+
+def get_source_agent() -> str:
+    """Get the name of the agent that initiated the request."""
+    from_agent = get_value(AGENTS_AGENT_NAME_KEY)
+    return from_agent if from_agent is not None else ""
+
+
+def get_description(instance) -> str:
+    """Get the description of an instance."""
+    if hasattr(instance, "description"):
+        return str(instance.description)
+    elif hasattr(instance, "handoff_description"):
+        return str(instance.handoff_description)
+    elif hasattr(instance, "__doc__") and instance.__doc__:
+        return str(instance.__doc__)
+    return ""
+
+
+def get_tool_description(instance) -> str:
+    """Get the description of a tool."""
+    return get_description(instance)
+
+
+def extract_handoff_target(arguments):
+    """Extract the target agent from a handoff operation."""
+    try:
+        # Check if this is a handoff by looking at the result
+        return arguments.get("result").name
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_handoff_target: %s", str(e))
+    return ""
+
+
+def update_span_from_agent_response(response):
+    """Update span with metadata from agent response."""
+    meta_dict = {}
+    try:
+        if response is not None and hasattr(response, "usage"):
+            usage = response.usage
+            if hasattr(usage, "completion_tokens"):
+                meta_dict.update({"completion_tokens": usage.completion_tokens})
+            if hasattr(usage, "prompt_tokens"):
+                meta_dict.update({"prompt_tokens": usage.prompt_tokens})
+            if hasattr(usage, "total_tokens"):
+                meta_dict.update({"total_tokens": usage.total_tokens})
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in update_span_from_agent_response: %s", str(e)
+        )
+    return meta_dict

monocle_apptrace/instrumentation/metamodel/agents/agents_processor.py

@@ -0,0 +1,174 @@
+import logging
+from types import SimpleNamespace
+from opentelemetry.context import attach, set_value, detach
+from opentelemetry.trace import Tracer
+from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY
+from monocle_apptrace.instrumentation.common.span_handler import (
+    SpanHandler as BaseSpanHandler,
+)
+from monocle_apptrace.instrumentation.common.utils import with_tracer_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper
+from monocle_apptrace.instrumentation.metamodel.agents._helper import (
+    AGENTS_AGENT_NAME_KEY,
+    DELEGATION_NAME_PREFIX,
+    get_runner_agent_name,
+)
+from monocle_apptrace.instrumentation.metamodel.agents.entities.inference import (
+    AGENT_DELEGATION,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@with_tracer_wrapper
+def constructor_wrapper(
+    tracer: Tracer,
+    handler: BaseSpanHandler,
+    to_wrap,
+    wrapped,
+    instance,
+    source_path,
+    args,
+    kwargs,
+):
+
+    original_func = kwargs.get("on_invoke_tool", None)
+    result = None
+    mcp_url = None
+    # kwargs.get("on_invoke_tool").args[0].params["url"]
+    if (
+        kwargs.get("on_invoke_tool")
+        and hasattr(kwargs.get("on_invoke_tool"), "args")
+        and len(kwargs.get("on_invoke_tool").args) > 0
+        and hasattr(kwargs.get("on_invoke_tool").args[0], "params")
+    ):
+        mcp_url = kwargs.get("on_invoke_tool").args[0].params.get("url", None)
+        if mcp_url:
+            logger.debug(f"Using MCP URL: {mcp_url}")
+    tool_instance = SimpleNamespace(
+        name=kwargs.get("name", "unknown_tool"),
+        description=kwargs.get("description", "No description provided"),
+    )
+    if original_func and not getattr(original_func, "_monocle_wrapped", False):
+        # Now wrap the function with our instrumentation
+
+        async def wrapped_func(*func_args, **func_kwargs):
+            token = None
+            try:
+                if mcp_url:
+                    token = attach(set_value("mcp.url", mcp_url))
+                # Use the handler to create spans when the decorated function is called
+                return await atask_wrapper(
+                    tracer=tracer, handler=handler, to_wrap=to_wrap
+                )(
+                    wrapped=original_func,
+                    instance=tool_instance,
+                    source_path=source_path,
+                    args=func_args,
+                    kwargs=func_kwargs,
+                )
+            finally:
+                if token:
+                    detach(token)
+
+        kwargs["on_invoke_tool"] = wrapped_func
+        # Preserve function metadata
+        wrapped_func.__name__ = getattr(wrapped, "__name__", "unknown_tool")
+        wrapped_func.__doc__ = getattr(wrapped, "__doc__", "")
+        # mark function as wrapped
+        setattr(wrapped_func, "_monocle_wrapped", True)
+
+    result = wrapped(*args, **kwargs)
+    return result
+
+
+@with_tracer_wrapper
+def handoff_constructor_wrapper(
+    tracer: Tracer,
+    handler: BaseSpanHandler,
+    to_wrap,
+    wrapped,
+    instance,
+    source_path,
+    args,
+    kwargs,
+):
+
+    original_func = kwargs.get("on_invoke_handoff", None)
+    result = None
+    tool_instance = SimpleNamespace(
+        name=kwargs.get("name", "unknown_handoff"),
+        description=kwargs.get("description", "No description provided"),
+    )
+    if original_func and not getattr(original_func, "_monocle_wrapped", False):
+        # Now wrap the function with our instrumentation
+        async def wrapped_func(*func_args, **func_kwargs):
+            # Use the handler to create spans when the decorated function is called
+            return await atask_wrapper(tracer=tracer, handler=handler, to_wrap=to_wrap)(
+                wrapped=original_func,
+                instance=tool_instance,
+                source_path=source_path,
+                args=func_args,
+                kwargs=func_kwargs,
+            )
+
+        kwargs["on_invoke_handoff"] = wrapped_func
+        # Preserve function metadata
+        wrapped_func.__name__ = getattr(wrapped, "__name__", "unknown_handoff")
+        wrapped_func.__doc__ = getattr(wrapped, "__doc__", "")
+        # mark function as wrapped
+        setattr(wrapped_func, "_monocle_wrapped", True)
+
+    result = wrapped(*args, **kwargs)
+    return result
+
+
+class AgentsSpanHandler(BaseSpanHandler):
+    """Span handler for OpenAI Agents SDK."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.agent_context_token = None
+
+    def set_agent_context(self, to_wrap, wrapped, instance, args, kwargs):
+        """Set the agent context for tracking across calls."""
+        try:
+            # For Runner.run, the agent is the first argument
+            if len(args) > 0:
+                agent = args[0]
+                agent_name = get_runner_agent_name(agent)
+                if agent_name:
+                    self.agent_context_token = attach(
+                        set_value(AGENTS_AGENT_NAME_KEY, agent_name)
+                    )
+        except Exception as e:
+            logger.warning("Warning: Error setting agent context: %s", str(e))
+
+    def clear_agent_context(self):
+        """Clear the agent context."""
+        if self.agent_context_token:
+            detach(self.agent_context_token)
+            self.agent_context_token = None
+
+    def pre_task_processing(
+        self, to_wrap, wrapped, instance, args, kwargs, *args1, **kwargs1
+    ):
+        """Pre-processing for agent tasks."""
+        self.set_agent_context(to_wrap, wrapped, instance, args, kwargs)
+        context = set_value(AGENT_PREFIX_KEY, DELEGATION_NAME_PREFIX)
+        attach(context)
+        return super().pre_task_processing(
+            to_wrap, wrapped, instance, args, kwargs, *args1, **kwargs1
+        )
+
+    def post_task_processing(
+        self, to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span
+    ):
+        """Post-processing for agent tasks."""
+        self.clear_agent_context()
+        context = set_value(AGENT_PREFIX_KEY, None)
+        attach(context)
+        return super().post_task_processing(
+            to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span
+        )
+