monocle-apptrace 0.3.0b4__py3-none-any.whl → 0.3.0b6__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/langgraph/_helper.py

@@ -0,0 +1,48 @@
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+import logging
+logger = logging.getLogger(__name__)
+
+def handle_openai_response(response):
+    try:
+        if 'messages' in response:
+            output = response["messages"][-1]
+            return str(output.content)
+    except Exception as e:
+        logger.warning("Warning: Error occurred in handle_openai_response: %s", str(e))
+        return ""
+
+def agent_instructions(arguments):
+    if callable(arguments['kwargs']['agent'].instructions):
+        return arguments['kwargs']['agent'].instructions(arguments['kwargs']['context_variables'])
+    else:
+        return arguments['kwargs']['agent'].instructions
+
+def extract_input(arguments):
+    history = arguments['result']['messages']
+    for message in history:
+        if hasattr(message, 'content') and hasattr(message, 'type') and message.type == "human":  # Check if the message is a HumanMessage
+            return message.content
+
+def get_inference_endpoint(arguments):
+    inference_endpoint = resolve_from_alias(arguments['instance'].client.__dict__, ['azure_endpoint', 'api_base', '_base_url'])
+    return str(inference_endpoint)
+
+def tools(instance):
+    if hasattr(instance,'nodes') and ('tools' in instance.nodes):
+        tools= instance.nodes['tools']
+        if hasattr(tools,'bound') and hasattr(tools.bound,'tools_by_name'):
+            return list(tools.bound.tools_by_name.keys())
+
+
+def update_span_from_llm_response(response):
+    meta_dict = {}
+    token_usage = None
+    if response is not None and "messages" in response:
+        token = response["messages"][-1]
+        if token.response_metadata is not None:
+            token_usage = token.response_metadata["token_usage"]
+        if token_usage is not None:
+            meta_dict.update({"completion_tokens": token_usage.get('completion_tokens')})
+            meta_dict.update({"prompt_tokens": token_usage.get('prompt_tokens')})
+            meta_dict.update({"total_tokens": token_usage.get('total_tokens')})
+    return meta_dict

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

@@ -0,0 +1,56 @@
+from monocle_apptrace.instrumentation.metamodel.langgraph import (
+    _helper
+)
+INFERENCE = {
+      "type": "agent",
+      "attributes": [
+        [
+              {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments:'agent.oai'
+              },
+              {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments:arguments['instance'].name
+              },
+              {
+                "_comment": "agent tools",
+                "attribute": "tools",
+                "accessor": lambda arguments: _helper.tools(arguments['instance'])
+              }
+        ]
+      ],
+      "events": [
+        {
+          "name":"data.input",
+          "attributes": [
+            {
+                "_comment": "this is LLM input",
+                "attribute": "query",
+                "accessor": lambda arguments: _helper.extract_input(arguments)
+            }
+          ]
+        },
+        {
+          "name":"data.output",
+          "attributes": [
+            {
+                "_comment": "this is response from LLM",
+                "attribute": "response",
+                "accessor": lambda arguments: _helper.handle_openai_response(arguments['result'])
+            }
+          ]
+       },
+       {
+            "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/langgraph/methods.py

@@ -0,0 +1,14 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.metamodel.langgraph.entities.inference import (
+    INFERENCE,
+)
+LANGGRAPH_METHODS = [
+    {
+        "package": "langgraph.graph.state",
+         "object": "CompiledStateGraph",
+         "method": "invoke",
+         "span_name": "langgraph.graph.invoke",
+         "wrapper_method": task_wrapper,
+         "output_processor": INFERENCE
+    }
+]

monocle_apptrace/instrumentation/metamodel/llamaindex/_helper.py

@@ -16,34 +16,50 @@ from monocle_apptrace.instrumentation.common.utils import (
 logger = logging.getLogger(__name__)
 
 
+def extract_tools(instance):
+    tools = []
+    if not hasattr(instance, 'state') or not hasattr(instance.state, 'task_dict'):
+        return []
+    try:
+        data = next(iter(instance.state.task_dict.values())).task
+    except (AttributeError, StopIteration):
+        return []
+
+    if hasattr(data,'extra_state') and 'sources' in data.extra_state:
+        for tool_output in data.extra_state['sources']:
+            tool_name = tool_output.tool_name
+            if tool_name:
+                tools.append(tool_name)
+    return tools
+
+
 def extract_messages(args):
     """Extract system and user messages"""
     try:
         messages = []
-        if args and isinstance(args, (list, tuple)) and args[0]:
+
+        def process_message(msg):
+            """Processes a single message and extracts relevant information."""
+            if hasattr(msg, 'content') and hasattr(msg, 'role'):
+                role = getattr(msg.role, 'value', msg.role)
+                content = msg.content if role == "system" else extract_query_from_content(msg.content)
+                messages.append({role: content})
+
+        if isinstance(args, (list, tuple)) and args:
             for msg in args[0]:
-                if hasattr(msg, 'content') and hasattr(msg, 'role'):
-                    role = getattr(msg.role, 'value', msg.role)
-                    if role == "system":
-                        messages.append({role: msg.content})
-                    elif role in ["user", "human"]:
-                        user_message = extract_query_from_content(msg.content)
-                        messages.append({role: user_message})
-        if args and isinstance(args, dict):
+                process_message(msg)
+        if isinstance(args, dict):
             for msg in args.get("messages", []):
-                if hasattr(msg, 'content') and hasattr(msg, 'role'):
-                    role = getattr(msg.role, 'value', msg.role)
-                    if role == "system":
-                        messages.append({role: msg.content})
-                    elif role in ["user", "human"]:
-                        user_message = msg.content
-                        messages.append({role: user_message})
+                process_message(msg)
+        if args and isinstance(args, tuple):
+            messages.append(args[0])
+
         return [str(message) for message in messages]
+
     except Exception as e:
-        logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
+        logger.warning("Error in extract_messages: %s", str(e))
         return []
 
-
 def extract_assistant_message(response):
     try:
         if isinstance(response, str):
@@ -52,6 +68,8 @@ def extract_assistant_message(response):
             return [response.content]
         if hasattr(response, "message") and hasattr(response.message, "content"):
             return [response.message.content]
+        if hasattr(response,"response") and isinstance(response.response, str):
+            return [response.response]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
         return []
@@ -63,7 +81,7 @@ def extract_query_from_content(content):
         answer_prefix = "Answer:"
         query_start = content.find(query_prefix)
         if query_start == -1:
-            return None
+            return content
 
         query_start += len(query_prefix)
         answer_start = content.find(answer_prefix, query_start)

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

@@ -0,0 +1,47 @@
+from monocle_apptrace.instrumentation.metamodel.llamaindex import (
+    _helper,
+)
+
+AGENT = {
+    "type": "agent",
+    "attributes": [
+        [
+            {
+                "_comment": "Agent name, type and Tools.",
+                "attribute": "name",
+                "accessor": lambda arguments: arguments['instance'].__class__.__name__
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: 'Agent.oai'
+            },
+            {
+                "attribute": "tools",
+                "accessor": lambda arguments: _helper.extract_tools(arguments['instance'])
+            }
+        ]
+
+    ],
+    "events": [
+        {"name": "data.input",
+         "attributes": [
+
+             {
+                 "_comment": "this is instruction and user query to LLM",
+                 "attribute": "input",
+                 "accessor": lambda arguments: _helper.extract_messages(arguments['args'])
+             }
+         ]
+         },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                }
+            ]
+        }
+    ]
+}

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

@@ -1,7 +1,7 @@
 from monocle_apptrace.instrumentation.metamodel.llamaindex import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias, get_llm_type
 
 INFERENCE = {
     "type": "inference",
@@ -10,11 +10,13 @@ INFERENCE = {
             {
                 "_comment": "provider type ,name , deployment , inference_endpoint",
                 "attribute": "type",
-                "accessor": lambda arguments: 'inference.azure_oai'
+                "accessor": lambda arguments: 'inference.' + (get_llm_type(arguments['instance']) or 'generic')
+
             },
             {
                 "attribute": "provider_name",
-                "accessor": lambda arguments: _helper.extract_provider_name(arguments['instance'])
+                "accessor": lambda arguments: arguments['kwargs'].get('provider_name') or _helper.extract_provider_name(arguments['instance'])
+
             },
             {
                 "attribute": "deployment",

monocle_apptrace/instrumentation/metamodel/llamaindex/methods.py

@@ -2,6 +2,7 @@ from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.inference import (
     INFERENCE,
 )
+from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.agent import AGENT
 from monocle_apptrace.instrumentation.metamodel.llamaindex.entities.retrieval import (
     RETRIEVAL,
 )
@@ -29,14 +30,16 @@ LLAMAINDEX_METHODS = [
         "object": "BaseQueryEngine",
         "method": "query",
         "span_name": "llamaindex.query",
-        "wrapper_method": task_wrapper
+        "wrapper_method": task_wrapper,
+        "span_type": "workflow"
     },
     {
         "package": "llama_index.core.base.base_query_engine",
         "object": "BaseQueryEngine",
         "method": "aquery",
         "span_name": "llamaindex.query",
-        "wrapper_method": atask_wrapper
+        "wrapper_method": atask_wrapper,
+        "span_type": "workflow"
     },
     {
         "package": "llama_index.core.llms.custom",
@@ -52,7 +55,8 @@ LLAMAINDEX_METHODS = [
         "method": "achat",
         "span_name": "llamaindex.llmchat",
         "wrapper_method": atask_wrapper,
-        "output_processor": INFERENCE
+        "output_processor": INFERENCE,
+        
     },
     {
         "package": "llama_index.llms.openai.base",
@@ -85,5 +89,13 @@ LLAMAINDEX_METHODS = [
         "span_name": "llamaindex.mistralai",
         "wrapper_method": atask_wrapper,
         "output_processor": INFERENCE
+    },
+    {
+        "package": "llama_index.core.agent",
+        "object": "ReActAgent",
+        "method": "chat",
+        "span_name": "react.agent",
+        "wrapper_method": task_wrapper,
+        "output_processor": AGENT
     }
 ]

monocle_apptrace/instrumentation/metamodel/openai/_helper.py

@@ -0,0 +1,112 @@
+"""
+This module provides utility functions for extracting system, user,
+and assistant messages from various input formats.
+"""
+
+import logging
+from monocle_apptrace.instrumentation.common.utils import (
+    Option,
+    get_keys_as_tuple,
+    get_nested_value,
+    try_option,
+)
+
+
+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 [str(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(response):
+    try:
+        if response is not None and hasattr(response,"choices") and len(response.choices) >0:
+            if hasattr(response.choices[0],"message"):
+                return response.choices[0].message.content
+    except (IndexError, AttributeError) as e:
+        logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
+        return None
+
+def extract_provider_name(instance):
+    provider_url: Option[str] = try_option(getattr, instance._client.base_url, 'host')
+    return provider_url.unwrap_or(None)
+
+
+def extract_inference_endpoint(instance):
+    inference_endpoint: Option[str] = try_option(getattr, instance._client, 'base_url').map(str)
+    if inference_endpoint.is_none() and "meta" in instance.client.__dict__:
+        inference_endpoint = try_option(getattr, instance.client.meta, 'endpoint_url').map(str)
+
+    return inference_endpoint.unwrap_or(extract_provider_name(instance))
+
+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_input_span_events(kwargs):
+    if 'input' in kwargs and isinstance(kwargs['input'], list):
+        query = ' '.join(kwargs['input'])
+        return query
+
+
+def update_output_span_events(results):
+    if hasattr(results,'data') and isinstance(results.data, list):
+        embeddings = results.data
+        embedding_strings = [f"index={e.index}, embedding={e.embedding}" for e in embeddings]
+        output = '\n'.join(embedding_strings)
+        if len(output) > 100:
+            output = output[:100] + "..."
+        return output
+
+
+def update_span_from_llm_response(response):
+    meta_dict = {}
+    if response is not None and hasattr(response, "usage"):
+        if hasattr(response, "usage") and response.usage is not None:
+            token_usage = response.usage
+        else:
+            response_metadata = response.response_metadata
+            token_usage = response_metadata.get("token_usage")
+        if token_usage is not None:
+            meta_dict.update(
+                {"completion_tokens": getattr(response.usage, "completion_tokens", None)})
+            meta_dict.update({"prompt_tokens": getattr(response.usage, "prompt_tokens", None)})
+            meta_dict.update({"total_tokens": getattr(response.usage, "total_tokens", None)})
+    return meta_dict
+
+def extract_vector_input(vector_input: dict):
+    if 'input' in vector_input:
+        return vector_input['input']
+    return ""
+
+def extract_vector_output(vector_output):
+    try:
+        if hasattr(vector_output, 'data') and len(vector_output.data) > 0:
+            return vector_output.data[0].embedding
+    except Exception as e:
+        pass
+    return ""
+
+def get_inference_type(instance):
+    inference_type: Option[str] = try_option(getattr, instance._client, '_api_version')
+    if inference_type.unwrap_or(None):
+        return 'azure_openai'
+    else:
+        return 'openai'

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

@@ -0,0 +1,71 @@
+from monocle_apptrace.instrumentation.metamodel.openai import (
+    _helper,
+)
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+
+INFERENCE = {
+    "type": "inference",
+    "attributes": [
+        [
+            {
+                "_comment": "provider type ,name , deployment , inference_endpoint",
+                "attribute": "type",
+                "accessor": lambda arguments: 'inference.' + (_helper.get_inference_type(arguments['instance'])) or 'openai'
+            },
+            {
+                "attribute": "provider_name",
+                "accessor": lambda arguments: _helper.extract_provider_name(arguments['instance'])
+            },
+            {
+                "attribute": "deployment",
+                "accessor": lambda arguments: resolve_from_alias(arguments['instance'].__dict__, ['engine', 'azure_deployment', 'deployment_name', 'deployment_id', 'deployment'])
+            },
+            {
+                "attribute": "inference_endpoint",
+                "accessor": lambda arguments: resolve_from_alias(arguments['instance'].__dict__, ['azure_endpoint', 'api_base', 'endpoint']) or _helper.extract_inference_endpoint(arguments['instance'])
+            }
+        ],
+        [
+            {
+                "_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": [
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                }
+            ]
+        },
+        {
+            "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/openai/entities/retrieval.py

@@ -0,0 +1,43 @@
+from monocle_apptrace.instrumentation.metamodel.openai import (
+    _helper,
+)
+from monocle_apptrace.instrumentation.common.utils import resolve_from_alias
+
+RETRIEVAL = {
+    "type": "retrieval",
+    "attributes": [
+        [
+            {
+                "_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.embedding.' + 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.update_input_span_events(arguments['kwargs'])
+                 }
+         ]
+         },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.update_output_span_events(arguments['result'])
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/openai/methods.py

@@ -0,0 +1,45 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.metamodel.openai.entities.inference import (
+    INFERENCE,
+)
+from monocle_apptrace.instrumentation.metamodel.openai.entities.retrieval import (
+    RETRIEVAL,
+)
+
+OPENAI_METHODS = [
+    {
+        "package": "openai.resources.chat.completions",
+        "object": "Completions",
+        "method": "create",
+        "wrapper_method": task_wrapper,
+        "span_handler": "non_framework_handler",
+        "output_processor": INFERENCE
+    },
+    {
+        "package": "openai.resources.chat.completions",
+        "object": "AsyncCompletions",
+        "method": "create",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "non_framework_handler",
+        "output_processor": INFERENCE
+    },
+    {
+        "package": "openai.resources.embeddings",
+        "object": "Embeddings",
+        "method": "create",
+        "wrapper_method": task_wrapper,
+        "span_name": "openai_embeddings",
+        "span_handler": "non_framework_handler",
+        "output_processor": RETRIEVAL
+    },
+    {
+        "package": "openai.resources.embeddings",
+        "object": "AsyncEmbeddings",
+        "method": "create",
+        "wrapper_method": atask_wrapper,
+        "span_name": "openai_embeddings",
+        "span_handler": "non_framework_handler",
+        "output_processor": RETRIEVAL
+    }
+
+]

monocle_apptrace/instrumentation/metamodel/requests/__init__.py

@@ -0,0 +1,4 @@
+from os import environ
+from monocle_apptrace.instrumentation.common.constants import TRACE_PROPOGATION_URLS
+allowed_url_str = environ.get(TRACE_PROPOGATION_URLS, "")
+allowed_urls:list[str] = [] if allowed_url_str == "" else allowed_url_str.split(',')

monocle_apptrace/instrumentation/metamodel/requests/_helper.py

@@ -0,0 +1,31 @@
+import os
+from  monocle_apptrace.instrumentation.metamodel.requests import allowed_urls
+from opentelemetry.propagate import inject
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+
+def request_pre_task_processor(kwargs):
+    # add traceparent to the request headers in kwargs
+    if 'headers' not in kwargs:
+        headers = {}
+    else:
+        headers = kwargs['headers'].copy()
+    inject(headers)
+    kwargs['headers'] = headers
+
+def request_skip_span(kwargs) -> bool:
+    # add traceparent to the request headers in kwargs
+    if 'url' in kwargs:
+        url:str = kwargs['url']
+        for allowed_url in allowed_urls:
+            if url.startswith(allowed_url.strip()):
+                return False
+    return True
+
+class RequestSpanHandler(SpanHandler):
+
+    def pre_task_processing(self, to_wrap, wrapped, instance, args,kwargs, span):
+        request_pre_task_processor(kwargs)
+        super().pre_task_processing(to_wrap, wrapped, instance, args,kwargs,span)
+
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        return request_skip_span(kwargs)

monocle_apptrace/instrumentation/metamodel/requests/methods.py

@@ -0,0 +1,12 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+
+REQUESTS_METHODS = [
+    {
+        "package": "requests.sessions",
+        "object": "Session",
+        "method": "request",
+        "span_name": "http_requests",
+        "wrapper_method": task_wrapper,
+        "span_handler":"request_handler",
+    }
+]