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

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/gemini/_helper.py

@@ -0,0 +1,76 @@
+import logging
+from monocle_apptrace.instrumentation.common.utils import (
+    get_exception_message,
+    get_status_code,
+)
+
+logger = logging.getLogger(__name__)
+
+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 extract_messages(kwargs):
+    """Extract system and user messages"""
+    try:
+        messages = []
+        config = kwargs.get('config')
+        if config and hasattr(config, 'system_instruction'):
+            system_instructions = getattr(config, 'system_instruction', None)
+            if system_instructions:
+                messages.append({'system': system_instructions})
+
+        contents = kwargs.get('contents')
+        if isinstance(contents, list):
+            for content in contents:
+                if hasattr(content, 'parts') and getattr(content, 'parts'):
+                    part = content.parts[0]
+                    if hasattr(part, 'text'):
+                        messages.append({getattr(content, 'role', 'user'): part.text})
+        elif isinstance(contents, str):
+            messages.append({'input': contents})
+
+        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(arguments):
+    try:
+        status = get_status_code(arguments)
+        response: str = ""
+        if status == 'success':
+            if hasattr(arguments['result'], "text") and len(arguments['result'].text):
+                response = arguments['result'].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
+
+def extract_inference_endpoint(instance):
+    try:
+        if hasattr(instance,'_api_client') and hasattr(instance._api_client, '_http_options'):
+            if hasattr(instance._api_client._http_options,'base_url'):
+                return instance._api_client._http_options.base_url
+    except Exception as e:
+        logger.warning("Warning: Error occurred in inference endpoint: %s", str(e))
+        return []
+
+def update_span_from_llm_response(response, instance):
+    meta_dict = {}
+    if response is not None and hasattr(response, "usage_metadata") and response.usage_metadata is not None:
+        token_usage = response.usage_metadata
+        if token_usage is not None:
+            meta_dict.update({"completion_tokens": token_usage.candidates_token_count})
+            meta_dict.update({"prompt_tokens": token_usage.prompt_token_count })
+            meta_dict.update({"total_tokens": token_usage.total_token_count})
+    return meta_dict

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

@@ -0,0 +1,75 @@
+from monocle_apptrace.instrumentation.metamodel.gemini import (
+    _helper,
+)
+from monocle_apptrace.instrumentation.common.utils import get_llm_type, get_status, get_status_code
+INFERENCE = {
+    "type": "inference",
+    "attributes": [
+        [
+            {
+                "_comment": "provider type  , inference_endpoint",
+                "attribute": "type",
+                "accessor": lambda arguments: 'inference.gemini'
+            },
+            {
+                "attribute": "inference_endpoint",
+                "accessor": lambda arguments: _helper.extract_inference_endpoint(arguments['instance'])
+            }
+        ],
+        [
+            {
+                "_comment": "LLM Model",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.resolve_from_alias(arguments['kwargs'],
+                                                                         ['model'])
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: 'model.llm.' + _helper.resolve_from_alias(arguments['kwargs'],
+                                                                                        ['model'])
+            }
+        ]
+    ],
+    "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": "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)
+                }
+            ]
+        },
+        {
+            "name": "metadata",
+            "attributes": [
+                {
+                    "_comment": "this is metadata usage from LLM",
+                    "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'], arguments['instance'])
+                }
+            ]
+        }
+
+
+    ]
+}

monocle_apptrace/instrumentation/metamodel/gemini/methods.py

@@ -0,0 +1,14 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.metamodel.gemini.entities.inference import (
+    INFERENCE,
+)
+
+GEMINI_METHODS = [
+    {
+      "package": "google.genai.models",
+      "object": "Models",
+      "method": "generate_content",
+      "wrapper_method": task_wrapper,
+      "output_processor": INFERENCE,
+    }
+]

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)
                 }
             ]
         },

monocle_apptrace/instrumentation/metamodel/langchain/_helper.py

@@ -9,6 +9,8 @@ from monocle_apptrace.instrumentation.common.utils import (
     get_keys_as_tuple,
     get_nested_value,
     try_option,
+    get_exception_message,
+    get_status_code,
 )
 
 
@@ -35,18 +37,23 @@ def extract_messages(args):
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
 
+def extract_assistant_message(arguments):
+    status = get_status_code(arguments)
+    response: str = ""
+    if status == 'success':
+        if isinstance(arguments['result'], str):
+            response = arguments['result']
+        if hasattr(arguments['result'], "content"):
+            response = arguments['result'].content
+        if hasattr(arguments['result'], "message") and hasattr(arguments['result'].message, "content"):
+            response =  arguments['result'].message.content
+    else:
+        if arguments["exception"] is not None:
+            response = get_exception_message(arguments)
+        elif hasattr(arguments["result"], "error"):
+            response = arguments["result"].error
 
-def extract_assistant_message(response):
-    try:
-        if isinstance(response, str):
-            return [response]
-        if hasattr(response, "content"):
-            return [response.content]
-        if hasattr(response, "message") and hasattr(response.message, "content"):
-            return [response.message.content]
-    except Exception as e:
-        logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
-        return []
+    return response
 
 
 def extract_provider_name(instance):
@@ -131,5 +138,4 @@ def update_span_from_llm_response(response, instance):
                 {"completion_tokens": token_usage.get("completion_tokens") or token_usage.get("output_tokens")})
             meta_dict.update({"prompt_tokens": token_usage.get("prompt_tokens") or token_usage.get("input_tokens")})
             meta_dict.update({"total_tokens": token_usage.get("total_tokens")})
-    return meta_dict
-
+    return meta_dict

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

@@ -1,7 +1,7 @@
 from monocle_apptrace.instrumentation.metamodel.langchain 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.framework",
@@ -54,8 +54,16 @@ INFERENCE = {
             "attributes": [
                 {
                     "_comment": "this is result 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)
                 }
             ]
         },