payi 0.1.0a110__py3-none-any.whl → 0.1.0a137__py3-none-any.whl

payi/lib/BedrockInstrumentor.py

@@ -1,18 +1,31 @@
+from __future__ import annotations
+
 import os
 import json
-from typing import Any, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence
 from functools import wraps
 from typing_extensions import override
 
 from wrapt import ObjectProxy, wrap_function_wrapper  # type: ignore
 
-from payi.lib.helpers import PayiCategories, PayiHeaderNames, payi_aws_bedrock_url
+from payi.lib.helpers import PayiCategories, PayiHeaderNames, PayiPropertyNames, payi_aws_bedrock_url
 from payi.types.ingest_units_params import Units
-from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
-from .instrument import _ChunkResult, _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
+from .instrument import (
+    PayiInstrumentAwsBedrockConfig,
+    _Context,
+    _IsStreaming,
+    _PayiInstrumentor,
+)
 from .version_helper import get_version_helper
+from .ProviderRequest import _ChunkResult, _StreamingType, _ProviderRequest
+
+if TYPE_CHECKING:
+    from tokenizers import Tokenizer  # type: ignore
+else:
+    Tokenizer = None  
 
+GUARDRAIL_SEMANTIC_FAILURE_DESCRIPTION = "Bedrock Guardrails intervened"
 
 class BedrockInstrumentor:
     _module_name: str = "boto3"
@@ -20,6 +33,37 @@ class BedrockInstrumentor:
 
     _instrumentor: _PayiInstrumentor
 
+    _guardrail_trace: bool = True
+
+    _model_mapping: Dict[str, _Context] = {}
+
+    _add_streaming_xproxy_result: bool = False
+
+    @staticmethod
+    def get_mapping(model_id: Optional[str]) -> _Context:
+        if not model_id:
+            return  {}
+
+        return BedrockInstrumentor._model_mapping.get(model_id, {})
+
+    @staticmethod
+    def configure(aws_config: Optional[PayiInstrumentAwsBedrockConfig]) -> None:
+        if not aws_config:
+            return
+        
+        trace = aws_config.get("guardrail_trace", True)
+        if trace is None:
+            trace = True
+        BedrockInstrumentor._guardrail_trace = trace
+
+        add_streaming_xproxy_result = aws_config.get("add_streaming_xproxy_result", False)
+        if add_streaming_xproxy_result:
+            BedrockInstrumentor._add_streaming_xproxy_result = add_streaming_xproxy_result
+
+        model_mappings = aws_config.get("model_mappings", [])
+        if model_mappings:
+            BedrockInstrumentor._model_mapping = _PayiInstrumentor._model_mapping_to_context_dict(model_mappings)
+
     @staticmethod
     def instrument(instrumentor: _PayiInstrumentor) -> None:
         BedrockInstrumentor._instrumentor = instrumentor
@@ -51,10 +95,10 @@ def create_client_wrapper(instrumentor: _PayiInstrumentor, wrapped: Any, instanc
 
     try:
         client: Any = wrapped(*args, **kwargs)
-        client.invoke_model = wrap_invoke(instrumentor, client.invoke_model)
-        client.invoke_model_with_response_stream = wrap_invoke_stream(instrumentor, client.invoke_model_with_response_stream)
-        client.converse = wrap_converse(instrumentor, client.converse)
-        client.converse_stream = wrap_converse_stream(instrumentor, client.converse_stream)
+        client.invoke_model = wrap_invoke(instrumentor, client.invoke_model, client)
+        client.invoke_model_with_response_stream = wrap_invoke_stream(instrumentor, client.invoke_model_with_response_stream, client)
+        client.converse = wrap_converse(instrumentor, client.converse, client)
+        client.converse_stream = wrap_converse_stream(instrumentor, client.converse_stream, client)
 
         instrumentor._logger.debug(f"Instrumented bedrock client")
 
@@ -100,17 +144,20 @@ def _redirect_to_payi(request: Any, event_name: str, **_: 'dict[str, Any]') -> N
     for key, value in extra_headers.items():
         request.headers[key] = value
 
-
 class InvokeResponseWrapper(ObjectProxy): # type: ignore
+    _cohere_embed_english_v3_tokenizer: Optional['Tokenizer'] = None
+
     def __init__(
         self,
-        response: Any,
+        response: 'dict[str, Any]',
+        body: Any,
         request: '_BedrockInvokeProviderRequest',
         log_prompt_and_response: bool
         ) -> None:
 
-        super().__init__(response) # type: ignore
+        super().__init__(body) # type: ignore
         self._response = response
+        self._body = body
         self._request = request
         self._log_prompt_and_response = log_prompt_and_response
 
@@ -160,14 +207,50 @@ class InvokeResponseWrapper(ObjectProxy): # type: ignore
 
             bedrock_converse_process_synchronous_function_call(self._request, response)
 
+        elif self._request._is_amazon_titan_embed_text_v1:
+            input = response.get('inputTextTokenCount', 0)
+            units["text"] = Units(input=input, output=0)
+
+        elif self._request._is_cohere_embed_english_v3:
+            texts: list[str] = response.get("texts", [])
+            if texts and len(texts) > 0:
+                text = " ".join(texts)
+
+                try:
+                    from tokenizers import Tokenizer  # type: ignore
+
+                    if self._cohere_embed_english_v3_tokenizer is None: # type: ignore
+                        current_dir = os.path.dirname(os.path.abspath(__file__))
+                        tokenizer_path = os.path.join(current_dir, "data", "cohere_embed_english_v3.json")
+                        self._cohere_embed_english_v3_tokenizer = Tokenizer.from_file(tokenizer_path) # type: ignore
+
+                    if self._cohere_embed_english_v3_tokenizer is not None and isinstance(self._cohere_embed_english_v3_tokenizer, Tokenizer): # type: ignore
+                        tokens: list = self._cohere_embed_english_v3_tokenizer.encode(text, add_special_tokens=False).tokens # type: ignore
+
+                        if tokens and isinstance(tokens, list):
+                            units["text"] = Units(input=len(tokens), output=0) # type: ignore
+
+                except ImportError:
+                    self._request._instrumentor._logger.warning("tokenizers module not found, caller must install the tokenizers module. Cannot record text tokens for Cohere embed english v3")
+                    pass
+                except Exception as e:
+                    self._request._instrumentor._logger.warning(f"Error processing Cohere embed english v3 response: {e}")
+                    pass
+
         if self._log_prompt_and_response:
             ingest["provider_response_json"] = data.decode('utf-8') # type: ignore
-            
-        self._request._instrumentor._ingest_units(self._request)
+
+        guardrails = response.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input", {})
+        self._request.process_guardrails(guardrails)
+
+        self._request.process_stop_action(response.get("amazon-bedrock-guardrailAction", ""))
+
+        xproxy_result = self._request._instrumentor._ingest_units(self._request)
+        self._request.assign_xproxy_result(self._response, xproxy_result)
 
         return data # type: ignore
 
-def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
+def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any, instance: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
@@ -176,14 +259,14 @@ def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
             _BedrockInvokeProviderRequest(instrumentor=instrumentor, model_id=modelId),
             _IsStreaming.false,
             wrapped,
-            None,
+            instance,
             args,
             kwargs,
         )   
     
     return invoke_wrapper
 
-def wrap_invoke_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
+def wrap_invoke_stream(instrumentor: _PayiInstrumentor, wrapped: Any, instance: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
         modelId: str = kwargs.get("modelId", "") # type: ignore
@@ -193,14 +276,14 @@ def wrap_invoke_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
             _BedrockInvokeProviderRequest(instrumentor=instrumentor, model_id=modelId),
             _IsStreaming.true,
             wrapped,
-            None,
+            instance,
             args,
             kwargs,
         )
 
     return invoke_wrapper
 
-def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
+def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any, instance: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
@@ -210,14 +293,14 @@ def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
             _BedrockConverseProviderRequest(instrumentor=instrumentor),
             _IsStreaming.false,
             wrapped,
-            None,
+            instance,
             args,
             kwargs,
         )
     
     return invoke_wrapper
 
-def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
+def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any, instance: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
         modelId: str = kwargs.get("modelId", "") # type: ignore
@@ -227,7 +310,7 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
             _BedrockConverseProviderRequest(instrumentor=instrumentor),
             _IsStreaming.true,
             wrapped,
-            None,
+            instance,
             args,
             kwargs,
         )
@@ -235,6 +318,7 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
     return invoke_wrapper
 
 class _BedrockProviderRequest(_ProviderRequest):
+
     def __init__(self, instrumentor: _PayiInstrumentor):
         super().__init__(
             instrumentor=instrumentor,
@@ -246,15 +330,40 @@ class _BedrockProviderRequest(_ProviderRequest):
             )
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
-        # boto3 doesn't allow extra_headers
-        kwargs.pop("extra_headers", None)
-        self._ingest["resource"] = kwargs.get("modelId", "")
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]',  args: Sequence[Any], kwargs: Any) -> bool:
+        modelId =  kwargs.get("modelId", "")
+        self._ingest["resource"] = modelId
+
+        if not self._price_as.resource and not self._price_as.category and BedrockInstrumentor._model_mapping:
+            deployment = BedrockInstrumentor._model_mapping.get(modelId, {})
+            self._price_as.category = deployment.get("price_as_category", "")
+            self._price_as.resource = deployment.get("price_as_resource", "")
+            self._price_as.resource_scope = deployment.get("resource_scope", None)
+
+        if self._price_as.resource_scope:
+            self._ingest["resource_scope"] = self._price_as.resource_scope
+        
+        # override defaults
+        if self._price_as.category:
+            self._ingest["category"] = self._price_as.category
+        if self._price_as.resource:
+            self._ingest["resource"] = self._price_as.resource
+
         return True
 
+    def process_response_metadata(self, metadata: 'dict[str, Any]') -> None:
+        request_id = metadata.get("RequestId", "")
+        if request_id:
+            self._ingest["provider_response_id"] = request_id
+
+        response_headers = metadata.get("HTTPHeaders", {})
+        if response_headers:
+            self.add_response_headers(response_headers)
+
     @override
     def process_initial_stream_response(self, response: Any) -> None:
-        self._ingest["provider_response_id"] = response.get("ResponseMetadata", {}).get("RequestId", None)
+        super().process_initial_stream_response(response)
+        self.process_response_metadata(response.get("ResponseMetadata", {}))
 
     @override
     def process_exception(self, exception: Exception, kwargs: Any, ) -> bool:
@@ -281,12 +390,73 @@ class _BedrockProviderRequest(_ProviderRequest):
             self._instrumentor._logger.debug(f"Error processing exception: {e}")
             return False
 
+    def process_guardrails(self, guardrails: 'dict[str, Any]') -> None:
+        units = self._ingest["units"]
+
+        # while we iterate over the entire dict, only one guardrail is expected and supported
+        for _, value in guardrails.items():
+            # _ (key) is the guardrail id
+            if not isinstance(value, dict):
+                continue
+
+            usage: dict[str, int] = value.get("invocationMetrics", {}).get("usage", {}) # type: ignore
+            if not usage:
+                continue
+
+            topicPolicyUnits: int  = usage.get("topicPolicyUnits", 0) # type: ignore
+            if topicPolicyUnits > 0:
+                units["guardrail_topic"] = Units(input=topicPolicyUnits, output=0) # type: ignore
+
+            contentPolicyUnits = usage.get("contentPolicyUnits", 0) # type: ignore
+            if contentPolicyUnits > 0:
+                units["guardrail_content"] = Units(input=contentPolicyUnits, output=0) # type: ignore
+
+            wordPolicyUnits = usage.get("wordPolicyUnits", 0) # type: ignore    
+            if wordPolicyUnits > 0:
+                units["guardrail_word_free"] = Units(input=wordPolicyUnits, output=0) # type: ignore
+
+            automatedReasoningPolicyUnits = usage.get("automatedReasoningPolicyUnits", 0) # type: ignore
+            if automatedReasoningPolicyUnits > 0:
+                units["guardrail_automated_reasoning"] = Units(input=automatedReasoningPolicyUnits, output=0) # type: ignore
+
+            sensitiveInformationPolicyUnits = usage.get("sensitiveInformationPolicyUnits", 0) # type: ignore
+            if sensitiveInformationPolicyUnits > 0:
+                units["guardrail_sensitive_information"] = Units(input=sensitiveInformationPolicyUnits, output=0) # type: ignore
+
+            sensitiveInformationPolicyFreeUnits = usage.get("sensitiveInformationPolicyFreeUnits", 0) # type: ignore
+            if sensitiveInformationPolicyFreeUnits > 0:
+                units["guardrail_sensitive_information_free"] = Units(input=sensitiveInformationPolicyFreeUnits, output=0) # type: ignore
+
+            contextualGroundingPolicyUnits = usage.get("contextualGroundingPolicyUnits", 0) # type: ignore
+            if contextualGroundingPolicyUnits > 0:
+                units["guardrail_contextual_grounding"] = Units(input=contextualGroundingPolicyUnits, output=0) # type: ignore
+
+            contentPolicyImageUnits = usage.get("contentPolicyImageUnits", 0) # type: ignore
+            if contentPolicyImageUnits > 0:
+                units["guardrail_content_image"] = Units(input=contentPolicyImageUnits, output=0) # type: ignore
+
 class _BedrockInvokeProviderRequest(_BedrockProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor, model_id: str):
         super().__init__(instrumentor=instrumentor)
-        self._is_anthropic: bool = 'anthropic' in model_id
-        self._is_nova: bool = 'nova' in model_id
-        self._is_meta: bool = 'meta' in model_id
+
+        price_as_resource = BedrockInstrumentor._model_mapping.get(model_id, {}).get("price_as_resource", None)
+        if price_as_resource:
+            model_id = price_as_resource
+
+        self._is_anthropic: bool = False
+        self._is_nova: bool = False
+        self._is_meta: bool = False
+        self._is_amazon_titan_embed_text_v1: bool = False
+        self._is_cohere_embed_english_v3: bool = False
+
+        self._assign_model_state(model_id=model_id)
+
+    def _assign_model_state(self, model_id: str) -> None:
+        self._is_anthropic = 'anthropic' in model_id
+        self._is_nova = 'nova' in model_id
+        self._is_meta = 'meta' in model_id
+        self._is_amazon_titan_embed_text_v1 = 'amazon.titan-embed-text-v1' == model_id
+        self._is_cohere_embed_english_v3 = 'cohere.embed-english-v3' == model_id
 
     @override
     def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
@@ -294,21 +464,54 @@ class _BedrockInvokeProviderRequest(_BedrockProviderRequest):
 
         super().process_request(instance, extra_headers, args, kwargs)
     
+        # super().process_request will assign price_as mapping from global state, so evaluate afterwards
+        if self._price_as.resource:
+            self._assign_model_state(model_id=self._price_as.resource)
+
+        guardrail_id = kwargs.get("guardrailIdentifier", "")
+        if guardrail_id:
+            self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_id, guardrail_id)
+
+        guardrail_version = kwargs.get("guardrailVersion", "")
+        if guardrail_version:
+            self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_version, guardrail_version)
+
+        if guardrail_id and guardrail_version and BedrockInstrumentor._guardrail_trace:
+            trace = kwargs.get("trace", None)
+            if not trace:
+                kwargs["trace"] = "ENABLED"
+
         if self._is_anthropic:
             try:
-                body = json.loads( kwargs.get("body", ""))
+                body = json.loads(kwargs.get("body", ""))
                 messages = body.get("messages", {})
                 if messages:
                     anthropic_has_image_and_get_texts(self, messages)
             except Exception as e:
                 self._instrumentor._logger.debug(f"Bedrock invoke error processing request body: {e}")
-
+        elif self._is_cohere_embed_english_v3:
+            try:
+                body = json.loads(kwargs.get("body", ""))
+                input_type = body.get("input_type", "")
+                if input_type == 'image':
+                    images = body.get("images", [])
+                    if (len(images) > 0):
+                        # only supports one image according to docs
+                        self._ingest["units"]["vision"] = Units(input=1, output=0)
+            except Exception as e:
+                self._instrumentor._logger.debug(f"Bedrock invoke error processing request body: {e}")
         return True
 
     @override
     def process_chunk(self, chunk: Any) -> _ChunkResult:
         chunk_dict = json.loads(chunk)
 
+        guardrails = chunk_dict.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input", {})
+        if guardrails:
+            self.process_guardrails(guardrails)
+    
+        self.process_stop_action(chunk_dict.get("amazon-bedrock-guardrailAction", ""))
+
         if self._is_anthropic:
             from .AnthropicInstrumentor import anthropic_process_chunk
             return anthropic_process_chunk(self, chunk_dict, assign_id=False)
@@ -347,23 +550,23 @@ class _BedrockInvokeProviderRequest(_BedrockProviderRequest):
         log_prompt_and_response: bool,
         kwargs: Any) -> Any:
 
-        metadata = response.get("ResponseMetadata", {})
-
-        request_id = metadata.get("RequestId", "")
-        if request_id:
-            self._ingest["provider_response_id"] = request_id
-
-        response_headers = metadata.get("HTTPHeaders", {}).copy()
-        if response_headers:
-            self._ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+        self.process_response_metadata(response.get("ResponseMetadata", {}))
 
         response["body"] = InvokeResponseWrapper(
-            response=response["body"],
+            response=response,
+            body=response["body"],
             request=self,
             log_prompt_and_response=log_prompt_and_response)
 
         return response
 
+    def process_stop_action(self, action: str) -> None:
+        # record both as a semantic failure and guardrail action so it is discoverable through both properties
+        if action == "INTERVENED":
+            self.add_internal_request_property(PayiPropertyNames.failure, action)
+            self.add_internal_request_property(PayiPropertyNames.failure_description, GUARDRAIL_SEMANTIC_FAILURE_DESCRIPTION)
+            self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_action, action)
+
     @override
     def remove_inline_data(self, prompt: 'dict[str, Any]') -> bool:# noqa: ARG002
         if not self._is_anthropic:
@@ -383,6 +586,27 @@ class _BedrockInvokeProviderRequest(_BedrockProviderRequest):
         return False
 
 class _BedrockConverseProviderRequest(_BedrockProviderRequest):
+    @override
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
+        super().process_request(instance, extra_headers, args, kwargs)
+
+        guardrail_config = kwargs.get("guardrailConfig", {})
+        if guardrail_config:
+            guardrailIdentifier = guardrail_config.get("guardrailIdentifier", "")
+            if guardrailIdentifier:
+                self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_id, guardrailIdentifier)
+
+            guardrailVersion = guardrail_config.get("guardrailVersion", "")
+            if guardrailVersion:
+                self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_version, guardrailVersion)
+
+            if guardrailIdentifier and guardrailVersion and BedrockInstrumentor._guardrail_trace:
+                trace = guardrail_config.get("trace", None)
+                if not trace:
+                    guardrail_config["trace"] = "enabled"
+
+        return True
+
     @override
     def process_synchronous_response(
         self,
@@ -390,22 +614,14 @@ class _BedrockConverseProviderRequest(_BedrockProviderRequest):
         log_prompt_and_response: bool,
         kwargs: Any) -> Any:
 
-        usage = response["usage"]
-        input = usage["inputTokens"]
-        output = usage["outputTokens"]
-        
+        usage = response.get("usage", {})
+        input = usage.get("inputTokens", 0)
+        output = usage.get("outputTokens", 0)
+
         units: dict[str, Units] = self._ingest["units"]
         units["text"] = Units(input=input, output=output)
 
-        metadata = response.get("ResponseMetadata", {})
-
-        request_id = metadata.get("RequestId", "")
-        if request_id:
-            self._ingest["provider_response_id"] = request_id
-
-        response_headers = metadata.get("HTTPHeaders", {})
-        if response_headers:
-            self._ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+        self.process_response_metadata(response.get("ResponseMetadata", {}))
 
         if log_prompt_and_response:
             response_without_metadata = response.copy()
@@ -414,6 +630,12 @@ class _BedrockConverseProviderRequest(_BedrockProviderRequest):
 
         bedrock_converse_process_synchronous_function_call(self, response)
 
+        guardrails = response.get("trace", {}).get("guardrail", {}).get("inputAssessment", {})
+        if guardrails:
+            self.process_guardrails(guardrails)
+
+        self.process_stop_reason(response.get("stopReason", ""))
+
         return None
 
     @override
@@ -422,17 +644,30 @@ class _BedrockConverseProviderRequest(_BedrockProviderRequest):
         metadata = chunk.get("metadata", {})
 
         if metadata:
-            usage = metadata['usage']
-            input = usage["inputTokens"]
-            output = usage["outputTokens"]
+            usage = metadata.get('usage', {})
+            input = usage.get("inputTokens", 0)
+            output = usage.get("outputTokens", 0)
             self._ingest["units"]["text"] = Units(input=input, output=output)
 
+            guardrail = metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment", {})
+            if guardrail:
+                self.process_guardrails(guardrail)
+
             ingest = True
 
+        self.process_stop_reason(chunk.get("messageStop", {}).get("stopReason", ""))
+
         bedrock_converse_process_streaming_for_function_call(self, chunk)
 
         return _ChunkResult(send_chunk_to_caller=True, ingest=ingest)
 
+    def process_stop_reason(self, reason: str) -> None:
+        if reason == "guardrail_intervened":
+            # record both as a semantic failure and guardrail action so it is discoverable through both properties
+            self.add_internal_request_property(PayiPropertyNames.failure, reason)
+            self.add_internal_request_property(PayiPropertyNames.failure_description, GUARDRAIL_SEMANTIC_FAILURE_DESCRIPTION)
+            self.add_internal_request_property(PayiPropertyNames.aws_bedrock_guardrail_action, reason)
+
 def bedrock_converse_process_streaming_for_function_call(request: _ProviderRequest, chunk: 'dict[str, Any]') -> None:  
     contentBlockStart = chunk.get("contentBlockStart", {})
     tool_use = contentBlockStart.get("start", {}).get("toolUse", {})

payi/lib/GoogleGenAiInstrumentor.py

@@ -1,11 +1,14 @@
+from __future__ import annotations
+
 from typing import Any, List, Union, Sequence
 from typing_extensions import override
 
 from wrapt import wrap_function_wrapper  # type: ignore
 
-from .instrument import _ChunkResult, _IsStreaming, _PayiInstrumentor
+from .instrument import _IsStreaming, _PayiInstrumentor
 from .VertexRequest import _VertexRequest
 from .version_helper import get_version_helper
+from .ProviderRequest import _ChunkResult
 
 
 class GoogleGenAiInstrumentor:
@@ -14,36 +17,20 @@ class GoogleGenAiInstrumentor:
 
     @staticmethod
     def instrument(instrumentor: _PayiInstrumentor) -> None:
-        try:
-            GoogleGenAiInstrumentor._module_version = get_version_helper(GoogleGenAiInstrumentor._module_name)
-
-            wrap_function_wrapper(
-                "google.genai.models",
-                "Models.generate_content",
-                generate_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "google.genai.models",
-                "Models.generate_content_stream",
-                generate_stream_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "google.genai.models",
-                "AsyncModels.generate_content",
-                agenerate_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "google.genai.models",
-                "AsyncModels.generate_content_stream",
-                agenerate_stream_wrapper(instrumentor),
-            )
-
-        except Exception as e:
-            instrumentor._logger.debug(f"Error instrumenting vertex: {e}")
-            return
+        GoogleGenAiInstrumentor._module_version = get_version_helper(GoogleGenAiInstrumentor._module_name)
+
+        wrappers = [
+            ("google.genai.models", "Models.generate_content", generate_wrapper(instrumentor)),
+            ("google.genai.models", "Models.generate_content_stream", generate_stream_wrapper(instrumentor)),
+            ("google.genai.models", "AsyncModels.generate_content", agenerate_wrapper(instrumentor)),
+            ("google.genai.models", "AsyncModels.generate_content_stream", agenerate_stream_wrapper(instrumentor)),
+        ]
+
+        for module, method, wrapper in wrappers:
+            try:
+                wrap_function_wrapper(module, method, wrapper)
+            except Exception as e:
+                instrumentor._logger.debug(f"Error wrapping {module}.{method}: {e}")
 
 @_PayiInstrumentor.payi_wrapper
 def generate_wrapper(