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

monocle_apptrace/instrumentation/metamodel/azureaiinference/_helper.py

@@ -1,10 +1,13 @@
+import json
 import logging
 from typing import Any, Dict, Optional
 from urllib.parse import urlparse
 from monocle_apptrace.instrumentation.common.utils import (
-    resolve_from_alias,
+    get_json_dumps,
     get_exception_message,
+    get_status_code,
 )
+from monocle_apptrace.instrumentation.metamodel.finish_types import map_azure_ai_inference_finish_reason_to_finish_type
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +25,7 @@ def extract_messages(args_or_kwargs: Any) -> str:
                 if msg.get("content") and msg.get("role"):
                     messages.append({msg["role"]: msg["content"]})
 
-        return [str(message) for message in messages]
+        return [get_json_dumps(message) for message in messages]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
@@ -83,20 +86,29 @@ def extract_assistant_message(arguments: Dict[str, Any]) -> str:
             return get_exception_message(arguments)
 
         result = arguments.get("result")
+        role = "assistant"
+        messages = []
         if not result:
             return ""
         if hasattr(result, "output_text"):
             # If the result has output_text attribute
-            return result.output_text
-        if (
-            result.choices
+            role = getattr(result, "role", role)
+            if "assistant" in role.lower():
+                # If the role is assistant, we can assume it's a chat completion
+                role = "assistant"
+            messages.append({role: result.output_text})
+        if (hasattr(result, "choices")
+            and result.choices
             and result.choices[0].message
             and result.choices[0].message.content
         ):
+            role = getattr(result.choices[0].message, "role", role)
+            if "assistant" in role.lower():
+                # If the role is assistant, we can assume it's a chat completion
+                role = "assistant"
             # If the result is a chat completion with content
-            return result.choices[0].message.content
-
-        return str(result)
+            messages.append({role: result.choices[0].message.content})
+        return get_json_dumps(messages[0]) if messages else ""
     except Exception as e:
         logger.warning(
             "Warning: Error occurred in extract_assistant_message: %s", str(e)
@@ -214,3 +226,71 @@ def get_provider_name(instance: Any) -> str:
     except Exception as e:
         logger.warning("Warning: Error occurred in get_provider_name: %s", str(e))
         return "azure_ai_inference"
+
+
+def extract_finish_reason(arguments: Dict[str, Any]) -> Optional[str]:
+    """Extract finish_reason from Azure AI Inference response."""
+    try:
+        # Handle exception cases first
+        if arguments.get("exception") is not None:
+            ex = arguments["exception"]
+            if hasattr(ex, "message") and isinstance(ex.message, str):
+                message = ex.message
+                if "content_filter" in message.lower():
+                    return "content_filter"
+            return "error"
+        
+        result = arguments.get("result")
+        if result is None:
+            return None
+            
+        # Check various possible locations for finish_reason in Azure AI Inference responses
+        
+        # Direct finish_reason attribute
+        if hasattr(result, "finish_reason") and result.finish_reason:
+            return result.finish_reason
+            
+        # Check for choices structure (OpenAI-compatible format)
+        if hasattr(result, "choices") and result.choices:
+            choice = result.choices[0]
+            if hasattr(choice, "finish_reason") and choice.finish_reason:
+                return choice.finish_reason
+        
+        # Check for additional metadata or response attributes
+        if hasattr(result, "additional_kwargs") and result.additional_kwargs:
+            kwargs = result.additional_kwargs
+            if isinstance(kwargs, dict):
+                for key in ["finish_reason", "stop_reason"]:
+                    if key in kwargs:
+                        return kwargs[key]
+        
+        # Check for response metadata
+        if hasattr(result, "response_metadata") and result.response_metadata:
+            metadata = result.response_metadata
+            if isinstance(metadata, dict):
+                for key in ["finish_reason", "stop_reason"]:
+                    if key in metadata:
+                        return metadata[key]
+        
+        # Check for streaming response with accumulated finish reason
+        if hasattr(result, "type") and result.type == "stream":
+            # For streaming responses, default to stop if completed successfully
+            return "stop"
+        
+        # If no specific finish reason found, infer from status
+        status_code = get_status_code(arguments)
+        if status_code == 'success':
+            return "stop"  # Default success finish reason
+        elif status_code == 'error':
+            return "error"
+            
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_finish_reason: %s", str(e))
+        return None
+    
+    return None
+
+
+def map_finish_reason_to_finish_type(finish_reason: Optional[str]) -> Optional[str]:
+    """Map Azure AI Inference finish_reason to finish_type."""
+    return map_azure_ai_inference_finish_reason_to_finish_type(finish_reason)

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

@@ -3,6 +3,7 @@ import time
 from types import SimpleNamespace
 from monocle_apptrace.instrumentation.metamodel.azureaiinference import _helper
 from monocle_apptrace.instrumentation.common.utils import (
+    get_error_message,
     resolve_from_alias, 
     patch_instance_method,
     get_status,
@@ -20,13 +21,14 @@ def process_stream(to_wrap, response, span_processor):
     stream_closed_time = None
     accumulated_response = ""
     token_usage = None
+    role = "assistant"
 
     # 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
+            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage, role
 
             try:
                 item = original_next()
@@ -34,6 +36,8 @@ def process_stream(to_wrap, response, span_processor):
                 # 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, 'role') and choice.delta.role:
+                        role = choice.delta.role
                     if hasattr(choice, 'delta') and hasattr(choice.delta, 'content') and choice.delta.content:
                         if waiting_for_first_token:
                             waiting_for_first_token = False
@@ -53,6 +57,7 @@ def process_stream(to_wrap, response, span_processor):
                 if span_processor:
                     ret_val = SimpleNamespace(
                         type="stream",
+                        role=role,
                         timestamps={
                             "data.input": int(stream_start_time),
                             "data.output": int(first_token_time),
@@ -77,7 +82,7 @@ def process_stream(to_wrap, response, span_processor):
         original_anext = response.__anext__
 
         async def new_anext(self):
-            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage
+            nonlocal waiting_for_first_token, first_token_time, stream_closed_time, accumulated_response, token_usage, role
 
             try:
                 item = await original_anext()
@@ -85,6 +90,8 @@ def process_stream(to_wrap, response, span_processor):
                 # 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, 'role') and choice.delta.role:
+                        role = choice.delta.role
                     if hasattr(choice, 'delta') and hasattr(choice.delta, 'content') and choice.delta.content:
                         if waiting_for_first_token:
                             waiting_for_first_token = False
@@ -104,6 +111,7 @@ def process_stream(to_wrap, response, span_processor):
                 if span_processor:
                     ret_val = SimpleNamespace(
                         type="stream",
+                        role=role,
                         timestamps={
                             "data.input": int(stream_start_time),
                             "data.output": int(first_token_time),
@@ -183,12 +191,8 @@ INFERENCE = {
                     "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)
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
                 }
             ]
         },
@@ -201,6 +205,16 @@ INFERENCE = {
                         arguments['result'], 
                         arguments.get('instance')
                     )
+                },
+                {
+                    "attribute": "finish_reason",
+                    "accessor": lambda arguments: _helper.extract_finish_reason(arguments)
+                },
+                {
+                    "attribute": "finish_type",
+                    "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(
+                        _helper.extract_finish_reason(arguments)
+                    )
                 }
             ]
         }

monocle_apptrace/instrumentation/metamodel/botocore/_helper.py

@@ -8,7 +8,8 @@ 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,)
+from monocle_apptrace.instrumentation.common.utils import ( get_exception_message, get_json_dumps, get_status_code,)
+from monocle_apptrace.instrumentation.metamodel.finish_types import map_bedrock_finish_reason_to_finish_type
 logger = logging.getLogger(__name__)
 
 
@@ -25,7 +26,7 @@ def extract_messages(args):
                 role = args['messages'][0]['role']
                 user_message = extract_query_from_content(args['messages'][0]['content'][0]['text'])
                 messages.append({role: user_message})
-        return [str(d) for d in messages]
+        return [get_json_dumps(message) for message in messages]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
         return []
@@ -39,18 +40,11 @@ def get_exception_status_code(arguments):
     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(arguments):
     try:
         status = get_status_code(arguments)
-        response: str = ""
+        messages = []
+        role = "assistant"
         if status == 'success':
             if "Body" in arguments['result'] and hasattr(arguments['result']['Body'], "_raw_stream"):
                 raw_stream = getattr(arguments['result']['Body'], "_raw_stream")
@@ -59,20 +53,20 @@ def extract_assistant_message(arguments):
                     response_str = response_bytes.decode('utf-8')
                     response_dict = json.loads(response_str)
                     arguments['result']['Body'] = BytesIO(response_bytes)
-                    response = response_dict["answer"]
+                    messages.append({role: 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
+                    messages.append({role: reply})
         else:
             if arguments["exception"] is not None:
-                response = get_exception_message(arguments)
+                return get_exception_message(arguments)
             elif hasattr(arguments["result"], "error"):
-                response = arguments["result"].error
-        return response
+                return arguments["result"].error
+        return get_json_dumps(messages[0]) if messages else ""
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
         return []
@@ -118,3 +112,85 @@ def update_span_from_llm_response(response, instance):
             meta_dict.update({"prompt_tokens": resolve_from_alias(token_usage,["prompt_tokens","input_tokens","inputTokens"])})
             meta_dict.update({"total_tokens": resolve_from_alias(token_usage,["total_tokens","totalTokens"])})
     return meta_dict
+
+
+def extract_finish_reason(arguments):
+    """Extract finish_reason/stopReason from Bedrock response."""
+    try:
+        # Handle exception cases first
+        if arguments.get("exception") is not None:
+            return "error"
+        
+        result = arguments.get("result")
+        if result is None:
+            return None
+            
+        # Check various possible locations for finish_reason in Bedrock responses
+        
+        # Direct stopReason attribute (Bedrock Converse API)
+        if "stopReason" in result:
+            return result["stopReason"]
+            
+        # Check for completionReason (some Bedrock models)
+        if "completionReason" in result:
+            return result["completionReason"]
+            
+        # Check for output structure (Bedrock Converse API)
+        if "output" in result and "message" in result["output"]:
+            message = result["output"]["message"]
+            if "stopReason" in message:
+                return message["stopReason"]
+                
+        # Check for nested result structure
+        if "result" in result:
+            nested_result = result["result"]
+            if "stopReason" in nested_result:
+                return nested_result["stopReason"]
+            if "completionReason" in nested_result:
+                return nested_result["completionReason"]
+                
+        # Check for streaming response accumulated finish reason
+        if "type" in result and result["type"] == "stream":
+            if "stopReason" in result:
+                return result["stopReason"]
+                
+        # Check for response metadata
+        if "ResponseMetadata" in result:
+            metadata = result["ResponseMetadata"]
+            if "stopReason" in metadata:
+                return metadata["stopReason"]
+                
+        # Check for Body content (for some Bedrock responses)
+        if "Body" in result:
+            body = result["Body"]
+            if hasattr(body, "_raw_stream"):
+                raw_stream = getattr(body, "_raw_stream")
+                if hasattr(raw_stream, "data"):
+                    response_bytes = getattr(raw_stream, "data")
+                    response_str = response_bytes.decode('utf-8')
+                    try:
+                        response_dict = json.loads(response_str)
+                        if "stopReason" in response_dict:
+                            return response_dict["stopReason"]
+                        if "completionReason" in response_dict:
+                            return response_dict["completionReason"]
+                    except json.JSONDecodeError:
+                        pass
+                        
+        # If no specific finish reason found, infer from status
+        status_code = get_status_code(arguments)
+        if status_code == 'success':
+            return "end_turn"  # Default successful completion
+        elif status_code == 'error':
+            return "error"
+            
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_finish_reason: %s", str(e))
+        return None
+    
+    return None
+
+
+def map_finish_reason_to_finish_type(finish_reason):
+    """Map Bedrock finish_reason/stopReason to finish_type."""
+    return map_bedrock_finish_reason_to_finish_type(finish_reason)

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, get_status,)
+from monocle_apptrace.instrumentation.common.utils import (get_error_message, get_llm_type, get_status,)
 INFERENCE = {
     "type": "inference",
     "attributes": [
@@ -43,14 +43,9 @@ 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)
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
                 },
                 {
                     "_comment": "this is response from LLM",
@@ -66,6 +61,16 @@ INFERENCE = {
                     "_comment": "this is metadata usage from LLM",
                     "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'],
                                                                                         arguments['instance'])
+                },
+                {
+                    "attribute": "finish_reason",
+                    "accessor": lambda arguments: _helper.extract_finish_reason(arguments)
+                },
+                {
+                    "attribute": "finish_type",
+                    "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(
+                        _helper.extract_finish_reason(arguments)
+                    )
                 }
             ]
         }

monocle_apptrace/instrumentation/metamodel/botocore/handlers/botocore_span_handler.py

@@ -20,7 +20,7 @@ class BotoCoreSpanHandler(SpanHandler):
                     instrumented_method = instrumentor(to_wrap, wrapped, span_name, return_value, original_method)
                     setattr(return_value, method_name, instrumented_method)
 
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value,token=None):
         self._botocore_processor(to_wrap=to_wrap, wrapped=wrapped, instance=instance, return_value=return_value, args=args,
                                  kwargs=kwargs)
         return super().post_tracing(to_wrap, wrapped, instance, args, kwargs,return_value)

monocle_apptrace/instrumentation/metamodel/fastapi/_helper.py

@@ -0,0 +1,82 @@
+import logging
+from threading import local
+from monocle_apptrace.instrumentation.common.utils import extract_http_headers, clear_http_scopes
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.constants import HTTP_SUCCESS_CODES
+from monocle_apptrace.instrumentation.common.utils import MonocleSpanException
+from opentelemetry.context import get_current
+from opentelemetry.trace import Span
+from opentelemetry.trace.propagation import _SPAN_KEY
+import json
+import urllib.parse
+
+logger = logging.getLogger(__name__)
+MAX_DATA_LENGTH = 1000
+token_data = local()
+token_data.current_token = None
+
+def get_route(scope) -> str:
+    return scope.get('path', '')
+
+def get_method(scope) -> str:
+    return scope.get('method', '')
+
+def get_params(args) -> dict:
+    try:
+        query_bytes = args.get("query_string", "")
+        query_str = query_bytes.decode('utf-8')
+        params = urllib.parse.parse_qs(query_str)
+        question = params.get('question', [''])[0]
+        return question
+    except Exception as e:
+        logger.warning(f"Error extracting params: {e}")
+        return {}
+
+def extract_response(response) -> str:
+    try:
+        if hasattr(response, 'body'):
+            data = response.body
+            answer = json.loads(data.decode("utf-8"))
+            return answer
+    except Exception as e:
+        logger.warning(f"Error extracting response: {e}")
+        return ""
+
+def extract_status(instance) -> str:
+    status = f"{instance.status_code}" if hasattr(instance, 'status_code') else ""
+    if status not in HTTP_SUCCESS_CODES:
+        error_message = extract_response(instance)
+        raise MonocleSpanException(f"error: {status} - {error_message}")
+    return status
+
+def fastapi_pre_tracing(scope):
+    headers = {k.decode('utf-8').lower(): v.decode('utf-8')
+               for k, v in scope.get('headers', [])}
+    token_data.current_token = extract_http_headers(headers)
+
+def fastapi_post_tracing():
+    clear_http_scopes(token_data.current_token)
+    token_data.current_token = None
+
+class FastAPISpanHandler(SpanHandler):
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        scope = args[0] if args else {}
+        fastapi_pre_tracing(scope)
+        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        fastapi_post_tracing()
+        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+
+class FastAPIResponseSpanHandler(SpanHandler):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        try:
+            ctx = get_current()
+            if ctx is not None:
+                parent_span: Span = ctx.get(_SPAN_KEY)
+                if parent_span is not None:
+                    self.hydrate_events(to_wrap, wrapped, instance, args, kwargs,
+                                        return_value, parent_span=parent_span)
+        except Exception as e:
+            logger.info(f"Failed to propagate fastapi response: {e}")
+        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)

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

@@ -0,0 +1,44 @@
+from monocle_apptrace.instrumentation.metamodel.fastapi import _helper
+
+FASTAPI_HTTP_PROCESSOR = {
+    "type": "http.process",
+    "attributes": [
+        [
+            {
+                "attribute": "method",
+                "accessor": lambda arguments: _helper.get_method(arguments['args'][0])
+            },
+            {
+                "attribute": "route",
+                "accessor": lambda arguments: _helper.get_route(arguments['args'][0])
+            },
+        ]
+    ]
+}
+
+FASTAPI_RESPONSE_PROCESSOR = {
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "attribute": "params",
+                    "accessor": lambda arguments: _helper.get_params(arguments['args'][0])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "attribute": "status",
+                    "accessor": lambda arguments: _helper.extract_status(arguments['instance'])
+                },
+                {
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_response(arguments['instance'])
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/fastapi/methods.py

@@ -0,0 +1,23 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper
+from monocle_apptrace.instrumentation.metamodel.fastapi.entities.http import FASTAPI_HTTP_PROCESSOR, FASTAPI_RESPONSE_PROCESSOR
+
+FASTAPI_METHODS = [
+    # {
+    #     "package": "fastapi",
+    #     "object": "FastAPI",
+    #     "method": "__call__",
+    #     "wrapper_method": atask_wrapper,
+    #     "span_name": "fastapi.request",
+    #     "span_handler": "fastapi_handler",
+    #     "output_processor": FASTAPI_HTTP_PROCESSOR,
+    # },
+    # {
+    #     "package": "starlette.responses",
+    #     "object": "Response",
+    #     "method": "__call__",
+    #     "span_name": "fastapi.response",
+    #     "wrapper_method": atask_wrapper,
+    #     "span_handler": "fastapi_response_handler",
+    #     "output_processor": FASTAPI_RESPONSE_PROCESSOR
+    # }
+]