monocle-apptrace 0.6.0__py3-none-any.whl → 0.6.6__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/adk/_helper.py

@@ -4,6 +4,7 @@ This module provides utility functions to extract various attributes from agent
 """
 
 from ast import arguments
+import json
 from typing import Any, Dict, Optional
 from monocle_apptrace.instrumentation.metamodel.finish_types import map_adk_finish_reason_to_finish_type
 
@@ -181,7 +182,7 @@ def extract_tool_input(arguments: Dict[str, Any]) -> Any:
     Returns:
         Any: The extracted input data
     """
-    return [str(arguments['kwargs'].get('args'))]
+    return json.dumps(arguments['kwargs'].get('args', {}))
 
 def extract_tool_response(result: Any) -> Any:
     """

monocle_apptrace/instrumentation/metamodel/agents/_helper.py

@@ -134,12 +134,12 @@ def extract_tool_input(arguments):
                 1
             ]  # Second argument is usually the JSON string with params
             if isinstance(tool_input, str):
-                return [tool_input]
+                return tool_input
             elif isinstance(tool_input, dict):
-                return [get_json_dumps(tool_input)]
+                return get_json_dumps(tool_input)
 
         # Fallback to all args
-        return [str(arg) for arg in arguments["args"]]
+        return str(arguments["args"])
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_tool_input: %s", str(e))
     return []

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

@@ -128,6 +128,7 @@ TOOLS = {
             {
                 "_comment": "tool type",
                 "attribute": "type",
+                "phase": "post_execution",
                 "accessor": lambda arguments: _helper.get_tool_type(arguments["span"]),
             },
             {
@@ -206,6 +207,7 @@ AGENT_DELEGATION = {
             {
                 "_comment": "name of the target agent",
                 "attribute": "to_agent",
+                "phase": "post_execution",
                 "accessor": lambda arguments: _helper.extract_handoff_target(arguments),
             },
         ]

monocle_apptrace/instrumentation/metamodel/aiohttp/_helper.py

@@ -67,7 +67,7 @@ def get_function_name(args) -> str:
 class aiohttpSpanHandler(SpanHandler):
 
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
-        return aiohttp_pre_tracing(args)
+        return aiohttp_pre_tracing(args), None
     
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         aiohttp_post_tracing(token)

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

@@ -6,6 +6,7 @@ from monocle_apptrace.instrumentation.common.utils import (get_error_message, re
 
 INFERENCE = {
     "type": SPAN_TYPES.INFERENCE,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "attributes": [
         [
             {
@@ -80,10 +81,6 @@ INFERENCE = {
                     "_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/azfunc/_helper.py

@@ -84,7 +84,7 @@ def azure_func_post_tracing(token):
 class azureSpanHandler(SpanHandler):
 
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
-        return azure_func_pre_tracing(kwargs)
+        return azure_func_pre_tracing(kwargs), None
     
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         azure_func_post_tracing(token)

monocle_apptrace/instrumentation/metamodel/botocore/_helper.py

@@ -7,6 +7,8 @@ import logging
 import json
 from io import BytesIO
 from functools import wraps
+
+from rfc3986 import urlparse
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
 from monocle_apptrace.instrumentation.common.utils import ( get_exception_message, get_json_dumps, get_status_code,)
 from monocle_apptrace.instrumentation.metamodel.finish_types import map_bedrock_finish_reason_to_finish_type
@@ -194,3 +196,6 @@ def extract_finish_reason(arguments):
 def map_finish_reason_to_finish_type(finish_reason):
     """Map Bedrock finish_reason/stopReason to finish_type."""
     return map_bedrock_finish_reason_to_finish_type(finish_reason)
+
+def extract_provider_name(instance):
+    return urlparse(instance.meta.endpoint_url).hostname

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

@@ -15,6 +15,10 @@ INFERENCE = {
             {
                 "attribute": "inference_endpoint",
                 "accessor": lambda arguments: arguments['instance'].meta.endpoint_url
+            },
+            {
+                "attribute": "provider_name",
+                "accessor": lambda arguments: _helper.extract_provider_name(arguments['instance'])
             }
         ],
         [

monocle_apptrace/instrumentation/metamodel/fastapi/_helper.py

@@ -77,12 +77,12 @@ class FastAPISpanHandler(SpanHandler):
         fastapi_pre_tracing(scope)
         return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
 
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         fastapi_post_tracing()
-        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value, token)
 
 class FastAPIResponseSpanHandler(SpanHandler):
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         try:
             ctx = get_current()
             if ctx is not None:
@@ -92,4 +92,4 @@ class FastAPIResponseSpanHandler(SpanHandler):
                                         return_value, parent_span=parent_span)
         except Exception as e:
             logger.info(f"Failed to propagate fastapi response: {e}")
-        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value, token)

monocle_apptrace/instrumentation/metamodel/fastapi/methods.py

@@ -3,10 +3,10 @@ from monocle_apptrace.instrumentation.metamodel.fastapi.entities.http import FAS
 
 FASTAPI_METHODS = [
     {
-        "package": "fastapi",
-        "object": "FastAPI",
-        "method": "__call__",
-        "wrapper_method": task_wrapper,
+        "package": "fastapi.routing",
+        "object": "APIRoute",
+        "method": "handle",
+        "wrapper_method": atask_wrapper,
         "span_name": "fastapi.request",
         "span_handler": "fastapi_handler",
         "output_processor": FASTAPI_HTTP_PROCESSOR,

monocle_apptrace/instrumentation/metamodel/flask/_helper.py

@@ -67,13 +67,13 @@ def flask_post_tracing(token):
 class FlaskSpanHandler(SpanHandler):
 
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
-        return flask_pre_tracing(args)
+        return flask_pre_tracing(args), None
     
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         flask_post_tracing(token)
 
 class FlaskResponseSpanHandler(SpanHandler):
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
         try:
             _parent_span_context = get_current()
             if _parent_span_context is not None:
@@ -82,4 +82,4 @@ class FlaskResponseSpanHandler(SpanHandler):
                     self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, return_value, parent_span=parent_span)
         except Exception as e:
             logger.info(f"Failed to propogate flask response: {e}")
-        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value, token)

monocle_apptrace/instrumentation/metamodel/hugging_face/_helper.py

@@ -95,7 +95,7 @@ def update_span_from_llm_response(result, include_token_counts=False):
         "total_tokens": getattr(result.usage, "total_tokens", 0),
     } if include_token_counts else {}
     # Add other metadata fields like finish_reason, etc.
-    return {**tokens, "inference_sub_type": "turn_end"}
+    return {**tokens}
 
 
 def get_exception_status_code(exc):

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

@@ -5,6 +5,7 @@ from monocle_apptrace.instrumentation.metamodel.hugging_face import _helper
 
 INFERENCE = {
     "type": SPAN_TYPES.INFERENCE,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "attributes": [
         [
             {
@@ -86,10 +87,6 @@ INFERENCE = {
                     "_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/lambdafunc/_helper.py

@@ -83,7 +83,7 @@ def lambda_func_post_tracing(token):
 
 class lambdaSpanHandler(SpanHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
-        return lambda_func_pre_tracing(kwargs)
+        return lambda_func_pre_tracing(kwargs), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value,token):
         lambda_func_post_tracing(token)

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

@@ -6,6 +6,7 @@ from monocle_apptrace.instrumentation.common.utils import get_error_message, res
 
 INFERENCE = {
     "type": SPAN_TYPES.INFERENCE_FRAMEWORK,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "attributes": [
         [
             {
@@ -79,10 +80,6 @@ INFERENCE = {
                     "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/langgraph/_helper.py

@@ -22,9 +22,19 @@ def agent_instructions(arguments):
     else:
         return arguments['kwargs']['agent'].instructions
 
+def extract_request_agent_input(arguments):
+    if arguments['kwargs'] is not None and 'input' in arguments['kwargs']:
+        history = arguments['kwargs']['input']['messages']
+        messages = []
+        for message in history:
+            if 'content' in message and 'role' in message and message['role'] == "user":  # Check if the message is a UserMessage
+                messages.append(message['content'])
+        return messages
+    return []
+
 def extract_agent_input(arguments):
-    if arguments['result'] is not None and 'messages' in arguments['result']:
-        history = arguments['result']['messages']
+    if arguments['args'] is not None and len(arguments['args']) > 0 and 'messages' in arguments['args'][0]:
+        history = arguments['args'][0]['messages']
         messages = []
         for message in history:
             if hasattr(message, 'content') and hasattr(message, 'type') and message.type == "human":  # Check if the message is a HumanMessage
@@ -58,6 +68,10 @@ def update_span_from_llm_response(response):
 def extract_tool_response(result):
     if result is not None and hasattr(result, 'content'):
         return result.content
+    if isinstance(result, str):
+        return result
+    if isinstance(result[0], str):
+        return result[0]
     return None
 
 def get_status(result):
@@ -66,11 +80,21 @@ def get_status(result):
     return None
 
 def extract_tool_input(arguments):
-    tool_input = arguments['args'][0]
-    if isinstance(tool_input, str):
-        return [tool_input]
+    if arguments['args'] and len(arguments['args']) > 0:
+        tool_input = arguments['args'][0]
     else:
-        return list(tool_input.values())
+        tool_input:dict = arguments['kwargs'].copy()
+        tool_input.pop('run_manager', None)  # remove run_manager if exists
+        tool_input.pop('config', None)  # remove config if exists
+    return str(tool_input)
+
+    # if isinstance(tool_input, str):
+    #     return [tool_input]
+    # elif isinstance(tool_input, dict):
+    #     # return array of key value pairs
+    #     return [f"'{k}': '{str(v)}'" for k, v in tool_input.items()]
+    # else:
+    #     return [str(tool_input)]
 
 def get_name(instance):
     return instance.name if hasattr(instance, 'name') else ""

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

@@ -73,7 +73,7 @@ AGENT_REQUEST = {
             {
                 "_comment": "this is Agent input",
                 "attribute": "input",
-                "accessor": lambda arguments: _helper.extract_agent_input(arguments)
+                "accessor": lambda arguments: _helper.extract_request_agent_input(arguments)
             }
           ]
         },

monocle_apptrace/instrumentation/metamodel/langgraph/langgraph_processor.py

@@ -1,3 +1,4 @@
+import logging
 from opentelemetry.context import set_value, attach, detach, get_value
 from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY, SCOPE_NAME
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
@@ -5,46 +6,114 @@ from monocle_apptrace.instrumentation.metamodel.langgraph._helper import (
    DELEGATION_NAME_PREFIX, get_name, is_root_agent_name, is_delegation_tool, LANGGRAPTH_AGENT_NAME_KEY
 )
 from monocle_apptrace.instrumentation.metamodel.langgraph.entities.inference import (
-    AGENT_DELEGATION, AGENT_REQUEST
+    AGENT_DELEGATION, AGENT_REQUEST, AGENT
 )
 from monocle_apptrace.instrumentation.common.scope_wrapper import start_scope, stop_scope
+from monocle_apptrace.instrumentation.common.utils import is_scope_set
+try:
+    from langgraph.errors import ParentCommand
+except ImportError:
+    ParentCommand = None
+
+logger = logging.getLogger(__name__)
+
+class ParentCommandFilterSpan:
+    """A wrapper for spans that filters out ParentCommand exceptions from being recorded."""
+    
+    def __init__(self, span):
+        self.span = span
+        self.original_record_exception = span.record_exception
+        
+    def record_exception(self, exception, attributes=None, timestamp=None, escaped=False):
+        """Filter out ParentCommand exceptions before recording them."""
+        try:
+            # Check if this is a ParentCommand exception
+            if ParentCommand is not None and isinstance(exception, ParentCommand):
+                logger.debug("Filtering out ParentCommand exception from span recording")
+                return  # Don't record ParentCommand exceptions
+            
+            # For all other exceptions, use the original record_exception method
+            return self.original_record_exception(exception, attributes, timestamp, escaped)
+        except Exception as e:
+            logger.debug(f"Error in ParentCommand filtering: {e}")
+            # If filtering fails, fall back to original behavior
+            return self.original_record_exception(exception, attributes, timestamp, escaped)
 
 class LanggraphAgentHandler(SpanHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         context = set_value(LANGGRAPTH_AGENT_NAME_KEY, get_name(instance))
         context = set_value(AGENT_PREFIX_KEY, DELEGATION_NAME_PREFIX, context)
         scope_name = AGENT_REQUEST.get("type")
-        if scope_name is not None and is_root_agent_name(instance) and get_value(scope_name, context) is None:
-            return start_scope(scope_name, scope_value=None, context=context)
+        if not is_scope_set(scope_name):
+            agent_request_wrapper = to_wrap.copy()
+            agent_request_wrapper["output_processor"] = AGENT_REQUEST
+#            return start_scope(scope_name, scope_value=None, context=context)
+            return attach(context), agent_request_wrapper
         else:
-            return attach(context)
+            return attach(context), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, result, token):
         if token is not None:
             detach(token)
 
+    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span):
+        """Apply ParentCommand filtering to the span before task execution."""
+        # Apply ParentCommand filtering to this span
+        self._apply_parent_command_filtering(span)
+        super().post_task_processing(to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span)
+
+    def _apply_parent_command_filtering(self, span):
+        """Apply ParentCommand exception filtering to a span."""
+        try:
+            if hasattr(span, 'record_exception'):
+                # Create a filtered wrapper and replace the record_exception method
+                filter_wrapper = ParentCommandFilterSpan(span)
+                span.record_exception = filter_wrapper.record_exception
+                logger.debug("Applied ParentCommand filtering to LangGraph agent span")
+        except Exception as e:
+            logger.debug(f"Failed to apply ParentCommand filtering: {e}")
+
     # In multi agent scenarios, the root agent is the one that orchestrates the other agents. LangGraph generates an extra root level invoke()
     # call on top of the supervisor agent invoke().
     # This span handler resets the parent invoke call as generic type to avoid duplicate attributes/events in supervisor span and this root span.
-    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
-        if is_root_agent_name(instance) and "parent.agent.span" in span.attributes:
-            agent_request_wrapper = to_wrap.copy()
-            agent_request_wrapper["output_processor"] = AGENT_REQUEST
-        else:
-            agent_request_wrapper = to_wrap
-            if hasattr(instance, 'name') and parent_span is not None and not SpanHandler.is_root_span(parent_span):
-                parent_span.set_attribute("parent.agent.span", True)
-        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex)
+
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> bool:
+        # Filter out ParentCommand exceptions as they are LangGraph control flow mechanisms, not actual errors
+        if ParentCommand is not None and isinstance(ex, ParentCommand):
+            ex = None  # Suppress the ParentCommand exception from being recorded
+
+        return super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)
 
 class LanggraphToolHandler(SpanHandler):
-    # LangGraph 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:
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         if is_delegation_tool(instance):
             agent_request_wrapper = to_wrap.copy()
             agent_request_wrapper["output_processor"] = AGENT_DELEGATION
         else:
-            agent_request_wrapper = to_wrap
+            agent_request_wrapper = None
+        return None, agent_request_wrapper
+
+    # LangGraph 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 post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span):
+        """Apply ParentCommand filtering to the span before task execution."""
+        # Apply ParentCommand filtering to this span
+        self._apply_parent_command_filtering(span)
+        super().post_task_processing(to_wrap, wrapped, instance, args, kwargs, result, ex, span, parent_span)
+
+    def _apply_parent_command_filtering(self, span):
+        """Apply ParentCommand exception filtering to a span."""
+        try:
+            if hasattr(span, 'record_exception'):
+                # Create a filtered wrapper and replace the record_exception method
+                filter_wrapper = ParentCommandFilterSpan(span)
+                span.record_exception = filter_wrapper.record_exception
+                logger.debug("Applied ParentCommand filtering to LangGraph tool span")
+        except Exception as e:
+            logger.debug(f"Failed to apply ParentCommand filtering: {e}")
 
-        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex)
-    
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> bool:
+        # Filter out ParentCommand exceptions as they are LangGraph control flow mechanisms, not actual errors
+        if ParentCommand is not None and isinstance(ex, ParentCommand):
+            ex = None  # Suppress the ParentCommand exception from being recorded
+        return super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)

monocle_apptrace/instrumentation/metamodel/langgraph/methods.py

@@ -24,17 +24,33 @@ LANGGRAPH_METHODS = [
         "output_processor": AGENT,
     },
     {
-        "package": "langchain_core.tools.base",
-        "object": "BaseTool",
-        "method": "run",
+        "package": "langchain_core.tools.simple",
+        "object": "Tool",
+        "method": "_run",
         "wrapper_method": task_wrapper,
         "span_handler": "langgraph_tool_handler",
         "output_processor": TOOLS,
     },
     {
-        "package": "langchain_core.tools.base",
-        "object": "BaseTool",
-        "method": "arun",
+        "package": "langchain_core.tools.structured",
+        "object": "StructuredTool",
+        "method": "_run",
+        "wrapper_method": task_wrapper,
+        "span_handler": "langgraph_tool_handler",
+        "output_processor": TOOLS,
+    },
+    {
+        "package": "langchain_core.tools.simple",
+        "object": "Tool",
+        "method": "_arun",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "langgraph_tool_handler",
+        "output_processor": TOOLS,
+    },
+    {
+        "package": "langchain_core.tools.structured",
+        "object": "StructuredTool",
+        "method": "_arun",
         "wrapper_method": atask_wrapper,
         "span_handler": "langgraph_tool_handler",
         "output_processor": TOOLS,

monocle_apptrace/instrumentation/metamodel/llamaindex/_helper.py

@@ -20,6 +20,11 @@ from monocle_apptrace.instrumentation.common.utils import (
 from monocle_apptrace.instrumentation.metamodel.finish_types import map_llamaindex_finish_reason_to_finish_type
 
 LLAMAINDEX_AGENT_NAME_KEY = "_active_agent_name"
+
+# Thread-local storage for current agent context
+import threading
+_thread_local = threading.local()
+
 logger = logging.getLogger(__name__)
 
 def get_status(result):
@@ -64,18 +69,18 @@ def get_tool_description(arguments):
         return ""
 
 def extract_tool_args(arguments):
-    tool_args = []
+    tool_args = {}
     if len(arguments['args']) > 1:
         for key, value in arguments['args'][2].items():
             # check if value is builtin type or a string
             if value is not None and isinstance(value, (str, int, float, bool)):
-                tool_args.append({key, value})
+                tool_args[key] = value
     else:
         for key, value in arguments['kwargs'].items():
             # check if value is builtin type or a string
             if value is not None and isinstance(value, (str, int, float, bool)):
-                tool_args.append({key, value})
-    return [get_json_dumps(tool_arg) for tool_arg in tool_args]
+                tool_args[key] = value
+    return get_json_dumps(tool_args)
 
 def extract_tool_response(response):
     if hasattr(response, 'raw_output'):
@@ -96,12 +101,27 @@ def get_agent_description(instance) -> str:
         return instance.description
     return ""
 
-def get_source_agent(parent_span:Span) -> str:
-    source_agent_name = parent_span.attributes.get(LLAMAINDEX_AGENT_NAME_KEY, "")
-    if source_agent_name == "" and parent_span.name.startswith("llama_index.core.agent.ReActAgent."):
-        # Fallback to the agent name from the parent span if not set
-        source_agent_name = "ReactAgent"
-    return source_agent_name
+def get_name(instance):
+    return instance.name if hasattr(instance, 'name') else ""
+
+def set_current_agent(agent_name: str):
+    """Set the current agent name in thread-local storage."""
+    _thread_local.current_agent = agent_name
+
+def get_current_agent() -> str:
+    """Get the current agent name from thread-local storage."""
+    return getattr(_thread_local, 'current_agent', '')
+
+def get_source_agent() -> str:
+    """Get the name of the agent that initiated the request."""
+    source_agent = get_value(LLAMAINDEX_AGENT_NAME_KEY)
+    if source_agent is not None and isinstance(source_agent,str) and source_agent != "":
+        return source_agent
+
+    source_agent = get_current_agent()
+    if source_agent:
+        return source_agent
+    return ""
 
 def get_target_agent(results) -> str:
     if hasattr(results, 'raw_input'):

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

@@ -111,7 +111,7 @@ TOOLS = {
               {
                 "_comment": "name of the agent",
                 "attribute": "name",
-                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+                "accessor": lambda arguments: _helper.get_source_agent()
               },          
               {
                 "_comment": "agent type",
@@ -157,12 +157,13 @@ AGENT_DELEGATION = {
               {
                 "_comment": "name of the agent",
                 "attribute": "from_agent",
-                "accessor": lambda arguments: _helper.get_source_agent(arguments['parent_span'])
+                "accessor": lambda arguments: _helper.get_source_agent()
               },
               {
                 "_comment": "name of the agent called",
                 "attribute": "to_agent",
-                "accessor": lambda arguments: _helper.get_target_agent(arguments['result'])
+                "accessor": lambda arguments: _helper.get_target_agent(arguments['result']),
+                "phase": "post_execution"
               }
         ]
       ]