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

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

@@ -1,24 +1,25 @@
 from monocle_apptrace.instrumentation.metamodel.langgraph import (
     _helper
 )
-INFERENCE = {
-      "type": "agent",
+
+AGENT = {
+      "type": "agentic.invocation",
       "attributes": [
         [
               {
                 "_comment": "agent type",
                 "attribute": "type",
-                "accessor": lambda arguments:'agent.oai'
+                "accessor": lambda arguments:'agent.langgraph'
               },
               {
                 "_comment": "name of the agent",
                 "attribute": "name",
-                "accessor": lambda arguments:arguments['instance'].name
+                "accessor": lambda arguments: _helper.get_agent_name(arguments['instance'])
               },
               {
-                "_comment": "agent tools",
-                "attribute": "tools",
-                "accessor": lambda arguments: _helper.tools(arguments['instance'])
+                "_comment": "agent description",
+                "attribute": "description",
+                "accessor": lambda arguments: _helper.get_agent_description(arguments['instance'])
               }
         ]
       ],
@@ -27,9 +28,9 @@ INFERENCE = {
           "name":"data.input",
           "attributes": [
             {
-                "_comment": "this is LLM input",
+                "_comment": "this is Agent input",
                 "attribute": "query",
-                "accessor": lambda arguments: _helper.extract_input(arguments)
+                "accessor": lambda arguments: _helper.extract_agent_input(arguments)
             }
           ]
         },
@@ -39,18 +40,124 @@ INFERENCE = {
             {
                 "_comment": "this is response from LLM",
                 "attribute": "response",
-                "accessor": lambda arguments: _helper.handle_openai_response(arguments['result'])
+                "accessor": lambda arguments: _helper.extract_agent_response(arguments['result'])
             }
           ]
-       },
-       {
-            "name": "metadata",
-            "attributes": [
-                {
-                    "_comment": "this is metadata usage from LLM",
-                    "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'])
-                }
-            ]
         }
       ]
-    }
+    }
+
+AGENT_REQUEST = {
+      "type": "agentic.request",
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.langgraph'
+              }
+        ],
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is Agent input",
+                "attribute": "input",
+                "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'])
+            }
+          ]
+        }
+      ]
+}
+
+TOOLS = {
+      "type": "agentic.tool.invocation",
+      "attributes": [
+        [
+              {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments:'tool.langgraph'
+              },
+              {
+                "_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.langgraph'
+              }
+        ]
+      ],
+      "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'])
+            }
+          ]
+        }
+      ]
+}
+
+AGENT_DELEGATION = {
+      "type": "agentic.delegation",
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.langgraph'
+              },
+              {
+                "_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/langgraph/langgraph_processor.py

@@ -0,0 +1,43 @@
+from opentelemetry.context import set_value, attach, detach
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.metamodel.langgraph._helper import (
+   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
+)
+
+class LanggraphAgentHandler(SpanHandler):
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        return attach(set_value(LANGGRAPTH_AGENT_NAME_KEY, get_name(instance)))
+
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, result, token):
+        if token is not None:
+            detach(token)
+
+    # 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)
+
+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:
+        if is_delegation_tool(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)
+    

monocle_apptrace/instrumentation/metamodel/langgraph/methods.py

@@ -1,14 +1,38 @@
-from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, atask_wrapper
 from monocle_apptrace.instrumentation.metamodel.langgraph.entities.inference import (
-    INFERENCE,
+    AGENT, TOOLS
 )
 LANGGRAPH_METHODS = [
     {
         "package": "langgraph.graph.state",
          "object": "CompiledStateGraph",
          "method": "invoke",
-         "span_name": "langgraph.graph.invoke",
          "wrapper_method": task_wrapper,
-         "output_processor": INFERENCE
-    }
+        "span_handler":"langgraph_agent_handler",
+         "output_processor": AGENT
+    },
+    {
+        "package": "langgraph.graph.state",
+         "object": "CompiledStateGraph",
+         "method": "ainvoke",
+         "wrapper_method": atask_wrapper,
+        "span_handler":"langgraph_agent_handler",
+         "output_processor": AGENT
+    },
+    {
+        "package": "langchain_core.tools.base",
+         "object": "BaseTool",
+         "method": "run",
+         "wrapper_method": task_wrapper,
+        "span_handler":"langgraph_tool_handler",
+         "output_processor": TOOLS
+    },
+    {
+        "package": "langchain_core.tools.base",
+         "object": "BaseTool",
+         "method": "arun",
+         "wrapper_method": atask_wrapper,
+        "span_handler":"langgraph_tool_handler",
+         "output_processor": TOOLS
+    },    
 ]

monocle_apptrace/instrumentation/metamodel/llamaindex/_helper.py

@@ -3,20 +3,29 @@ This module provides utility functions for extracting system, user,
 and assistant messages from various input formats.
 """
 
+from ast import arguments
 import logging
 from urllib.parse import urlparse
 from opentelemetry.sdk.trace import Span
+from opentelemetry.context import get_value
 from monocle_apptrace.instrumentation.common.utils import (
     Option,
+    get_json_dumps,
     get_keys_as_tuple,
     get_nested_value,
     try_option,
     get_exception_message,
     get_status_code,
 )
+from monocle_apptrace.instrumentation.metamodel.finish_types import map_llamaindex_finish_reason_to_finish_type
 
+LLAMAINDEX_AGENT_NAME_KEY = "_active_agent_name"
 logger = logging.getLogger(__name__)
 
+def get_status(result):
+    if result is not None and hasattr(result, 'status'):
+        return result.status
+    return None
 
 def extract_tools(instance):
     tools = []
@@ -34,6 +43,70 @@ def extract_tools(instance):
                 tools.append(tool_name)
     return tools
 
+def get_tool_name(args, instance):
+    if len(args) > 1:
+        if hasattr(args[1], 'metadata') and hasattr(args[1].metadata, 'name'):
+            return args[1].metadata.name
+        return ""
+    else:
+        if hasattr(instance, 'metadata') and hasattr(instance.metadata, 'name'):
+            return instance.metadata.name
+        return ""
+
+def get_tool_description(arguments):
+    if len(arguments['args']) > 1:
+        if hasattr(arguments['args'][1], 'metadata') and hasattr(arguments['args'][1].metadata, 'description'):
+            return arguments['args'][1].metadata.description
+        return ""
+    else:
+        if hasattr(arguments['instance'], 'metadata') and hasattr(arguments['instance'].metadata, 'description'):
+            return arguments['instance'].metadata.description
+        return ""
+
+def extract_tool_args(arguments):
+    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})
+    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]
+
+def extract_tool_response(response):
+    if hasattr(response, 'raw_output'):
+        return response.raw_output
+    return ""
+
+def is_delegation_tool(args, instance) -> bool:
+    return get_tool_name(args, instance) == "handoff"
+
+def get_agent_name(instance) -> str:
+    if hasattr(instance, 'name'):
+        return instance.name
+    else:
+        return instance.__class__.__name__
+
+def get_agent_description(instance) -> str:
+    if hasattr(instance, 'description'):
+        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_target_agent(results) -> str:
+    if hasattr(results, 'raw_input'):
+        return results.raw_input.get('kwargs', {}).get("to_agent", "")
+    return ""
 
 def extract_messages(args):
     """Extract system and user messages"""
@@ -50,38 +123,69 @@ def extract_messages(args):
         if isinstance(args, (list, tuple)) and args:
             for msg in args[0]:
                 process_message(msg)
+        elif args and isinstance(args, tuple):
+            messages.append(args[0])
         if isinstance(args, dict):
             for msg in args.get("messages", []):
                 process_message(msg)
-        if args and isinstance(args, tuple):
-            messages.append(args[0])
+        
 
-        return [str(message) for message in messages]
+        return [get_json_dumps(message) for message in messages]
 
     except Exception as e:
         logger.warning("Error in extract_messages: %s", str(e))
         return []
 
+def extract_agent_input(args):
+    if isinstance(args, (list, tuple)):
+        input_args = []
+        for arg in args:
+            if isinstance(arg, (str, dict)):
+                input_args.append(arg)
+            elif hasattr(arg, 'raw') and isinstance(arg.raw, str):
+                input_args.append(arg.raw)
+        return input_args
+    elif isinstance(args, str):
+        return [args]
+    return ""
+
+def extract_agent_response(arguments):
+    status = get_status_code(arguments)
+    if status == 'success':
+        if hasattr(arguments['result'], 'response'):
+            if hasattr(arguments['result'].response, 'content'):
+                return arguments['result'].response.content
+            return arguments['result'].response
+        return ""
+    else:
+        if arguments["exception"] is not None:
+            return get_exception_message(arguments)
+        elif hasattr(arguments['result'], "error"):
+            return arguments['result'].error
+
 def extract_assistant_message(arguments):
     status = get_status_code(arguments)
-    response: str = ""
+    messages = []
+    role = "assistant"
     if status == 'success':
         if isinstance(arguments['result'], str):
-            response = arguments['result']
+            messages.append({role: arguments['result']})
         if hasattr(arguments['result'], "content"):
-            response = arguments['result'].content
+            messages.append({role: arguments['result'].content})
         if hasattr(arguments['result'], "message") and hasattr(arguments['result'].message, "content"):
-            response = arguments['result'].message.content
+            role = getattr(arguments['result'].message, 'role', role)
+            if hasattr(role, 'value'):
+                role = role.value
+            messages.append({role: arguments['result'].message.content})
         if hasattr(arguments['result'],"response") and isinstance(arguments['result'].response, str):
-            response = arguments['result'].response
+            messages.append({role: arguments['result'].response})
     else:
         if arguments["exception"] is not None:
-            response = get_exception_message(arguments)
-        elif hasattr(response, "error"):
-            response = arguments['result'].error
-
-    return response
+            return get_exception_message(arguments)
+        elif hasattr(arguments['result'], "error"):
+            return arguments['result'].error
 
+    return get_json_dumps(messages[0]) if messages else ""
 
 def extract_query_from_content(content):
     try:
@@ -108,6 +212,8 @@ def extract_provider_name(instance):
         provider_url: Option[str]= try_option(getattr, instance, 'api_base').and_then(lambda url: urlparse(url).hostname)
     if hasattr(instance,'_client'):
         provider_url:Option[str] = try_option(getattr, instance._client.base_url,'host')
+    if hasattr(instance, 'model') and isinstance(instance.model, str) and 'gemini' in instance.model.lower():
+        provider_url: Option[str] = try_option(lambda: 'gemini.googleapis.com')
     return provider_url.unwrap_or(None)
 
 
@@ -117,6 +223,8 @@ def extract_inference_endpoint(instance):
             inference_endpoint: Option[str] = try_option(getattr, instance._client.sdk_configuration, 'server_url').map(str)
         if hasattr(instance._client,'base_url'):
             inference_endpoint: Option[str] = try_option(getattr, instance._client, 'base_url').map(str)
+    if hasattr(instance, 'model') and isinstance(instance.model, str) and 'gemini' in instance.model.lower():
+        inference_endpoint = try_option(lambda: f"https://generativelanguage.googleapis.com/v1beta/models/{instance.model}:generateContent")
     return inference_endpoint.unwrap_or(extract_provider_name(instance))
 
 
@@ -175,10 +283,113 @@ def update_span_from_llm_response(response, instance):
     if response is not None and hasattr(response, "raw"):
         if response.raw is not None:
             token_usage = response.raw.get("usage") if isinstance(response.raw, dict) else getattr(response.raw, "usage", None)
+            if token_usage is None:
+                token_usage = response.raw.get("usage_metadata") if isinstance(response.raw, dict) else getattr(response.raw,
+                                                                                                       "usage_metadata", None)
             if token_usage is not None:
                 temperature = instance.__dict__.get("temperature", None)
                 meta_dict.update({"temperature": temperature})
-                meta_dict.update({"completion_tokens": getattr(token_usage, "completion_tokens",None) or getattr(token_usage,"output_tokens",None)})
-                meta_dict.update({"prompt_tokens": getattr(token_usage, "prompt_tokens",None) or getattr(token_usage,"input_tokens",None)})
-                meta_dict.update({"total_tokens": getattr(token_usage, "total_tokens",None) or getattr(token_usage,"output_tokens",None)+getattr(token_usage,"input_tokens",None)})
+                meta_dict.update({"completion_tokens": getattr(token_usage, "completion_tokens",None) or getattr(token_usage,"output_tokens",None) or token_usage.get("candidates_token_count",None)})
+                meta_dict.update({"prompt_tokens": getattr(token_usage, "prompt_tokens",None) or getattr(token_usage,"input_tokens",None) or token_usage.get("prompt_token_count",None)})
+                total_tokens = getattr(token_usage, "total_tokens", None)
+                if total_tokens is not None:
+                    meta_dict.update({"total_tokens": total_tokens})
+                else:
+                    output_tokens = getattr(token_usage, "output_tokens", None)
+                    input_tokens = getattr(token_usage, "input_tokens", None)
+                    if output_tokens is not None and input_tokens is not None:
+                        meta_dict.update({"total_tokens": output_tokens + input_tokens})
+                    else:
+                        meta_dict.update({ "total_tokens": token_usage.get("total_token_count", None)})
+
     return meta_dict
+
+def extract_finish_reason(arguments):
+    """Extract finish_reason from LlamaIndex response."""
+    try:
+        # Handle exception cases first
+        if arguments.get("exception") is not None:
+            return "error"
+        
+        response = arguments.get("result")
+        if response is None:
+            return None
+            
+        # Check various possible locations for finish_reason in LlamaIndex responses
+        
+        # Direct finish_reason attribute
+        if hasattr(response, "finish_reason") and response.finish_reason:
+            return response.finish_reason
+            
+        # Check if response has raw attribute (common in LlamaIndex)
+        if hasattr(response, "raw") and response.raw:
+            raw_response = response.raw
+            if isinstance(raw_response, dict):
+                # Check for finish_reason in raw response
+                if "finish_reason" in raw_response:
+                    return raw_response["finish_reason"]
+                if "stop_reason" in raw_response:
+                    return raw_response["stop_reason"]
+                # Check for choices structure (OpenAI-style)
+                if "choices" in raw_response and raw_response["choices"]:
+                    choice = raw_response["choices"][0]
+                    if isinstance(choice, dict) and "finish_reason" in choice:
+                        return choice["finish_reason"]
+            elif hasattr(raw_response, "choices") and raw_response.choices:
+                # Handle object-style raw response
+                choice = raw_response.choices[0]
+                if hasattr(choice, "finish_reason"):
+                    return choice.finish_reason
+        
+        # Check for additional metadata
+        if hasattr(response, "additional_kwargs") and response.additional_kwargs:
+            kwargs = response.additional_kwargs
+            if isinstance(kwargs, dict):
+                for key in ["finish_reason", "stop_reason"]:
+                    if key in kwargs:
+                        return kwargs[key]
+        
+        # Check for response metadata
+        if hasattr(response, "response_metadata") and response.response_metadata:
+            metadata = response.response_metadata
+            if isinstance(metadata, dict):
+                for key in ["finish_reason", "stop_reason"]:
+                    if key in metadata:
+                        return metadata[key]
+        
+        # Check for source nodes or other LlamaIndex-specific attributes
+        if hasattr(response, "source_nodes") and response.source_nodes:
+            # If we have source nodes, it's likely a successful retrieval
+            return "stop"
+        
+        # If no specific finish reason found, infer from status
+        status_code = get_status_code(arguments)
+        if status_code == 'success':
+            return "stop"  # Default success finish reason
+        elif status_code == 'error':
+            return "error"
+            
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_finish_reason: %s", str(e))
+        return None
+    
+    return None
+
+
+def map_finish_reason_to_finish_type(finish_reason):
+    """Map LlamaIndex finish_reason to finish_type."""
+    return map_llamaindex_finish_reason_to_finish_type(finish_reason)
+
+def extract_agent_request_input(kwargs):
+    if "user_msg" in kwargs:
+        return kwargs["user_msg"]
+    return ""
+
+def extract_agent_request_output(arguments):
+    if hasattr(arguments['result'], 'response'):
+        if hasattr(arguments['result'].response, 'content'):
+            return arguments['result'].response.content
+        return arguments['result'].response
+    elif hasattr(arguments['result'], 'raw_output'):
+        return arguments['result'].raw_output
+    return ""