monocle-apptrace 0.4.0b3__py3-none-any.whl → 0.4.1__py3-none-any.whl

monocle_apptrace/instrumentation/__init__.py

@@ -1 +1,2 @@
-from .common import *
+from .common import *
+from .metamodel.azfunc.wrapper import monocle_trace_azure_function_route

monocle_apptrace/instrumentation/common/constants.py

@@ -54,6 +54,9 @@ llm_type_map = {
     "anthropic": "anthropic",
     "chatanthropic":"anthropic",
     "anthropicchatgenerator":"anthropic",
+    "chatcompletionsclient": "azure_ai_inference",
+    "embeddingsclient": "azure_ai_inference",
+    "imageembeddingsclient": "azure_ai_inference",
 }
 
 MONOCLE_INSTRUMENTOR = "monocle_apptrace"

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -362,7 +362,7 @@ def monocle_trace_http_route(func):
     if inspect.iscoroutinefunction(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            return http_async_route_handler(func, *args, **kwargs)
+            return await http_async_route_handler(func, *args, **kwargs)
         return wrapper
     else:
         @wraps(func)

monocle_apptrace/instrumentation/common/span_handler.py

@@ -159,7 +159,7 @@ class SpanHandler:
                                 result = accessor(arguments)
                                 if result and isinstance(result, dict):
                                     result = dict((key, value) for key, value in result.items() if value is not None)
-                                if result and isinstance(result, (str, list, dict)):
+                                if result and isinstance(result, (int, str, list, dict)):
                                     if attribute_key is not None:
                                         event_attributes[attribute_key] = result
                                     else:

monocle_apptrace/instrumentation/common/utils.py

@@ -376,10 +376,12 @@ def get_status(arguments):
         return 'success'
 
 def get_exception_status_code(arguments):
-    if arguments['exception'] is not None and hasattr(arguments['exception'], 'code'):
+    if arguments['exception'] is not None and hasattr(arguments['exception'], 'code') and arguments['exception'].code is not None:
         return arguments['exception'].code
-    else:
+    elif arguments['exception'] is not None:
         return 'error'
+    else:
+        return 'success'
 
 def get_exception_message(arguments):
     if arguments['exception'] is not None:
@@ -390,6 +392,22 @@ def get_exception_message(arguments):
     else:
         return ''
 
+def get_status_code(arguments):
+    if arguments["exception"] is not None:
+        return get_exception_status_code(arguments)
+    elif hasattr(arguments["result"], "status"):
+        return arguments["result"].status
+    else:
+        return 'success'
+
+def get_status(arguments):
+    if arguments["exception"] is not None:
+        return 'error'
+    elif get_status_code(arguments) == 'success':
+        return 'success'
+    else:
+        return 'error'
+
 def patch_instance_method(obj, method_name, func):
     """
     Patch a special method (like __iter__) for a single instance.

monocle_apptrace/instrumentation/common/wrapper_method.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict
 from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, scope_wrapper
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler, NonFrameworkSpanHandler
+from monocle_apptrace.instrumentation.metamodel.azureaiinference.methods import AZURE_AI_INFERENCE_METHODS
 from monocle_apptrace.instrumentation.metamodel.botocore.methods import BOTOCORE_METHODS
 from monocle_apptrace.instrumentation.metamodel.botocore.handlers.botocore_span_handler import BotoCoreSpanHandler
 from monocle_apptrace.instrumentation.metamodel.langchain.methods import (
@@ -20,6 +21,8 @@ from monocle_apptrace.instrumentation.metamodel.teamsai.methods import (TEAMAI_M
 from monocle_apptrace.instrumentation.metamodel.anthropic.methods import (ANTHROPIC_METHODS, )
 from monocle_apptrace.instrumentation.metamodel.aiohttp.methods import (AIOHTTP_METHODS, )
 from monocle_apptrace.instrumentation.metamodel.aiohttp._helper import aiohttpSpanHandler
+from monocle_apptrace.instrumentation.metamodel.azfunc._helper import (azureSpanHandler)
+from monocle_apptrace.instrumentation.metamodel.azfunc.methods import AZFUNC_HTTP_METHODS
 class WrapperMethod:
     def __init__(
             self,
@@ -68,7 +71,7 @@ class WrapperMethod:
     def get_span_handler(self) -> SpanHandler:
         return self.span_handler()
 
-DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS + FLASK_METHODS + REQUESTS_METHODS + LANGGRAPH_METHODS + OPENAI_METHODS + TEAMAI_METHODS + ANTHROPIC_METHODS + AIOHTTP_METHODS
+DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS + FLASK_METHODS + REQUESTS_METHODS + LANGGRAPH_METHODS + OPENAI_METHODS + TEAMAI_METHODS + ANTHROPIC_METHODS + AIOHTTP_METHODS + AZURE_AI_INFERENCE_METHODS + AZFUNC_HTTP_METHODS
 
 MONOCLE_SPAN_HANDLERS: Dict[str, SpanHandler] = {
     "default": SpanHandler(),
@@ -79,4 +82,5 @@ MONOCLE_SPAN_HANDLERS: Dict[str, SpanHandler] = {
     "request_handler": RequestSpanHandler(),
     "non_framework_handler": NonFrameworkSpanHandler(),
     "openai_handler": OpenAISpanHandler(),
+    "azure_func_handler": azureSpanHandler(),
 }

monocle_apptrace/instrumentation/metamodel/anthropic/_helper.py

@@ -9,6 +9,7 @@ from monocle_apptrace.instrumentation.common.utils import (
     get_keys_as_tuple,
     get_nested_value,
     try_option,
+    get_exception_message,
 )
 
 
@@ -39,12 +40,36 @@ def extract_messages(kwargs):
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
 
+def get_exception_status_code(arguments):
+    if arguments['exception'] is not None and hasattr(arguments['exception'], 'status_code') and arguments['exception'].status_code is not None:
+        return arguments['exception'].status_code
+    elif arguments['exception'] is not None:
+        return 'error'
+    else:
+        return 'success'
 
-def extract_assistant_message(response):
+def get_status_code(arguments):
+    if arguments["exception"] is not None:
+        return get_exception_status_code(arguments)
+    elif hasattr(arguments["result"], "status"):
+        return arguments["result"].status
+    else:
+        return 'success'
+
+def extract_assistant_message(arguments):
     try:
-        if response is not None and hasattr(response,"content") and len(response.content) >0:
-            if hasattr(response.content[0],"text"):
-                return response.content[0].text
+        status = get_status_code(arguments)
+        response: str = ""
+        if status == 'success':
+            if arguments['result'] is not None and hasattr(arguments['result'],"content") and len(arguments['result'].content) >0:
+                if hasattr(arguments['result'].content[0],"text"):
+                    response = arguments['result'].content[0].text
+        else:
+            if arguments["exception"] is not None:
+                response = get_exception_message(arguments)
+            elif hasattr(arguments["result"], "error"):
+                response = arguments["result"].error
+        return response
     except (IndexError, AttributeError) as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
         return None

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

@@ -1,7 +1,9 @@
 from monocle_apptrace.instrumentation.metamodel.anthropic import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import resolve_from_alias, get_llm_type
+from monocle_apptrace.instrumentation.common.utils import (resolve_from_alias, get_llm_type,
+    get_status, get_status_code
+)
 
 INFERENCE = {
     "type": "inference",
@@ -52,10 +54,18 @@ INFERENCE = {
         {
             "name": "data.output",
             "attributes": [
+                {
+                    "attribute": "status",
+                    "accessor": lambda arguments: get_status(arguments)
+                },
+                {
+                    "attribute": "status_code",
+                    "accessor": lambda arguments: _helper.get_status_code(arguments)
+                },
                 {
                     "_comment": "this is result from LLM",
                     "attribute": "response",
-                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments)
                 }
             ]
         },

monocle_apptrace/instrumentation/metamodel/azfunc/_helper.py

@@ -0,0 +1,78 @@
+import logging
+from threading import local
+from monocle_apptrace.instrumentation.common.utils import extract_http_headers, clear_http_scopes, try_option, Option, MonocleSpanException
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.constants import HTTP_SUCCESS_CODES
+from urllib.parse import unquote, urlparse, ParseResult
+
+
+logger = logging.getLogger(__name__)
+MAX_DATA_LENGTH = 1000
+token_data = local()
+token_data.current_token = None
+
+def get_url(kwargs) -> ParseResult:
+    url_str = try_option(getattr, kwargs['req'], 'url')
+    url = url_str.unwrap_or(None)
+    if url is not None:
+        return urlparse(url)
+    else:
+        return None
+
+def get_route(kwargs) -> str:
+    url:ParseResult = get_url(kwargs)
+    if url is not None:
+        return url.path
+
+def get_method(kwargs) -> str:
+#    return args[0]['method'] if 'method' in args[0] else ""
+    http_method: Option[str] = try_option(getattr, kwargs['req'], 'method')
+    return http_method.unwrap_or("")
+
+def get_params(kwargs) -> dict:
+    url:ParseResult = get_url(kwargs)
+    if url is not None:
+        return unquote(url.query)
+
+def get_body(kwargs) -> dict:
+    if hasattr(kwargs['req'], 'get_body'):
+        response = kwargs.get_body()
+        if isinstance(response, bytes):
+            response = response.decode('utf-8', errors='ignore')
+    else:
+        response = ""
+    return response
+
+def extract_response(result) -> str:
+    if hasattr(result, 'get_body'):
+        response = result.get_body() #  text[0:max(result.text.__len__(), MAX_DATA_LENGTH)]
+        if isinstance(response, bytes):
+            response = response.decode('utf-8', errors='ignore')
+    else:
+        response = ""
+    return response
+
+def extract_status(result) -> str:
+    status = f"{result.status_code}" if hasattr(result, 'status_code') else ""
+    if status not in HTTP_SUCCESS_CODES:
+        error_message = extract_response(result)
+        raise MonocleSpanException(f"error: {status} - {error_message}")
+    return status
+
+def azure_func_pre_tracing(kwargs):
+    headers = kwargs['req'].headers if hasattr(kwargs['req'], 'headers') else {}
+    token_data.current_token = extract_http_headers(headers)
+
+def azure_func_post_tracing():
+    clear_http_scopes(token_data.current_token)
+    token_data.current_token = None
+
+class azureSpanHandler(SpanHandler):
+
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        azure_func_pre_tracing(kwargs)
+        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+    
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        azure_func_post_tracing()
+        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)

monocle_apptrace/instrumentation/metamodel/azfunc/entities/http.py

@@ -0,0 +1,51 @@
+from monocle_apptrace.instrumentation.metamodel.azfunc import _helper
+AZFUNC_HTTP_PROCESSOR = {
+    "type": "http.process",
+    "attributes": [
+        [
+            {
+                "_comment": "request method, request URI",
+                "attribute": "method",
+                "accessor": lambda arguments: _helper.get_method(arguments['kwargs'])
+            },
+            {
+                "_comment": "request method, request URI",
+                "attribute": "route",
+                "accessor": lambda arguments: _helper.get_route(arguments['kwargs'])
+            },
+            {
+                "_comment": "request method, request URI",
+                "attribute": "body",
+                "accessor": lambda arguments: _helper.get_body(arguments['kwargs'])
+            },
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "route params",
+                    "attribute": "params",
+                    "accessor": lambda arguments: _helper.get_params(arguments['kwargs'])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "status from HTTP response",
+                    "attribute": "status",
+                    "accessor": lambda arguments: _helper.extract_status(arguments['result'])
+                },
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_response(arguments['result'])
+                }
+            ]
+        }
+
+    ]
+}

monocle_apptrace/instrumentation/metamodel/azfunc/methods.py

@@ -0,0 +1,23 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.metamodel.azfunc.entities.http import AZFUNC_HTTP_PROCESSOR
+
+AZFUNC_HTTP_METHODS = [
+    {
+        "package": "monocle_apptrace.instrumentation.metamodel.azfunc.wrapper",
+        "object": "AzureFunctionRouteWrapper",
+        "method": "run_async",
+        "span_name": "azure_function_route",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "azure_func_handler",
+        "output_processor": AZFUNC_HTTP_PROCESSOR
+    },
+    {
+        "package": "monocle_apptrace.instrumentation.metamodel.azfunc.wrapper",
+        "object": "AzureFunctionRouteWrapper",
+        "method": "run_sync",
+        "span_name": "azure_function_route",
+        "wrapper_method": task_wrapper,
+        "span_handler": "azure_func_handler",
+        "output_processor": AZFUNC_HTTP_PROCESSOR
+    }
+]

monocle_apptrace/instrumentation/metamodel/azfunc/wrapper.py

@@ -0,0 +1,23 @@
+from functools import wraps
+import inspect
+
+def monocle_trace_azure_function_route(func):
+    if inspect.iscoroutinefunction(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            return await AzureFunctionRouteWrapper.run_async(func, *args, **kwargs)
+        return wrapper
+    else:
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return AzureFunctionRouteWrapper.run_sync(func ,*args, **kwargs)
+        return wrapper
+
+class AzureFunctionRouteWrapper:
+    @staticmethod
+    async def run_async(func, *args, **kwargs):
+        return await func(*args, **kwargs)
+
+    @staticmethod
+    def run_sync(func, *args, **kwargs):
+        return func(*args, **kwargs)

monocle_apptrace/instrumentation/metamodel/azureaiinference/__init__.py

@@ -0,0 +1 @@
+# Azure AI Inference instrumentation module

monocle_apptrace/instrumentation/metamodel/azureaiinference/_helper.py

@@ -0,0 +1,216 @@
+import logging
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+from monocle_apptrace.instrumentation.common.utils import (
+    resolve_from_alias,
+    get_exception_message,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def extract_messages(args_or_kwargs: Any) -> str:
+    """Extract messages from azure-ai-inference request arguments."""
+    try:
+        messages = []
+        if "instructions" in args_or_kwargs:
+            messages.append({"instructions": args_or_kwargs.get("instructions", {})})
+        if "input" in args_or_kwargs:
+            messages.append({"input": args_or_kwargs.get("input", {})})
+        if "messages" in args_or_kwargs and len(args_or_kwargs["messages"]) > 0:
+            for msg in args_or_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_inference_endpoint(instance: Any) -> str:
+    """Extract inference endpoint from azure-ai-inference client instance."""
+    try:
+        return instance._config.endpoint
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in extract_inference_endpoint: %s", str(e)
+        )
+        return ""
+
+
+def extract_embeddings_input(args_or_kwargs: Any) -> str:
+    """Extract input text from azure-ai-inference embeddings request."""
+    try:
+        # Handle both args and kwargs scenarios
+        if isinstance(args_or_kwargs, dict):
+            if "input" in args_or_kwargs:
+                input_data = args_or_kwargs["input"]
+            elif len(args_or_kwargs) > 0:
+                first_arg = list(args_or_kwargs.values())[0]
+                if isinstance(first_arg, dict) and "input" in first_arg:
+                    input_data = first_arg["input"]
+                else:
+                    input_data = first_arg
+            else:
+                return ""
+        elif hasattr(args_or_kwargs, "__iter__") and len(args_or_kwargs) > 0:
+            first_arg = args_or_kwargs[0]
+            if hasattr(first_arg, "get") and "input" in first_arg:
+                input_data = first_arg["input"]
+            else:
+                input_data = first_arg
+        else:
+            return ""
+
+        # Format input for display
+        if isinstance(input_data, (list, tuple)):
+            return " | ".join(str(item) for item in input_data)
+        else:
+            return str(input_data)
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in extract_embeddings_input: %s", str(e)
+        )
+        return ""
+
+
+def extract_assistant_message(arguments: Dict[str, Any]) -> str:
+    """Extract assistant response from azure-ai-inference completion result."""
+    try:
+        # Check for exception first
+        if arguments.get("exception") is not None:
+            return get_exception_message(arguments)
+
+        result = arguments.get("result")
+        if not result:
+            return ""
+        if hasattr(result, "output_text"):
+            # If the result has output_text attribute
+            return result.output_text
+        if (
+            result.choices
+            and result.choices[0].message
+            and result.choices[0].message.content
+        ):
+            # If the result is a chat completion with content
+            return result.choices[0].message.content
+
+        return str(result)
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in extract_assistant_message: %s", str(e)
+        )
+        return ""
+
+
+def extract_embeddings_output(arguments: Dict[str, Any]) -> str:
+    """Extract embeddings from azure-ai-inference embeddings result."""
+    try:
+        result = arguments.get("result")
+        if not result:
+            return ""
+
+        if hasattr(result, "data") and result.data:
+            # Format as summary of embeddings data
+            embeddings_info = []
+            for i, item in enumerate(result.data):
+                if hasattr(item, "embedding") and hasattr(item, "index"):
+                    embedding_length = len(item.embedding) if item.embedding else 0
+                    embeddings_info.append(
+                        f"index={item.index}, embedding=[{embedding_length} dimensions]"
+                    )
+            return " | ".join(embeddings_info)
+
+        return str(result)
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in extract_embeddings_output: %s", str(e)
+        )
+        return ""
+
+
+def update_span_from_llm_response(result: Any, instance: Any = None) -> Dict[str, Any]:
+    """Extract usage metadata from azure-ai-inference response."""
+    try:
+        attributes = {}
+
+        # Handle streaming responses with accumulated usage data
+        if hasattr(result, "usage") and result.usage:
+            usage = result.usage
+            if hasattr(usage, "completion_tokens"):
+                attributes["completion_tokens"] = usage.completion_tokens
+            if hasattr(usage, "prompt_tokens"):
+                attributes["prompt_tokens"] = usage.prompt_tokens
+            if hasattr(usage, "total_tokens"):
+                attributes["total_tokens"] = usage.total_tokens
+
+        # Handle regular response usage
+        elif hasattr(result, "usage"):
+            usage = result.usage
+            if hasattr(usage, "completion_tokens"):
+                attributes["completion_tokens"] = usage.completion_tokens
+            if hasattr(usage, "prompt_tokens"):
+                attributes["prompt_tokens"] = usage.prompt_tokens
+            if hasattr(usage, "total_tokens"):
+                attributes["total_tokens"] = usage.total_tokens
+
+        # Extract model information if available
+        if hasattr(result, "model"):
+            attributes["model"] = result.model
+
+        return attributes
+    except Exception as e:
+        logger.warning(
+            "Warning: Error occurred in update_span_from_llm_response: %s", str(e)
+        )
+        return {}
+
+
+def get_model_name(arguments: Dict[str, Any]) -> str:
+    """Extract model name from azure-ai-inference request arguments."""
+    try:
+
+        # Try to get from instance
+        instance = arguments.get("instance")
+        if arguments.get('kwargs') and arguments.get('kwargs').get('model'):
+            return arguments['kwargs'].get('model')
+        if instance and hasattr(instance, "_config") and hasattr(instance._config, "model"):
+            return instance._config.endpoint.split("/")[-1]
+
+        return ""
+    except Exception as e:
+        logger.warning("Warning: Error occurred in get_model_name: %s", str(e))
+        return ""
+
+
+def get_inference_type(arguments) -> str:
+    instance = arguments.get("instance")
+    if instance and hasattr(instance, "_config") and hasattr(instance._config, "endpoint"):
+        endpoint = instance._config.endpoint
+        try:
+            parsed = urlparse(endpoint)
+            hostname = parsed.hostname or endpoint
+            hostname = hostname.lower()
+        except Exception:
+            hostname = str(endpoint).lower()
+        if hostname.endswith("services.ai.azure.com"):
+            return "azure_ai_inference"
+        if hostname.endswith("openai.azure.com"):
+            return "azure_openai"
+    return "azure_ai_inference"
+
+
+def get_provider_name(instance: Any) -> str:
+    """Extract provider name from azure-ai-inference client instance."""
+    try:
+        # extract hostname from instance._config.endpoint
+        # https://okahu-openai-dev.openai.azure.com/openai/deployments/kshitiz-gpt => okahu-openai-dev.openai.azure.com
+        endpoint = instance._config.endpoint
+        if endpoint:
+            # Extract the hostname part
+            provider_name = endpoint.split("/")[2] if "/" in endpoint else endpoint
+            return provider_name
+    except Exception as e:
+        logger.warning("Warning: Error occurred in get_provider_name: %s", str(e))
+        return "azure_ai_inference"