payi 0.1.0a68__py3-none-any.whl → 0.1.0a70__py3-none-any.whl

payi/lib/AnthropicInstrumentor.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Union
+from typing import Any, Union, Optional
 from typing_extensions import override
 
 import tiktoken
@@ -54,7 +54,6 @@ def chat_wrapper(
     **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        "system.anthropic",
         _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -72,7 +71,6 @@ async def achat_wrapper(
     **kwargs: Any,
 ) -> Any:
     return await instrumentor.achat_wrapper(
-        "system.anthropic",
         _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -82,6 +80,9 @@ async def achat_wrapper(
     )
 
 class _AnthropicProviderRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor=instrumentor, category="system.anthropic")
+
     @override
     def process_chunk(self, chunk: Any) -> bool:
         if chunk.type == "message_start":
@@ -135,25 +136,56 @@ class _AnthropicProviderRequest(_ProviderRequest):
         return None
 
     @override
-    def process_request(self, kwargs: Any) -> None:
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+        self._ingest["resource"] = kwargs.get("model", "")
         messages = kwargs.get("messages")
-        if not messages or len(messages) == 0:
-            return
-        
-        estimated_token_count = 0 
-        has_image = False
+        if messages:
+            estimated_token_count = 0 
+            has_image = False
+
+            enc = tiktoken.get_encoding("cl100k_base")
+            
+            for message in messages:
+                msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
+                if msg_has_image:
+                    has_image = True
+                    estimated_token_count += msg_prompt_tokens
+            
+            if has_image and estimated_token_count > 0:
+                self._estimated_prompt_tokens = estimated_token_count
+        return True
+
+    @override
+    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool:
+        try:
+            status_code: Optional[int] = None
+
+            if hasattr(exception, "status_code"):
+                status_code = getattr(exception, "status_code", None)
+                if isinstance(status_code, int):
+                    self._ingest["http_status_code"] = status_code
+
+            if not status_code:
+                return False
+
+            if hasattr(exception, "request_id"):
+                request_id = getattr(exception, "request_id", None)
+                if isinstance(request_id, str):
+                    self._ingest["provider_response_id"] = request_id
+
+            if hasattr(exception, "response"):
+                response = getattr(exception, "response", None)
+                if hasattr(response, "text"):
+                    text = getattr(response, "text", None)
+                    if isinstance(text, str):
+                        self._ingest["provider_response_json"] = text
+
+        except Exception as e:
+            logging.debug(f"Error processing exception: {e}")
+            return False
+
+        return True
 
-        enc = tiktoken.get_encoding("cl100k_base")
-        
-        for message in messages:
-            msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
-            if msg_has_image:
-                has_image = True
-                estimated_token_count += msg_prompt_tokens
-        
-        if not has_image or estimated_token_count == 0:
-            return
-        self._estimated_prompt_tokens = estimated_token_count
 
 def has_image_and_get_texts(encoding: tiktoken.Encoding, content: Union[str, 'list[Any]']) -> 'tuple[bool, int]':
     if isinstance(content, str):

payi/lib/BedrockInstrumentor.py

@@ -1,3 +1,4 @@
+import os
 import json
 import logging
 from typing import Any
@@ -6,15 +7,21 @@ from typing_extensions import override
 
 from wrapt import ObjectProxy, wrap_function_wrapper  # type: ignore
 
+from payi.lib.helpers import PayiHeaderNames, payi_aws_bedrock_url
 from payi.types.ingest_units_params import Units, IngestUnitsParams
 from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
 from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
 
+_supported_model_prefixes = ["meta.llama3", "anthropic.", "amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro"]
 
 class BedrockInstrumentor:
+    _instrumentor: _PayiInstrumentor
+
     @staticmethod
     def instrument(instrumentor: _PayiInstrumentor) -> None:
+        BedrockInstrumentor._instrumentor = instrumentor
+
         try:
             import boto3  # type: ignore #  noqa: F401  I001
 
@@ -46,12 +53,48 @@ def create_client_wrapper(instrumentor: _PayiInstrumentor, wrapped: Any, instanc
         client.converse = wrap_converse(instrumentor, client.converse)
         client.converse_stream = wrap_converse_stream(instrumentor, client.converse_stream)
 
+        if BedrockInstrumentor._instrumentor._proxy_default:
+            # Register client callbacks to handle the Pay-i extra_headers parameter in the inference calls and redirect the request to the Pay-i endpoint
+            _register_bedrock_client_callbacks(client)
+
         return client
     except Exception as e:
         logging.debug(f"Error instrumenting bedrock client: {e}")
     
     return wrapped(*args, **kwargs)
+
+BEDROCK_REQUEST_NAMES = [
+    'request-created.bedrock-runtime.Converse',
+    'request-created.bedrock-runtime.ConverseStream',
+    'request-created.bedrock-runtime.InvokeModel',
+    'request-created.bedrock-runtime.InvokeModelWithResponseStream',
+]
+
+def _register_bedrock_client_callbacks(client: Any) -> None:
+    # Pass a unqiue_id to avoid registering the same callback multiple times in case this cell executed more than once
+    # Redirect the request to the Pay-i endpoint after the request has been signed. 
+    client.meta.events.register_last('request-created', _redirect_to_payi, unique_id=_redirect_to_payi)
+
+def _redirect_to_payi(request: Any, event_name: str, **_: 'dict[str, Any]') -> None:
+    from urllib3.util import parse_url
+    from urllib3.util.url import Url
+
+    if not event_name in BEDROCK_REQUEST_NAMES:
+        return
     
+    parsed_url: Url = parse_url(request.url)
+    route_path = parsed_url.path
+    request.url = f"{payi_aws_bedrock_url()}{route_path}"
+
+    request.headers[PayiHeaderNames.api_key] = os.environ.get("PAYI_API_KEY", "")
+    request.headers[PayiHeaderNames.provider_base_uri] = parsed_url.scheme + "://" + parsed_url.host # type: ignore
+    
+    extra_headers = BedrockInstrumentor._instrumentor._create_extra_headers()
+
+    for key, value in extra_headers.items():
+        request.headers[key] = value
+
+
 class InvokeResponseWrapper(ObjectProxy): # type: ignore
     def __init__(
         self,
@@ -67,10 +110,10 @@ class InvokeResponseWrapper(ObjectProxy): # type: ignore
         self._ingest = ingest
         self._log_prompt_and_response = log_prompt_and_response
 
-    def read(self, amt: Any =None): # type: ignore
+    def read(self, amt: Any =None) -> Any: # type: ignore
         # data is array of bytes
-        data: Any = self.__wrapped__.read(amt) # type: ignore
-        response = json.loads(data)
+        data: bytes = self.__wrapped__.read(amt) # type: ignore
+        response = json.loads(data) # type: ignore
 
         resource = self._ingest["resource"]
         if not resource:
@@ -90,27 +133,29 @@ class InvokeResponseWrapper(ObjectProxy): # type: ignore
         units["text"] = Units(input=input, output=output)
 
         if self._log_prompt_and_response:
-            self._ingest["provider_response_json"] = data.decode('utf-8')
+            self._ingest["provider_response_json"] = data.decode('utf-8') # type: ignore
             
         self._instrumentor._ingest_units(self._ingest)
 
-        return data
+        return data # type: ignore
+
+def _is_supported_model(modelId: str) -> bool:
+    return any(prefix in modelId for prefix in _supported_model_prefixes)
 
 def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
 
-        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+        if _is_supported_model(modelId):
             return instrumentor.chat_wrapper(
-                "system.aws.bedrock",
-                _BedrockInvokeSynchronousProviderRequest(instrumentor),
+                _BedrockInvokeSynchronousProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.false,
                 wrapped,
                 None,
                 args,
                 kwargs,
-        )
+            )   
         return wrapped(*args, **kwargs)
     
     return invoke_wrapper
@@ -118,12 +163,11 @@ def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 def wrap_invoke_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
-        model_id: str = kwargs.get("modelId", "") # type: ignore
+        modelId: str = kwargs.get("modelId", "") # type: ignore
 
-        if model_id.startswith("meta.llama3") or model_id.startswith("anthropic."):
+        if _is_supported_model(modelId):
             return instrumentor.chat_wrapper(
-                "system.aws.bedrock",
-                _BedrockInvokeStreamingProviderRequest(instrumentor, model_id),
+                _BedrockInvokeStreamingProviderRequest(instrumentor=instrumentor, model_id=modelId),
                 _IsStreaming.true,
                 wrapped,
                 None,
@@ -139,10 +183,9 @@ def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
     def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
 
-        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+        if _is_supported_model(modelId):
             return instrumentor.chat_wrapper(
-                "system.aws.bedrock",
-                _BedrockConverseSynchronousProviderRequest(instrumentor),
+                _BedrockConverseSynchronousProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.false,
                 wrapped,
                 None,
@@ -156,12 +199,11 @@ def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
     @wraps(wrapped)
     def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
-        model_id: str = kwargs.get("modelId", "") # type: ignore
+        modelId: str = kwargs.get("modelId", "") # type: ignore
 
-        if model_id.startswith("meta.llama3") or model_id.startswith("anthropic."):
+        if _is_supported_model(modelId):
             return instrumentor.chat_wrapper(
-                "system.aws.bedrock",
-                _BedrockConverseStreamingProviderRequest(instrumentor),
+                _BedrockConverseStreamingProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.true,
                 wrapped,
                 None,
@@ -172,9 +214,45 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 
     return invoke_wrapper
 
-class _BedrockInvokeStreamingProviderRequest(_ProviderRequest):
+class _BedrockProviderRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor=instrumentor, category="system.aws.bedrock")
+
+    @override
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+        # boto3 doesn't allow extra_headers
+        kwargs.pop("extra_headers", None)
+        self._ingest["resource"] = kwargs.get("modelId", "")
+        return True
+
+    @override
+    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool:
+        try:
+            if hasattr(exception, "response"):
+                response: dict[str, Any] = getattr(exception, "response", {})
+                status_code: int = response.get('ResponseMetadata', {}).get('HTTPStatusCode', 0)
+                if status_code == 0:
+                    return False
+
+                self._ingest["http_status_code"] = status_code
+                
+                request_id = response.get('ResponseMetadata', {}).get('RequestId', "")
+                if request_id:
+                    self._ingest["provider_response_id"] = request_id
+
+                error = response.get('Error', "")
+                if error:
+                    self._ingest["provider_response_json"] = json.dumps(error)
+
+            return True
+
+        except Exception as e:
+            logging.debug(f"Error processing exception: {e}")
+            return False
+
+class _BedrockInvokeStreamingProviderRequest(_BedrockProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor, model_id: str):
-        super().__init__(instrumentor)
+        super().__init__(instrumentor=instrumentor)
         self._is_anthropic: bool = model_id.startswith("anthropic.")
 
     @override
@@ -220,7 +298,7 @@ class _BedrockInvokeStreamingProviderRequest(_ProviderRequest):
 
         return True
 
-class _BedrockInvokeSynchronousProviderRequest(_ProviderRequest):
+class _BedrockInvokeSynchronousProviderRequest(_BedrockProviderRequest):
     @override
     def process_synchronous_response(
         self,
@@ -246,7 +324,7 @@ class _BedrockInvokeSynchronousProviderRequest(_ProviderRequest):
 
         return response
 
-class _BedrockConverseSynchronousProviderRequest(_ProviderRequest):
+class _BedrockConverseSynchronousProviderRequest(_BedrockProviderRequest):
     @override
     def process_synchronous_response(
         self,
@@ -278,7 +356,7 @@ class _BedrockConverseSynchronousProviderRequest(_ProviderRequest):
 
         return None    
 
-class _BedrockConverseStreamingProviderRequest(_ProviderRequest):
+class _BedrockConverseStreamingProviderRequest(_BedrockProviderRequest):
     @override
     def process_chunk(self, chunk: 'dict[str, Any]') -> bool:
         metadata = chunk.get("metadata", {})

payi/lib/OpenAIInstrumentor.py

@@ -8,6 +8,7 @@ import tiktoken  # type: ignore
 from wrapt import wrap_function_wrapper  # type: ignore
 
 from payi.types import IngestUnitsParams
+from payi.lib.helpers import PayiCategories, PayiHeaderNames
 from payi.types.ingest_units_params import Units
 
 from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
@@ -63,7 +64,6 @@ def embeddings_wrapper(
     **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        "system.openai",
         _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
@@ -81,7 +81,6 @@ async def aembeddings_wrapper(
     **kwargs: Any,
 ) -> Any:
     return await instrumentor.achat_wrapper(
-        "system.openai",
         _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
@@ -99,7 +98,6 @@ def chat_wrapper(
     **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        "system.openai",
         _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -117,7 +115,6 @@ async def achat_wrapper(
     **kwargs: Any,
 ) -> Any:
     return await instrumentor.achat_wrapper(
-        "system.openai",
         _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -126,7 +123,79 @@ async def achat_wrapper(
         kwargs,
     )
 
-class _OpenAiEmbeddingsProviderRequest(_ProviderRequest):
+class _OpenAiProviderRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor=instrumentor, category="system.openai")
+
+    @override
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+        self._ingest["resource"] = kwargs.get("model", "")
+
+        if not (instance and hasattr(instance, "_client")) or OpenAiInstrumentor.is_azure(instance) is False:
+            return True
+
+        context = self._instrumentor.get_context_safe()
+        route_as_resource = extra_headers.get(PayiHeaderNames.route_as_resource) or context.get("route_as_resource")
+        resource_scope = extra_headers.get(PayiHeaderNames.resource_scope) or context.get("resource_scope")
+
+        if PayiHeaderNames.route_as_resource in extra_headers:
+            del extra_headers[PayiHeaderNames.route_as_resource]
+        if PayiHeaderNames.resource_scope in extra_headers:
+            del extra_headers[PayiHeaderNames.resource_scope]
+            
+        if not route_as_resource:
+            logging.error("Azure OpenAI route as resource not found, not ingesting")
+            return False
+
+        if resource_scope:
+            if not(resource_scope in ["global", "datazone"] or resource_scope.startswith("region")):
+                logging.error("Azure OpenAI invalid resource scope, not ingesting")
+                return False
+
+            self._ingest["resource_scope"] = resource_scope
+
+        self._category = PayiCategories.azure_openai
+
+        self._ingest["category"] = self._category
+        self._ingest["resource"] = route_as_resource
+ 
+        return True
+
+    @override
+    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool:
+        try:
+            status_code: Optional[int] = None
+
+            if hasattr(exception, "status_code"):
+                status_code = getattr(exception, "status_code", None)
+                if isinstance(status_code, int):
+                    self._ingest["http_status_code"] = status_code
+
+            if not status_code:
+                return False
+
+            if hasattr(exception, "request_id"):
+                request_id = getattr(exception, "request_id", None)
+                if isinstance(request_id, str):
+                    self._ingest["provider_response_id"] = request_id
+
+            if hasattr(exception, "response"):
+                response = getattr(exception, "response", None)
+                if hasattr(response, "text"):
+                    text = getattr(response, "text", None)
+                    if isinstance(text, str):
+                        self._ingest["provider_response_json"] = text
+
+        except Exception as e:
+            logging.debug(f"Error processing exception: {e}")
+            return False
+
+        return True
+
+class _OpenAiEmbeddingsProviderRequest(_OpenAiProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor=instrumentor)
+
     @override
     def process_synchronous_response(
         self,
@@ -135,9 +204,9 @@ class _OpenAiEmbeddingsProviderRequest(_ProviderRequest):
         kwargs: Any) -> Any:
         return process_chat_synchronous_response(response, self._ingest, log_prompt_and_response, self._estimated_prompt_tokens)
 
-class _OpenAiChatProviderRequest(_ProviderRequest):
+class _OpenAiChatProviderRequest(_OpenAiProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor)
+        super().__init__(instrumentor=instrumentor)
         self._include_usage_added = False
 
     @override
@@ -163,39 +232,42 @@ class _OpenAiChatProviderRequest(_ProviderRequest):
         return send_chunk_to_client
 
     @override
-    def process_request(self, kwargs: Any) -> None: # noqa: ARG001
-        messages = kwargs.get("messages", None)
-        if not messages or len(messages) == 0:
-            return
-        
-        estimated_token_count = 0 
-        has_image = False
-
-        try: 
-            enc = tiktoken.encoding_for_model(kwargs.get("model")) # type: ignore
-        except KeyError:
-            enc = tiktoken.get_encoding("o200k_base") # type: ignore
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+        result = super().process_request(instance, extra_headers, kwargs)
+        if result is False:
+            return result
         
-        for message in messages:
-            msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
-            if msg_has_image:
-                has_image = True
-                estimated_token_count += msg_prompt_tokens
-        
-        if has_image and estimated_token_count > 0:
-            self._estimated_prompt_tokens = estimated_token_count
+        messages = kwargs.get("messages", None)
+        if messages:
+            estimated_token_count = 0 
+            has_image = False
+
+            try: 
+                enc = tiktoken.encoding_for_model(kwargs.get("model")) # type: ignore
+            except KeyError:
+                enc = tiktoken.get_encoding("o200k_base") # type: ignore
+            
+            for message in messages:
+                msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
+                if msg_has_image:
+                    has_image = True
+                    estimated_token_count += msg_prompt_tokens
+            
+            if has_image and estimated_token_count > 0:
+                self._estimated_prompt_tokens = estimated_token_count
 
-        stream: bool = kwargs.get("stream", False)
-        if stream:
-            add_include_usage = True
+            stream: bool = kwargs.get("stream", False)
+            if stream:
+                add_include_usage = True
 
-            stream_options: dict[str, Any] = kwargs.get("stream_options", None)
-            if stream_options and "include_usage" in stream_options:
-                add_include_usage = stream_options["include_usage"] == False
+                stream_options: dict[str, Any] = kwargs.get("stream_options", None)
+                if stream_options and "include_usage" in stream_options:
+                    add_include_usage = stream_options["include_usage"] == False
 
-            if add_include_usage:
-                kwargs['stream_options'] = {"include_usage": True}
-                self._include_usage_added = True
+                if add_include_usage:
+                    kwargs['stream_options'] = {"include_usage": True}
+                    self._include_usage_added = True
+        return True
 
     @override
     def process_synchronous_response(

payi/lib/helpers.py

@@ -15,7 +15,8 @@ class PayiHeaderNames:
     route_as_resource:str = "xProxy-RouteAs-Resource"
     provider_base_uri = "xProxy-Provider-BaseUri"
     resource_scope:str = "xProxy-Resource-Scope"
-
+    api_key:str = "xProxy-Api-Key"
+    
 class PayiCategories:
     anthropic:str  = "system.anthropic"
     openai:str = "system.openai"