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

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"

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

@@ -0,0 +1,208 @@
+import logging
+import time
+from types import SimpleNamespace
+from monocle_apptrace.instrumentation.metamodel.azureaiinference import _helper
+from monocle_apptrace.instrumentation.common.utils import (
+    resolve_from_alias, 
+    patch_instance_method,
+    get_status,
+    get_exception_status_code
+)
+
+logger = logging.getLogger(__name__)
+
+
+def process_stream(to_wrap, response, span_processor):
+    """Process streaming responses from Azure AI Inference."""
+    waiting_for_first_token = True
+    stream_start_time = time.time_ns()
+    first_token_time = stream_start_time
+    stream_closed_time = None
+    accumulated_response = ""
+    token_usage = None
+
+    # For sync iteration - patch __next__ instead of __iter__
+    if to_wrap and hasattr(response, "__next__"):
+        original_next = response.__next__
+
+        def new_next(self):
+            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage
+
+            try:
+                item = original_next()
+                
+                # Handle Azure AI Inference streaming chunks
+                if hasattr(item, 'choices') and item.choices:
+                    choice = item.choices[0]
+                    if hasattr(choice, 'delta') and hasattr(choice.delta, 'content') and choice.delta.content:
+                        if waiting_for_first_token:
+                            waiting_for_first_token = False
+                            first_token_time = time.time_ns()
+
+                        accumulated_response += choice.delta.content
+                
+                # Check for usage information at the end of stream
+                if hasattr(item, 'usage') and item.usage:
+                    token_usage = item.usage
+                    stream_closed_time = time.time_ns()
+
+                return item
+
+            except StopIteration:
+                # Stream is complete, process final span
+                if span_processor:
+                    ret_val = SimpleNamespace(
+                        type="stream",
+                        timestamps={
+                            "data.input": int(stream_start_time),
+                            "data.output": int(first_token_time),
+                            "metadata": int(stream_closed_time or time.time_ns()),
+                        },
+                        output_text=accumulated_response,
+                        usage=token_usage,
+                    )
+                    span_processor(ret_val)
+                raise
+            except Exception as e:
+                logger.warning(
+                    "Warning: Error occurred while processing item in new_next: %s",
+                    str(e),
+                )
+                raise
+
+        patch_instance_method(response, "__next__", new_next)
+        
+    # For async iteration - patch __anext__ instead of __aiter__
+    if to_wrap and hasattr(response, "__anext__"):
+        original_anext = response.__anext__
+
+        async def new_anext(self):
+            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage
+
+            try:
+                item = await original_anext()
+                
+                # Handle Azure AI Inference streaming chunks
+                if hasattr(item, 'choices') and item.choices:
+                    choice = item.choices[0]
+                    if hasattr(choice, 'delta') and hasattr(choice.delta, 'content') and choice.delta.content:
+                        if waiting_for_first_token:
+                            waiting_for_first_token = False
+                            first_token_time = time.time_ns()
+
+                        accumulated_response += choice.delta.content
+                
+                # Check for usage information at the end of stream
+                if hasattr(item, 'usage') and item.usage:
+                    token_usage = item.usage
+                    stream_closed_time = time.time_ns()
+
+                return item
+
+            except StopAsyncIteration:
+                # Stream is complete, process final span
+                if span_processor:
+                    ret_val = SimpleNamespace(
+                        type="stream",
+                        timestamps={
+                            "data.input": int(stream_start_time),
+                            "data.output": int(first_token_time),
+                            "metadata": int(stream_closed_time or time.time_ns()),
+                        },
+                        output_text=accumulated_response,
+                        usage=token_usage,
+                    )
+                    span_processor(ret_val)
+                raise
+            except Exception as e:
+                logger.warning(
+                    "Warning: Error occurred while processing item in new_anext: %s",
+                    str(e),
+                )
+                raise
+
+        patch_instance_method(response, "__anext__", new_anext)
+
+
+INFERENCE = {
+    "type": "inference",
+    "is_auto_close": lambda kwargs: kwargs.get("stream", False) is False,
+    "response_processor": process_stream,
+    "attributes": [
+        [
+            {
+                "_comment": "Azure AI Inference provider type, endpoint",
+                "attribute": "type",
+                "accessor": lambda arguments: f"inference.{_helper.get_inference_type(arguments)}"
+            },
+            {
+                "attribute": "provider_name",
+                "accessor": lambda arguments: _helper.get_provider_name(arguments['instance'])
+            },
+            {
+                "attribute": "inference_endpoint",
+                "accessor": lambda arguments: _helper.extract_inference_endpoint(arguments['instance'])
+            },
+            {
+                "attribute": "deployment",
+                "accessor": lambda arguments: resolve_from_alias(
+                    arguments['instance'].__dict__,
+                    ['deployment', 'deployment_name', 'azure_deployment', '_deployment']
+                )
+            }
+        ],
+        [
+            {
+                "_comment": "LLM Model information",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_model_name(arguments)
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: f"model.llm.{_helper.get_model_name(arguments)}" if _helper.get_model_name(arguments) else "model.llm.unknown"
+            }
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "Chat messages input to Azure AI Inference",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.extract_messages(arguments['kwargs'])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "Response from Azure AI Inference",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments)
+                },
+                {
+                    "attribute": "status",
+                    "accessor": lambda arguments: get_status(arguments)
+                },
+                {
+                    "attribute": "status_code",
+                    "accessor": lambda arguments: get_exception_status_code(arguments)
+                }
+            ]
+        },
+        {
+            "name": "metadata",
+            "attributes": [
+                {
+                    "_comment": "Usage metadata from Azure AI Inference",
+                    "accessor": lambda arguments: _helper.update_span_from_llm_response(
+                        arguments['result'], 
+                        arguments.get('instance')
+                    )
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/azureaiinference/methods.py

@@ -0,0 +1,23 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.metamodel.azureaiinference.entities.inference import INFERENCE
+
+AZURE_AI_INFERENCE_METHODS = [
+    # Chat Completions - Synchronous
+    {
+        "package": "azure.ai.inference",
+        "object": "ChatCompletionsClient",
+        "method": "complete",
+        "wrapper_method": task_wrapper,
+        "span_handler": "non_framework_handler",
+        "output_processor": INFERENCE
+    },
+    # Chat Completions - Asynchronous
+    {
+        "package": "azure.ai.inference.aio",
+        "object": "ChatCompletionsClient",
+        "method": "complete",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "non_framework_handler",
+        "output_processor": INFERENCE
+    }
+]

monocle_apptrace/instrumentation/metamodel/botocore/_helper.py

@@ -8,7 +8,7 @@ import json
 from io import BytesIO
 from functools import wraps
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
-
+from monocle_apptrace.instrumentation.common.utils import ( get_exception_message,)
 logger = logging.getLogger(__name__)
 
 
@@ -30,24 +30,49 @@ def extract_messages(args):
         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'], 'response') and arguments['exception'].response is not None:
+        if "ResponseMetadata" in arguments['exception'].response and "HTTPStatusCode" in arguments['exception'].response["ResponseMetadata"]:
+            return arguments['exception'].response["ResponseMetadata"]["HTTPStatusCode"]
+    elif arguments['exception'] is not None:
+        return 'error'
+    else:
+        return 'success'
+
+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(response):
+def extract_assistant_message(arguments):
     try:
-        if "Body" in response and hasattr(response['Body'], "_raw_stream"):
-            raw_stream = getattr(response['Body'], "_raw_stream")
-            if hasattr(raw_stream, "data"):
-                response_bytes = getattr(raw_stream, "data")
-                response_str = response_bytes.decode('utf-8')
-                response_dict = json.loads(response_str)
-                response['Body'] = BytesIO(response_bytes)
-                return [response_dict["answer"]]
-        if "output" in response:
-            output = response.get("output", {})
-            message = output.get("message", {})
-            content = message.get("content", [])
-            if isinstance(content, list) and len(content) > 0 and "text" in content[0]:
-                reply = content[0]["text"]
-                return [reply]
+        status = get_status_code(arguments)
+        response: str = ""
+        if status == 'success':
+            if "Body" in arguments['result'] and hasattr(arguments['result']['Body'], "_raw_stream"):
+                raw_stream = getattr(arguments['result']['Body'], "_raw_stream")
+                if hasattr(raw_stream, "data"):
+                    response_bytes = getattr(raw_stream, "data")
+                    response_str = response_bytes.decode('utf-8')
+                    response_dict = json.loads(response_str)
+                    arguments['result']['Body'] = BytesIO(response_bytes)
+                    response = response_dict["answer"]
+            if "output" in arguments['result']:
+                output = arguments['result'].get("output", {})
+                message = output.get("message", {})
+                content = message.get("content", [])
+                if isinstance(content, list) and len(content) > 0 and "text" in content[0]:
+                    reply = content[0]["text"]
+                    response = reply
+        else:
+            if arguments["exception"] is not None:
+                response = get_exception_message(arguments)
+            elif hasattr(arguments["result"], "error"):
+                response = arguments["result"].error
+        return response
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
         return []

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

@@ -1,7 +1,7 @@
 from monocle_apptrace.instrumentation.metamodel.botocore import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import get_llm_type
+from monocle_apptrace.instrumentation.common.utils import (get_llm_type, get_status,)
 INFERENCE = {
     "type": "inference",
     "attributes": [
@@ -33,7 +33,6 @@ INFERENCE = {
     "events": [
         {"name": "data.input",
          "attributes": [
-
              {
                  "_comment": "this is instruction and user query to LLM",
                  "attribute": "input",
@@ -44,10 +43,19 @@ INFERENCE = {
         {
             "name": "data.output",
             "attributes": [
+            {
+                    "_comment": "this is result from LLM",
+                    "attribute": "status",
+                    "accessor": lambda arguments: get_status(arguments)
+                },
+                {
+                    "attribute": "status_code",
+                    "accessor": lambda arguments: _helper.get_status_code(arguments)
+                },
                 {
                     "_comment": "this is response 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/flask/_helper.py

@@ -70,7 +70,7 @@ class FlaskResponseSpanHandler(SpanHandler):
             if _parent_span_context is not None:
                 parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
                 if parent_span is not None:
-                    self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, return_value, parent_span)
+                    self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, return_value, parent_span=parent_span)
         except Exception as e:
             logger.info(f"Failed to propogate flask response: {e}")
         super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)

monocle_apptrace/instrumentation/metamodel/haystack/_helper.py

@@ -5,6 +5,8 @@ from monocle_apptrace.instrumentation.common.utils import (
     get_keys_as_tuple,
     get_nested_value,
     try_option,
+    get_exception_message,
+    get_status_code,
 )
 logger = logging.getLogger(__name__)
 
@@ -52,19 +54,25 @@ def extract_question_from_prompt(content):
         logger.warning("Warning: Error occurred in extract_question_from_prompt: %s", str(e))
         return ""
 
-
-def extract_assistant_message(response):
-    try:
-        if "replies" in response:
-            reply = response["replies"][0]
+def extract_assistant_message(arguments):
+    status = get_status_code(arguments)
+    response: str = ""
+    if status == 'success':
+        if "replies" in arguments['result']:
+            reply = arguments['result']["replies"][0]
             if hasattr(reply, 'content'):
-                return [reply.content]
-            if hasattr(reply, 'text'):
-                return [reply.text]
-            return [reply]
-    except Exception as e:
-        logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
-        return []
+                response = reply.content
+            elif hasattr(reply, 'text'):
+                response = reply.text
+            else:
+                response = reply
+    else:
+        if arguments["exception"] is not None:
+            response = get_exception_message(arguments)
+        elif hasattr(response, "error"):
+            response = arguments['result'].error
+
+    return response
 
 
 def get_vectorstore_deployment(my_map):

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

@@ -1,7 +1,7 @@
 from monocle_apptrace.instrumentation.metamodel.haystack import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import get_llm_type
+from monocle_apptrace.instrumentation.common.utils import get_llm_type, get_status, get_status_code
 
 INFERENCE = {
     "type": "inference.framework",
@@ -60,8 +60,16 @@ INFERENCE = {
             "attributes": [
                 {
                     "_comment": "this is response from LLM",
+                    "attribute": "status",
+                    "accessor": lambda arguments: get_status(arguments)
+                },
+                {
+                    "attribute": "status_code",
+                    "accessor": lambda arguments: get_status_code(arguments)
+                },
+                {
                     "attribute": "response",
-                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments)
                 }
             ]
         },