monocle-apptrace 0.4.2__py3-none-any.whl → 0.5.0b1__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/finish_types.py

@@ -0,0 +1,387 @@
+"""
+This module provides common finish reason mappings and finish type enums
+for different AI providers (OpenAI, Anthropic, Gemini, LangChain, LlamaIndex, Azure AI Inference).
+"""
+
+from enum import Enum
+
+class FinishType(Enum):
+    """Enum for standardized finish types across all AI providers."""
+    SUCCESS = "success"
+    TRUNCATED = "truncated"
+    CONTENT_FILTER = "content_filter"
+    ERROR = "error"
+    REFUSAL = "refusal"
+    RATE_LIMITED = "rate_limited"
+
+# OpenAI finish reason mapping
+OPENAI_FINISH_REASON_MAPPING = {
+    "stop": FinishType.SUCCESS.value,
+    "tool_calls": FinishType.SUCCESS.value,
+    "function_call": FinishType.SUCCESS.value,  # deprecated but still possible
+    "length": FinishType.TRUNCATED.value,
+    "content_filter": FinishType.CONTENT_FILTER.value
+}
+
+# Anthropic finish reason mapping
+ANTHROPIC_FINISH_REASON_MAPPING = {
+    "end_turn": FinishType.SUCCESS.value,         # Natural completion
+    "max_tokens": FinishType.TRUNCATED.value,     # Hit max_tokens limit
+    "stop_sequence": FinishType.SUCCESS.value,    # Hit user stop sequence
+    "tool_use": FinishType.SUCCESS.value,         # Tool use triggered
+    "pause_turn": FinishType.SUCCESS.value,       # Paused for tool or server action
+    "refusal": FinishType.REFUSAL.value,          # Refused for safety/ethics
+}
+
+# Gemini finish reason mapping
+GEMINI_FINISH_REASON_MAPPING = {
+    "STOP": FinishType.SUCCESS.value,
+    "MAX_TOKENS": FinishType.TRUNCATED.value,
+    "SAFETY": FinishType.CONTENT_FILTER.value,
+    "RECITATION": FinishType.CONTENT_FILTER.value,
+    "OTHER": FinishType.ERROR.value,
+    "FINISH_REASON_UNSPECIFIED": None
+}
+
+# LlamaIndex finish reason mapping
+# LlamaIndex often wraps underlying provider responses, similar to LangChain
+LLAMAINDEX_FINISH_REASON_MAPPING = {
+    # Standard completion reasons
+    "stop": FinishType.SUCCESS.value,
+    "complete": FinishType.SUCCESS.value,
+    "finished": FinishType.SUCCESS.value,
+    "success": FinishType.SUCCESS.value,
+    
+    # Token limits
+    "length": FinishType.TRUNCATED.value,
+    "max_tokens": FinishType.TRUNCATED.value,
+    "token_limit": FinishType.TRUNCATED.value,
+    "truncated": FinishType.TRUNCATED.value,
+    
+    # Tool/function calling
+    "tool_calls": FinishType.SUCCESS.value,
+    "function_call": FinishType.SUCCESS.value,
+    "agent_finish": FinishType.SUCCESS.value,
+    
+    # Content filtering and safety
+    "content_filter": FinishType.CONTENT_FILTER.value,
+    "safety": FinishType.CONTENT_FILTER.value,
+    "filtered": FinishType.CONTENT_FILTER.value,
+    
+    # Errors
+    "error": FinishType.ERROR.value,
+    "failed": FinishType.ERROR.value,
+    "exception": FinishType.ERROR.value,
+    
+    # Provider-specific reasons that might pass through LlamaIndex
+    # OpenAI reasons
+    "end_turn": FinishType.SUCCESS.value,  # Anthropic
+    "stop_sequence": FinishType.SUCCESS.value,  # Anthropic
+    "STOP": FinishType.SUCCESS.value,  # Gemini
+    "SAFETY": FinishType.CONTENT_FILTER.value,  # Gemini
+    "RECITATION": FinishType.CONTENT_FILTER.value,  # Gemini
+    "OTHER": FinishType.ERROR.value,  # Gemini
+}
+
+# Azure AI Inference finish reason mapping
+AZURE_AI_INFERENCE_FINISH_REASON_MAPPING = {
+    # Standard completion reasons
+    "stop": FinishType.SUCCESS.value,
+    "completed": FinishType.SUCCESS.value,
+    "finished": FinishType.SUCCESS.value,
+    
+    # Token limits
+    "length": FinishType.TRUNCATED.value,
+    "max_tokens": FinishType.TRUNCATED.value,
+    "token_limit": FinishType.TRUNCATED.value,
+    "max_completion_tokens": FinishType.TRUNCATED.value,
+    
+    # Tool/function calling
+    "tool_calls": FinishType.SUCCESS.value,
+    "function_call": FinishType.SUCCESS.value,
+    
+    # Content filtering and safety
+    "content_filter": FinishType.CONTENT_FILTER.value,
+    "content_filtered": FinishType.CONTENT_FILTER.value,
+    "safety": FinishType.CONTENT_FILTER.value,
+    "responsible_ai_policy": FinishType.CONTENT_FILTER.value,
+    
+    # Errors
+    "error": FinishType.ERROR.value,
+    "failed": FinishType.ERROR.value,
+    "exception": FinishType.ERROR.value,
+    "timeout": FinishType.ERROR.value,
+    
+    # Azure-specific reasons
+    "model_error": FinishType.ERROR.value,
+    "service_unavailable": FinishType.ERROR.value,
+    "rate_limit": FinishType.ERROR.value,
+}
+
+# AWS Bedrock finish reason mapping
+# Based on AWS Bedrock Converse API and model-specific APIs
+BEDROCK_FINISH_REASON_MAPPING = {
+    # Standard completion reasons
+    "end_turn": FinishType.SUCCESS.value,           # Natural completion
+    "stop": FinishType.SUCCESS.value,               # Hit stop sequence
+    "stop_sequence": FinishType.SUCCESS.value,      # Stop sequence triggered
+    "completed": FinishType.SUCCESS.value,          # Completion finished successfully
+    
+    # Token limits
+    "max_tokens": FinishType.TRUNCATED.value,       # Hit max_tokens limit
+    "length": FinishType.TRUNCATED.value,           # Token length limit
+    "max_length": FinishType.TRUNCATED.value,       # Maximum length reached
+    "token_limit": FinishType.TRUNCATED.value,      # Token limit reached
+    
+    # Tool/function calling
+    "tool_use": FinishType.SUCCESS.value,           # Tool use triggered
+    "function_call": FinishType.SUCCESS.value,      # Function call triggered
+    
+    # Content filtering and safety
+    "content_filter": FinishType.CONTENT_FILTER.value,    # Content filtered
+    "content_filtered": FinishType.CONTENT_FILTER.value,  # Content was filtered
+    "safety": FinishType.CONTENT_FILTER.value,            # Safety filter triggered
+    "guardrails": FinishType.CONTENT_FILTER.value,        # Bedrock guardrails triggered
+    "blocked": FinishType.CONTENT_FILTER.value,           # Request blocked
+    
+    # Errors
+    "error": FinishType.ERROR.value,                # General error
+    "failed": FinishType.ERROR.value,               # Request failed
+    "exception": FinishType.ERROR.value,            # Exception occurred
+    "timeout": FinishType.ERROR.value,              # Request timeout
+    "model_error": FinishType.ERROR.value,          # Model-specific error
+    "service_unavailable": FinishType.ERROR.value,  # Service unavailable
+    "throttled": FinishType.ERROR.value,            # Request throttled
+    "rate_limit": FinishType.ERROR.value,           # Rate limit exceeded
+    "validation_error": FinishType.ERROR.value,     # Validation error
+    
+    # Model-specific reasons (various Bedrock models)
+    # Claude models via Bedrock
+    "end_turn": FinishType.SUCCESS.value,           # Already defined above
+    "max_tokens": FinishType.TRUNCATED.value,       # Already defined above
+    "stop_sequence": FinishType.SUCCESS.value,      # Already defined above
+    "tool_use": FinishType.SUCCESS.value,           # Already defined above
+    
+    # AI21 models via Bedrock
+    "endoftext": FinishType.SUCCESS.value,          # AI21 end of text
+    "length": FinishType.TRUNCATED.value,           # AI21 length limit
+    
+    # Cohere models via Bedrock
+    "COMPLETE": FinishType.SUCCESS.value,           # Cohere completion
+    "MAX_TOKENS": FinishType.TRUNCATED.value,       # Cohere max tokens
+    "ERROR": FinishType.ERROR.value,                # Cohere error
+    
+    # Meta Llama models via Bedrock
+    "stop": FinishType.SUCCESS.value,               # Already defined above
+    "length": FinishType.TRUNCATED.value,           # Already defined above
+    
+    # Amazon Titan models via Bedrock
+    "FINISH": FinishType.SUCCESS.value,             # Titan finish
+    "LENGTH": FinishType.TRUNCATED.value,           # Titan length limit
+    "CONTENT_FILTERED": FinishType.CONTENT_FILTER.value,  # Titan content filter
+}
+
+# LangChain finish reason mapping
+# LangChain often wraps underlying provider responses, so we include common finish reasons
+# that might appear in LangChain response objects
+LANGCHAIN_FINISH_REASON_MAPPING = {
+    # Standard completion reasons
+    "stop": FinishType.SUCCESS.value,
+    "complete": FinishType.SUCCESS.value,
+    "finished": FinishType.SUCCESS.value,
+    
+    # Token limits
+    "length": FinishType.TRUNCATED.value,
+    "max_tokens": FinishType.TRUNCATED.value,
+    "token_limit": FinishType.TRUNCATED.value,
+    
+    # Tool/function calling
+    "tool_calls": FinishType.SUCCESS.value,
+    "function_call": FinishType.SUCCESS.value,
+    
+    # Content filtering and safety
+    "content_filter": FinishType.CONTENT_FILTER.value,
+    "safety": FinishType.CONTENT_FILTER.value,
+    "filtered": FinishType.CONTENT_FILTER.value,
+    
+    # Errors
+    "error": FinishType.ERROR.value,
+    "failed": FinishType.ERROR.value,
+    "exception": FinishType.ERROR.value,
+    
+    # Provider-specific reasons that might pass through LangChain
+    # OpenAI reasons
+    "stop": FinishType.SUCCESS.value,  # Already defined above
+    
+    # Anthropic reasons
+    "end_turn": FinishType.SUCCESS.value,
+    "stop_sequence": FinishType.SUCCESS.value,
+    
+    # Gemini reasons
+    "STOP": FinishType.SUCCESS.value,
+    "SAFETY": FinishType.CONTENT_FILTER.value,
+    "RECITATION": FinishType.CONTENT_FILTER.value,
+    "OTHER": FinishType.ERROR.value,
+}
+
+TEAMSAI_FINISH_REASON_MAPPING = {
+    "success": FinishType.SUCCESS.value,
+    "error": FinishType.ERROR.value,
+    "too_long": FinishType.TRUNCATED.value,
+    "rate_limited": FinishType.RATE_LIMITED.value,
+    "invalid_response": FinishType.ERROR.value,
+}
+
+def map_openai_finish_reason_to_finish_type(finish_reason):
+    """Map OpenAI finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    return OPENAI_FINISH_REASON_MAPPING.get(finish_reason, None)
+
+
+def map_anthropic_finish_reason_to_finish_type(finish_reason):
+    """Map Anthropic stop_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    return ANTHROPIC_FINISH_REASON_MAPPING.get(finish_reason, None)
+
+
+def map_gemini_finish_reason_to_finish_type(finish_reason):
+    """Map Gemini finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    return GEMINI_FINISH_REASON_MAPPING.get(finish_reason, None)
+
+
+def map_langchain_finish_reason_to_finish_type(finish_reason):
+    """Map LangChain finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+    
+    # Try direct mapping first
+    if finish_reason in LANGCHAIN_FINISH_REASON_MAPPING:
+        return LANGCHAIN_FINISH_REASON_MAPPING[finish_reason]
+    
+    # Try lowercase mapping
+    if finish_reason_lower in LANGCHAIN_FINISH_REASON_MAPPING:
+        return LANGCHAIN_FINISH_REASON_MAPPING[finish_reason_lower]
+    
+    # If no direct mapping, try to infer from common patterns
+    if any(keyword in finish_reason_lower for keyword in ['stop', 'complete', 'success', 'done']):
+        return FinishType.SUCCESS.value
+    elif any(keyword in finish_reason_lower for keyword in ['length', 'token', 'limit', 'truncat']):
+        return FinishType.TRUNCATED.value
+    elif any(keyword in finish_reason_lower for keyword in ['filter', 'safety', 'block']):
+        return FinishType.CONTENT_FILTER.value
+    elif any(keyword in finish_reason_lower for keyword in ['error', 'fail', 'exception']):
+        return FinishType.ERROR.value
+    
+    return None
+
+
+def map_llamaindex_finish_reason_to_finish_type(finish_reason):
+    """Map LlamaIndex finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+    
+    # Try direct mapping first
+    if finish_reason in LLAMAINDEX_FINISH_REASON_MAPPING:
+        return LLAMAINDEX_FINISH_REASON_MAPPING[finish_reason]
+    
+    # Try lowercase mapping
+    if finish_reason_lower in LLAMAINDEX_FINISH_REASON_MAPPING:
+        return LLAMAINDEX_FINISH_REASON_MAPPING[finish_reason_lower]
+    
+    # If no direct mapping, try to infer from common patterns
+    if any(keyword in finish_reason_lower for keyword in ['stop', 'complete', 'success', 'done', 'finish']):
+        return FinishType.SUCCESS.value
+    elif any(keyword in finish_reason_lower for keyword in ['length', 'token', 'limit', 'truncat']):
+        return FinishType.TRUNCATED.value
+    elif any(keyword in finish_reason_lower for keyword in ['filter', 'safety', 'block']):
+        return FinishType.CONTENT_FILTER.value
+    elif any(keyword in finish_reason_lower for keyword in ['error', 'fail', 'exception']):
+        return FinishType.ERROR.value
+    
+    return None
+
+
+def map_azure_ai_inference_finish_reason_to_finish_type(finish_reason):
+    """Map Azure AI Inference finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+    
+    # Try direct mapping first
+    if finish_reason in AZURE_AI_INFERENCE_FINISH_REASON_MAPPING:
+        return AZURE_AI_INFERENCE_FINISH_REASON_MAPPING[finish_reason]
+    
+    # Try lowercase mapping
+    if finish_reason_lower in AZURE_AI_INFERENCE_FINISH_REASON_MAPPING:
+        return AZURE_AI_INFERENCE_FINISH_REASON_MAPPING[finish_reason_lower]
+    
+    # If no direct mapping, try to infer from common patterns
+    if any(keyword in finish_reason_lower for keyword in ['stop', 'complete', 'success', 'done', 'finish']):
+        return FinishType.SUCCESS.value
+    elif any(keyword in finish_reason_lower for keyword in ['length', 'token', 'limit', 'truncat']):
+        return FinishType.TRUNCATED.value
+    elif any(keyword in finish_reason_lower for keyword in ['filter', 'safety', 'block', 'responsible_ai', 'content_filter']):
+        return FinishType.CONTENT_FILTER.value
+    elif any(keyword in finish_reason_lower for keyword in ['error', 'fail', 'exception', 'timeout', 'unavailable', 'rate_limit']):
+        return FinishType.ERROR.value
+    
+    return None
+
+
+def map_bedrock_finish_reason_to_finish_type(finish_reason):
+    """Map AWS Bedrock finish_reason/stopReason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+    
+    # Try direct mapping first
+    if finish_reason in BEDROCK_FINISH_REASON_MAPPING:
+        return BEDROCK_FINISH_REASON_MAPPING[finish_reason]
+    
+    # Try lowercase mapping
+    if finish_reason_lower in BEDROCK_FINISH_REASON_MAPPING:
+        return BEDROCK_FINISH_REASON_MAPPING[finish_reason_lower]
+    
+    # If no direct mapping, try to infer from common patterns
+    if any(keyword in finish_reason_lower for keyword in ['stop', 'complete', 'success', 'done', 'finish', 'end_turn', 'endoftext']):
+        return FinishType.SUCCESS.value
+    elif any(keyword in finish_reason_lower for keyword in ['length', 'token', 'limit', 'truncat', 'max_tokens']):
+        return FinishType.TRUNCATED.value
+    elif any(keyword in finish_reason_lower for keyword in ['filter', 'safety', 'block', 'guardrails', 'content_filter']):
+        return FinishType.CONTENT_FILTER.value
+    elif any(keyword in finish_reason_lower for keyword in ['error', 'fail', 'exception', 'timeout', 'unavailable', 'rate_limit', 'throttled', 'validation']):
+        return FinishType.ERROR.value
+    
+    return None
+
+def map_teamsai_finish_reason_to_finish_type(finish_reason):
+    """Map TeamsAI finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+
+    # Try direct mapping first
+    if finish_reason in TEAMSAI_FINISH_REASON_MAPPING:
+        return TEAMSAI_FINISH_REASON_MAPPING[finish_reason]
+
+    # Try lowercase mapping
+    if finish_reason_lower in TEAMSAI_FINISH_REASON_MAPPING:
+        return TEAMSAI_FINISH_REASON_MAPPING[finish_reason_lower]
+
+    return None

monocle_apptrace/instrumentation/metamodel/flask/_helper.py

@@ -11,8 +11,6 @@ from opentelemetry.trace.propagation import _SPAN_KEY
 
 logger = logging.getLogger(__name__)
 MAX_DATA_LENGTH = 1000
-token_data = local()
-token_data.current_token = None
 
 def get_route(args) -> str:
     return args[0]['PATH_INFO'] if 'PATH_INFO' in args[0] else ""
@@ -47,21 +45,18 @@ def flask_pre_tracing(args):
         if key.startswith("HTTP_"):
             new_key = key[5:].lower().replace("_", "-")
             headers[new_key] = value
-    token_data.current_token = extract_http_headers(headers)
+    return extract_http_headers(headers)
 
-def flask_post_tracing():
-    clear_http_scopes(token_data.current_token)
-    token_data.current_token = None
+def flask_post_tracing(token):
+    clear_http_scopes(token)
 
 class FlaskSpanHandler(SpanHandler):
 
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
-        flask_pre_tracing(args)
-        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        return flask_pre_tracing(args)
     
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
-        flask_post_tracing()
-        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token):
+        flask_post_tracing(token)
 
 class FlaskResponseSpanHandler(SpanHandler):
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):

monocle_apptrace/instrumentation/metamodel/gemini/_helper.py

@@ -1,8 +1,13 @@
 import logging
 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_gemini_finish_reason_to_finish_type,
+    GEMINI_FINISH_REASON_MAPPING
+)
 
 logger = logging.getLogger(__name__)
 
@@ -32,9 +37,9 @@ def extract_messages(kwargs):
                     if hasattr(part, 'text'):
                         messages.append({getattr(content, 'role', 'user'): part.text})
         elif isinstance(contents, str):
-            messages.append({'input': contents})
+            messages.append({'user': contents})
 
-        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 []
@@ -42,20 +47,35 @@ def extract_messages(kwargs):
 def extract_assistant_message(arguments):
     try:
         status = get_status_code(arguments)
-        response: str = ""
+        messages = []
+        role = "assistant"
+        if hasattr(arguments['result'], "candidates") and len(arguments['result'].candidates) > 0 and hasattr(arguments['result'].candidates[0], "content") and hasattr(arguments['result'].candidates[0].content, "role"):
+                role = arguments["result"].candidates[0].content.role
         if status == 'success':
             if hasattr(arguments['result'], "text") and len(arguments['result'].text):
-                response = arguments['result'].text
+                messages.append({role: arguments['result'].text})
         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 (IndexError, AttributeError) as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
         return None
 
+def update_input_span_events(kwargs):
+    if 'contents' in kwargs and isinstance(kwargs['contents'], list) and len(kwargs['contents']) > 0:
+        query = kwargs['contents'][0]
+        return query
+
+def update_output_span_events(results):
+    if hasattr(results,'embeddings') and isinstance(results.embeddings, list) and len(results.embeddings) > 0:
+        embeddings = results.embeddings[0]
+        if hasattr(embeddings, 'values') and isinstance(embeddings.values, list) and len(embeddings.values) > 100:
+            output = str(results.embeddings[0].values[:100]) + "..."
+            return output
+
 def extract_inference_endpoint(instance):
     try:
         if hasattr(instance,'_api_client') and hasattr(instance._api_client, '_http_options'):
@@ -74,3 +94,27 @@ def update_span_from_llm_response(response, instance):
             meta_dict.update({"prompt_tokens": token_usage.prompt_token_count })
             meta_dict.update({"total_tokens": token_usage.total_token_count})
     return meta_dict
+
+def extract_finish_reason(arguments):
+    """Extract finish_reason from Gemini response"""
+    try:
+        if arguments["exception"] is not None:
+            return None
+            
+        response = arguments["result"]
+        
+        # Handle Gemini response structure
+        if (response is not None and 
+            hasattr(response, "candidates") and 
+            len(response.candidates) > 0 and 
+            hasattr(response.candidates[0], "finish_reason")):
+            return response.candidates[0].finish_reason
+            
+    except (IndexError, AttributeError) 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 Gemini finish_reason to finish_type based on the possible errors mapping"""
+    return map_gemini_finish_reason_to_finish_type(finish_reason)

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

@@ -1,7 +1,8 @@
 from monocle_apptrace.instrumentation.metamodel.gemini import (
     _helper,
 )
-from monocle_apptrace.instrumentation.common.utils import get_llm_type, get_status, get_status_code
+from monocle_apptrace.instrumentation.common.utils import get_error_message
+
 INFERENCE = {
     "type": "inference",
     "attributes": [
@@ -9,7 +10,7 @@ INFERENCE = {
             {
                 "_comment": "provider type  , inference_endpoint",
                 "attribute": "type",
-                "accessor": lambda arguments: 'inference.gemini'
+                "accessor": lambda arguments: 'inference.vertexai' if hasattr(arguments['instance'],"vertexai") and arguments['instance'].vertexai else 'inference.gemini'
             },
             {
                 "attribute": "inference_endpoint",
@@ -46,13 +47,8 @@ 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: get_status_code(arguments)
+                    "attribute": "error_code",
+                    "accessor": lambda arguments: get_error_message(arguments)
                 },
                 {
                     "attribute": "response",
@@ -66,6 +62,18 @@ INFERENCE = {
                 {
                     "_comment": "this is metadata usage from LLM",
                     "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'], arguments['instance'])
+                },
+                {
+                    "_comment": "finish reason from Gemini response",
+                    "attribute": "finish_reason",
+                    "accessor": lambda arguments: _helper.extract_finish_reason(arguments)
+                },
+                {
+                    "_comment": "finish type mapped from finish reason",
+                    "attribute": "finish_type",
+                    "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(
+                        _helper.extract_finish_reason(arguments)
+                    )
                 }
             ]
         }

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

@@ -0,0 +1,43 @@
+from monocle_apptrace.instrumentation.metamodel.gemini import (
+    _helper,
+)
+
+RETRIEVAL = {
+    "type": "retrieval",
+    "attributes": [
+        [
+            {
+                "_comment": "Embedding Model",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.resolve_from_alias(arguments['kwargs'],
+                                                                         ['model'])
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: 'model.embedding.' + _helper.resolve_from_alias(arguments['kwargs'],
+                                                                                        ['model'])
+            }
+        ]
+    ],
+    "events": [
+        {
+         "name": "data.input",
+         "attributes": [
+             {
+                 "attribute": "input",
+                 "accessor": lambda arguments: _helper.update_input_span_events(arguments['kwargs'])
+             }
+         ]
+         },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.update_output_span_events(arguments['result'])
+                }
+            ]
+        }
+
+    ]
+}

monocle_apptrace/instrumentation/metamodel/gemini/methods.py

@@ -2,6 +2,9 @@ from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
 from monocle_apptrace.instrumentation.metamodel.gemini.entities.inference import (
     INFERENCE,
 )
+from monocle_apptrace.instrumentation.metamodel.gemini.entities.retrieval import (
+    RETRIEVAL,
+)
 
 GEMINI_METHODS = [
     {
@@ -10,5 +13,12 @@ GEMINI_METHODS = [
       "method": "generate_content",
       "wrapper_method": task_wrapper,
       "output_processor": INFERENCE,
+    },
+    {
+      "package": "google.genai.models",
+      "object": "Models",
+      "method": "embed_content",
+      "wrapper_method": task_wrapper,
+      "output_processor": RETRIEVAL,
     }
 ]