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

monocle_apptrace/instrumentation/metamodel/llamaindex/llamaindex_processor.py

@@ -2,7 +2,7 @@ from opentelemetry.context import attach, detach, get_current, get_value, set_va
 from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY
 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
+    is_delegation_tool, LLAMAINDEX_AGENT_NAME_KEY, get_agent_name, get_name, set_current_agent
 )
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import (
     AGENT_DELEGATION
@@ -13,14 +13,14 @@ 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:
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> 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)
+        return super().hydrate_span(agent_request_wrapper, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)
 
 # 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.
@@ -29,7 +29,10 @@ class LlamaIndexToolHandler(DelegationHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         cur_context = get_current()
         cur_context = set_value(TOOL_INVOCATION_STARTED, True, cur_context)
-        return attach(cur_context)
+        current_agent = get_value(LLAMAINDEX_AGENT_NAME_KEY)
+        if current_agent is not None:
+            cur_context = set_value(LLAMAINDEX_AGENT_NAME_KEY, current_agent, cur_context)
+        return attach(cur_context), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
         if token:
@@ -44,8 +47,13 @@ class LlamaIndexSingleAgenttToolHandlerWrapper(DelegationHandler):
 class LlamaIndexAgentHandler(SpanHandler):
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         cur_context = get_current()
+        agent_name = get_name(instance)
+
+        # Set both OpenTelemetry context and thread-local storage
+        set_current_agent(agent_name)
+        cur_context = set_value(LLAMAINDEX_AGENT_NAME_KEY, agent_name, cur_context)
         cur_context = set_value(AGENT_PREFIX_KEY, "handoff", cur_context)
-        return attach(cur_context)
+        return attach(cur_context), None
 
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
         if token:
@@ -53,8 +61,8 @@ 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)
+    def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None, is_post_exec:bool= False) -> bool:
+        retval = super().hydrate_span(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex, is_post_exec)
         if SpanHandler.is_root_span(parent_span):
             span.set_attribute(LLAMAINDEX_AGENT_NAME_KEY, "")
         else:

monocle_apptrace/instrumentation/metamodel/llamaindex/methods.py

@@ -89,17 +89,10 @@ LLAMAINDEX_METHODS = [
     {
         "package": "llama_index.core.agent",
         "object": "ReActAgent",
-        "method": "chat",
+        "method": "run",
         "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",

monocle_apptrace/instrumentation/metamodel/mcp/_helper.py

@@ -37,7 +37,7 @@ def get_output_text(arguments):
 
 def get_name(arguments):
     """Get the name of the tool from the instance."""
-    if 'parent_span' in arguments:
+    if 'parent_span' in arguments and arguments['parent_span']:
         arguments['parent_span'].set_attribute("is_mcp", True)
     args = arguments["args"]
     if (

monocle_apptrace/instrumentation/metamodel/mistral/_helper.py

@@ -146,7 +146,7 @@ def update_span_from_llm_response(result, include_token_counts=False):
         "total_tokens": getattr(result, "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 extract_finish_reason(arguments):

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

@@ -4,6 +4,7 @@ from monocle_apptrace.instrumentation.common.utils import get_error_message, res
 
 MISTRAL_INFERENCE = {
     "type": SPAN_TYPES.INFERENCE,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "attributes": [
         [
             {
@@ -83,10 +84,6 @@ MISTRAL_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/mistral/methods.py

@@ -43,14 +43,6 @@ MISTRAL_METHODS = [
         "wrapper_method": task_wrapper,
         "output_processor": MISTRAL_RETRIEVAL
     },
-    {
-        "package": "mistralai.embeddings",    # where Embeddings is defined
-        "object": "AsyncEmbeddings",          # async embeddings client
-        "method": "create",                   # async create
-        "span_handler": "non_framework_handler",
-        "wrapper_method": atask_wrapper,
-        "output_processor": MISTRAL_RETRIEVAL
-    }
 ]
 
  

monocle_apptrace/instrumentation/metamodel/openai/_helper.py

@@ -5,6 +5,7 @@ and assistant messages from various input formats.
 
 import json
 import logging
+from urllib.parse import urlparse
 from opentelemetry.context import get_value
 from monocle_apptrace.instrumentation.common.utils import (
     Option,
@@ -23,6 +24,12 @@ from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY,
 
 logger = logging.getLogger(__name__)
 
+# Mapping of URL substrings to provider names
+URL_MAP = {
+    "deepseek.com": "deepseek",
+    # add more providers here as needed
+}
+
 def extract_messages(kwargs):
     """Extract system and user messages"""
     try:
@@ -184,8 +191,19 @@ def extract_assistant_message(arguments):
 
 
 def extract_provider_name(instance):
+    # Try to get host from base_url if it's a parsed object
     provider_url: Option[str] = try_option(getattr, instance._client.base_url, 'host')
-    return provider_url.unwrap_or(None)
+    if provider_url.unwrap_or(None) is not None:
+        return provider_url.unwrap_or(None)
+
+    # If base_url is just a string (e.g., "https://api.deepseek.com")
+    base_url = getattr(instance._client, "base_url", None)
+    if isinstance(base_url, str):
+        parsed = urlparse(base_url)
+        if parsed.hostname:
+            return parsed.hostname
+
+    return None
 
 
 def extract_inference_endpoint(instance):
@@ -248,29 +266,49 @@ def extract_vector_output(vector_output):
     return ""
 
 def get_inference_type(instance):
+    # Check if it's Azure OpenAI first
     inference_type: Option[str] = try_option(getattr, instance._client, '_api_version')
     if inference_type.unwrap_or(None):
         return 'azure_openai'
-    else:
-        return 'openai'
+
+    # Check based on base_url using the mapping
+    base_url = getattr(instance, "base_url", None) or getattr(instance._client, "base_url", None)
+    
+    if base_url:
+        base_url_str = str(base_url).lower()
+        for key, name in URL_MAP.items():
+            if key in base_url_str:
+                return name
+
+    # fallback default
+    return "openai"
 
 class OpenAISpanHandler(NonFrameworkSpanHandler):
     def is_teams_span_in_progress(self) -> bool:
         return self.is_framework_span_in_progress() and self.get_workflow_name_in_progress() == WORKFLOW_TYPE_MAP["teams.ai"]
+    
+    def is_llamaindex_span_in_progress(self) -> bool:
+        return self.is_framework_span_in_progress() and self.get_workflow_name_in_progress() == WORKFLOW_TYPE_MAP["llama_index"]
 
-    # If openAI is being called by Teams AI SDK, then retain the metadata part of the span events
+    # If openAI is being called by Teams AI SDK or LlamaIndex, customize event processing
     def skip_processor(self, to_wrap, wrapped, instance, span, args, kwargs) -> list[str]:
         if self.is_teams_span_in_progress():
             return ["attributes", "events.data.input", "events.data.output"]
+        elif self.is_llamaindex_span_in_progress():
+            # For LlamaIndex, we want to keep all inference span attributes and events
+            return []
         else:
             return super().skip_processor(to_wrap, wrapped, instance, span, args, kwargs)
 
-    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=None, ex:Exception=None) -> bool:
+    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=None, ex:Exception=None, is_post_exec:bool=False) -> bool:
         # If openAI is being called by Teams AI SDK, then copy parent
         if self.is_teams_span_in_progress() and ex is None:
-            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=None, ex=ex, is_post_exec=is_post_exec)
+        # For LlamaIndex, process events normally on the inference span
+        elif self.is_llamaindex_span_in_progress():
+            return super().hydrate_events(to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=parent_span, ex=ex, is_post_exec=is_post_exec)
 
-        return super().hydrate_events(to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=parent_span, ex=ex)
+        return super().hydrate_events(to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=parent_span, ex=ex, is_post_exec=is_post_exec)
 
     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.
@@ -316,3 +354,5 @@ def agent_inference_type(arguments):
             return INFERENCE_AGENT_DELEGATION
         return INFERENCE_TOOL_CALL
     return INFERENCE_TURN_END
+
+

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

@@ -11,6 +11,7 @@ from monocle_apptrace.instrumentation.common.utils import (
     patch_instance_method,
     resolve_from_alias,
 )
+from monocle_apptrace.instrumentation.common.constants import PROVIDER_BASE_URLS
 
 logger = logging.getLogger(__name__)
 
@@ -120,6 +121,24 @@ def _create_span_result(state, stream_start_time):
         finish_reason=state["finish_reason"],
     )
 
+# Registry mapping client detection functions → entity_type
+CLIENT_ENTITY_MAP = {
+    "deepseek": "inference.deepseek",
+    # add more clients in future
+}
+
+def get_entity_type(response, helper=None):
+    for client_name, entity in CLIENT_ENTITY_MAP.items():
+        check_fn = globals().get(f"is_{client_name}_client")
+        if check_fn and check_fn(response):
+            return entity
+
+    # fallback to helper if available
+    if helper and hasattr(helper, "get_inference_type"):
+        return "inference." + helper.get_inference_type(response)
+
+    # default fallback
+    return "inference.openai"
 
 def process_stream(to_wrap, response, span_processor):
     stream_start_time = time.time_ns()
@@ -167,6 +186,7 @@ def process_stream(to_wrap, response, span_processor):
 
 INFERENCE = {
     "type": SPAN_TYPES.INFERENCE,
+    "subtype": lambda arguments: _helper.agent_inference_type(arguments),
     "is_auto_close": lambda kwargs: kwargs.get("stream", False) is False,
     "response_processor": process_stream,
     "attributes": [
@@ -276,10 +296,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/openai/methods.py

@@ -56,4 +56,4 @@ OPENAI_METHODS = [
         "output_processor": INFERENCE
     }
 
-]
+]

monocle_apptrace/instrumentation/metamodel/strands/_helper.py

@@ -0,0 +1,44 @@
+
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
+
+
+def get_agent_name(instance):
+    return instance.name
+
+def get_agent_description(instance):
+    return instance.description
+
+def extract_agent_input(arguments):
+    return arguments['args'][0]
+
+def extract_agent_response(result):
+    return result.message['content'][0]['text']
+
+def get_tool_type(arguments):
+    ## TODO: check for MCP type
+    return "tool.strands"
+
+def get_tool_name(arguments):
+    return arguments.get('args')[1][0]['name']
+
+def get_tool_description(arguments):
+    return arguments.get('args')[1][0]['description']
+
+def get_source_agent(arguments):
+    return arguments.get('args')[0].name
+
+def extract_tool_input(arguments):
+    return str(arguments.get('args')[1][0]['input'])
+
+def extract_tool_response(result):
+    return result.tool_result['content'][0]['text']
+
+def should_skip_delegation(arguments):
+    if arguments.get('parent_span') and arguments.get('parent_span').attributes.get("span.type") != SPAN_TYPES.AGENTIC_TOOL_INVOCATION:
+        return True
+    return False
+
+def should_skip_request(arguments):
+    if arguments.get('parent_span') and arguments.get('parent_span').attributes.get("span.type") in [SPAN_TYPES.AGENTIC_TOOL_INVOCATION, SPAN_TYPES.AGENTIC_REQUEST]:
+        return True
+    return False

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

@@ -0,0 +1,179 @@
+from monocle_apptrace.instrumentation.common.constants import AGENT_REQUEST_SPAN_NAME, SPAN_SUBTYPES, SPAN_TYPES
+from monocle_apptrace.instrumentation.metamodel.strands import (
+    _helper
+)
+from monocle_apptrace.instrumentation.common.utils import get_error_message
+
+AGENT = {
+      "type": SPAN_TYPES.AGENTIC_INVOCATION,
+      "subtype": SPAN_SUBTYPES.CONTENT_PROCESSING,
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.strands'
+              },
+              {
+                "_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'])
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Agent input",
+                "attribute": "input",
+                "accessor": lambda arguments: _helper.extract_agent_input(arguments)
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
+            },
+            {
+                "_comment": "this is response from LLM",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_agent_response(arguments['result'])
+            }
+          ]
+        }
+      ]
+    }
+
+AGENT_REQUEST = {
+      "type": SPAN_TYPES.AGENTIC_REQUEST,
+      "subtype": SPAN_SUBTYPES.PLANNING,
+      "should_skip": lambda arguments: _helper.should_skip_request(arguments),
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.strands'
+              }
+        ],
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Agent input",
+                "attribute": "input",
+                "accessor": lambda arguments: _helper.extract_agent_input(arguments)
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
+            },
+            {
+                "_comment": "this is response from LLM",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.extract_agent_response(arguments['result'])
+            }
+          ]
+        }
+      ]
+}
+
+TOOLS = {
+      "type": SPAN_TYPES.AGENTIC_TOOL_INVOCATION,
+      "subtype": SPAN_SUBTYPES.ROUTING,
+      "attributes": [
+        [
+              {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.get_tool_type(arguments['span'])
+              },
+              {
+                "_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()
+              },          
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.strands'
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Tool input",
+                "attribute": "input",
+                "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'])
+            }
+          ]
+        }
+      ]
+}
+
+AGENT_DELEGATION = {
+      "type": SPAN_TYPES.AGENTIC_DELEGATION,
+      "subtype": SPAN_SUBTYPES.ROUTING,
+      "should_skip": lambda arguments: _helper.should_skip_delegation(arguments),
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.strands'
+              },
+              {
+                "_comment": "name of the agent",
+                "attribute": "from_agent",
+                "accessor": lambda arguments: _helper.get_source_agent()
+              },
+              {
+                "_comment": "name of the agent called",
+                "attribute": "to_agent",
+                "accessor": lambda arguments: _helper.get_target_agent(arguments['instance'])
+              }
+        ]
+      ]
+}

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

@@ -0,0 +1,62 @@
+from monocle_apptrace.instrumentation.common.constants import AGENT_REQUEST_SPAN_NAME, SPAN_SUBTYPES, SPAN_TYPES
+from monocle_apptrace.instrumentation.metamodel.strands import (
+    _helper
+)
+from monocle_apptrace.instrumentation.common.utils import get_error_message
+TOOL = {
+      "type": SPAN_TYPES.AGENTIC_TOOL_INVOCATION,
+      "subtype": SPAN_SUBTYPES.CONTENT_GENERATION,
+      "attributes": [
+        [
+              {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.get_tool_type(arguments)
+              },
+              {
+                "_comment": "name of the tool",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_tool_name(arguments)
+              },
+              {
+                  "_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)
+              },          
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.stands'
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Tool input",
+                "attribute": "input",
+                "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/strands/methods.py

@@ -0,0 +1,20 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, atask_wrapper, atask_iter_wrapper
+from monocle_apptrace.instrumentation.metamodel.strands.entities.agent import AGENT, AGENT_REQUEST, AGENT_DELEGATION
+from monocle_apptrace.instrumentation.metamodel.strands.entities.tool import TOOL
+
+STRAND_METHODS = [
+    {
+        "package": "strands.agent.agent",
+        "object": "Agent",
+        "method": "__call__",
+        "wrapper_method": task_wrapper,
+        "output_processor_list": [AGENT_REQUEST, AGENT_DELEGATION, AGENT]
+    },
+    {
+        "package": "strands.tools.executors.concurrent",
+        "object": "ConcurrentToolExecutor",
+        "method": "_execute",
+        "wrapper_method": atask_iter_wrapper,
+        "output_processor": TOOL,
+    }
+]