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

monocle_apptrace/instrumentation/metamodel/litellm/_helper.py

@@ -0,0 +1,89 @@
+"""
+This module provides utility functions for extracting system, user,
+and assistant messages from various input formats.
+"""
+import json
+import logging
+from monocle_apptrace.instrumentation.common.utils import (
+    Option,
+    get_json_dumps,
+    try_option,
+    get_exception_message,
+    get_status_code,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def extract_messages(kwargs):
+    """Extract system and user messages"""
+    try:
+        messages = []
+        if 'messages' in kwargs and len(kwargs['messages']) > 0:
+            for msg in kwargs['messages']:
+                if msg.get('content') and msg.get('role'):
+                    messages.append({msg['role']: msg['content']})
+
+        return [get_json_dumps(message) for message in messages]
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
+        return []
+
+
+def extract_assistant_message(arguments):
+    try:
+        messages = []
+        status = get_status_code(arguments)
+        if status == 'success' or status == 'completed':
+            response = arguments["result"]
+            if (response is not None and hasattr(response, "choices") and len(response.choices) > 0):
+                if hasattr(response.choices[0], "message"):
+                    role = (
+                        response.choices[0].message.role
+                        if hasattr(response.choices[0].message, "role")
+                        else "assistant"
+                    )
+                    messages.append({role: response.choices[0].message.content})
+            return get_json_dumps(messages[0]) if messages else ""
+        else:
+            if arguments["exception"] is not None:
+                return get_exception_message(arguments)
+            elif hasattr(arguments["result"], "error"):
+                return arguments["result"].error
+
+    except (IndexError, AttributeError) as e:
+        logger.warning(
+            "Warning: Error occurred in extract_assistant_message: %s", str(e)
+        )
+        return None
+
+def extract_provider_name(url):
+    """Extract host from a URL string (e.g., https://api.openai.com/v1/ -> api.openai.com)"""
+    if not url:
+        return None
+    return url.split("//")[-1].split("/")[0]
+
+def resolve_from_alias(my_map, alias):
+    """Find a alias that is not none from list of aliases"""
+
+    for i in alias:
+        if i in my_map.keys():
+            return my_map[i]
+    return None
+
+
+def update_span_from_llm_response(response):
+    meta_dict = {}
+    token_usage = None
+    if response is not None:
+        if token_usage is None and hasattr(response, "usage") and response.usage is not None:
+            token_usage = response.usage
+        elif token_usage is None and hasattr(response, "response_metadata"):
+            token_usage = getattr(response.response_metadata, "token_usage", None) \
+                if hasattr(response.response_metadata, "token_usage") \
+                else response.response_metadata.get("token_usage", None)
+        if token_usage is not None:
+            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")})
+    return meta_dict

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

@@ -0,0 +1,108 @@
+from monocle_apptrace.instrumentation.metamodel.litellm import (
+    _helper,
+)
+from monocle_apptrace.instrumentation.common.utils import (
+    get_error_message,
+    resolve_from_alias,
+    get_llm_type,
+)
+INFERENCE = {
+    "type": "inference",
+    "attributes": [
+        [
+            {
+                "_comment": "provider type ,name , deployment , inference_endpoint",
+                "attribute": "type",
+                "accessor": lambda arguments: "inference."
+                + (get_llm_type(arguments['instance']) or 'generic')
+            },
+            {
+                "attribute": "provider_name",
+                "accessor": lambda arguments: _helper.extract_provider_name(
+                    resolve_from_alias(
+                        arguments["kwargs"],
+                        ["azure_endpoint", "api_base", "endpoint"],
+                    )
+                ),
+            },
+            {
+                "attribute": "deployment",
+                "accessor": lambda arguments: resolve_from_alias(
+                    arguments["kwargs"].__dict__,
+                    [
+                        "engine",
+                        "azure_deployment",
+                        "deployment_name",
+                        "deployment_id",
+                        "deployment",
+                    ],
+                ),
+            },
+            {
+                "attribute": "inference_endpoint",
+                "accessor": lambda arguments: resolve_from_alias(
+                    arguments["kwargs"],
+                    ["azure_endpoint", "api_base", "endpoint"],
+                )
+            },
+        ],
+        [
+            {
+                "_comment": "LLM Model",
+                "attribute": "name",
+                "accessor": lambda arguments: resolve_from_alias(
+                    arguments["kwargs"],
+                    ["model", "model_name", "endpoint_name", "deployment_name"],
+                ),
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: "model.llm."
+                + resolve_from_alias(
+                    arguments["kwargs"],
+                    ["model", "model_name", "endpoint_name", "deployment_name"],
+                ),
+            },
+        ],
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is instruction and user query to LLM",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.extract_messages(
+                        arguments["kwargs"]
+                    ),
+                }
+            ],
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+
+                {
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
+                },
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments),
+                }
+            ],
+        },
+        {
+            "name": "metadata",
+            "attributes": [
+                {
+                    "_comment": "this is metadata usage from LLM",
+                    "accessor": lambda arguments: _helper.update_span_from_llm_response(
+                        arguments["result"]
+                    ),
+                }
+            ],
+        },
+    ],
+}

monocle_apptrace/instrumentation/metamodel/litellm/methods.py

@@ -0,0 +1,19 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.metamodel.litellm.entities.inference import INFERENCE
+
+LITELLM_METHODS = [
+    {
+        "package": "litellm.llms.openai.openai",
+        "object": "OpenAIChatCompletion",
+        "method": "completion",
+        "wrapper_method": task_wrapper,
+        "output_processor": INFERENCE
+    },
+    {
+        "package": "litellm.llms.azure.azure",
+        "object": "AzureChatCompletion",
+        "method": "completion",
+        "wrapper_method": task_wrapper,
+        "output_processor": INFERENCE
+    }
+]

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 ""