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

payi/lib/OpenAIInstrumentor.py

@@ -1,22 +1,32 @@
+from __future__ import annotations
+
 import json
-from typing import Any, Union, Optional, Sequence
+from typing import Any, Dict, Union, Optional, Sequence
 from typing_extensions import override
 from importlib.metadata import version
 
 import tiktoken  # type: ignore
 from wrapt import wrap_function_wrapper  # type: ignore
 
-from payi.lib.helpers import PayiCategories, PayiHeaderNames
+from payi.lib.helpers import PayiCategories
 from payi.types.ingest_units_params import Units
 
-from .instrument import _ChunkResult, _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
+from .instrument import (
+    PayiInstrumentAzureOpenAiConfig,
+    _Context,
+    _IsStreaming,
+    _PayiInstrumentor,
+)
 from .version_helper import get_version_helper
+from .ProviderRequest import _ChunkResult, _StreamingType, _ProviderRequest
 
 
 class OpenAiInstrumentor:
     _module_name: str = "openai"
     _module_version: str = ""
 
+    _azure_openai_deployments: Dict[str, _Context] = {}
+
     @staticmethod
     def is_azure(instance: Any) -> bool:
         from openai import AzureOpenAI, AsyncAzureOpenAI # type: ignore # noqa: I001
@@ -24,52 +34,31 @@ class OpenAiInstrumentor:
         return isinstance(instance._client, (AsyncAzureOpenAI, AzureOpenAI))
 
     @staticmethod
-    def instrument(instrumentor: _PayiInstrumentor) -> None:
-        try:
-            OpenAiInstrumentor._module_version = get_version_helper(OpenAiInstrumentor._module_name)
-
-            wrap_function_wrapper(
-                "openai.resources.chat.completions",
-                "Completions.create",
-                chat_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "openai.resources.chat.completions",
-                "AsyncCompletions.create",
-                achat_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "openai.resources.embeddings",
-                "Embeddings.create",
-                embeddings_wrapper(instrumentor),
-            )
-
-            wrap_function_wrapper(
-                "openai.resources.embeddings",
-                 "AsyncEmbeddings.create",
-                aembeddings_wrapper(instrumentor),
-            )
-        except Exception as e:
-            instrumentor._logger.debug(f"Error instrumenting openai: {e}")
-
-        # responses separately as they are relatively new and the client may not be using the latest openai module
-        try:            
-            wrap_function_wrapper(
-                "openai.resources.responses",
-                "Responses.create",
-                responses_wrapper(instrumentor),
-            )
+    def configure(azure_openai_config: Optional[PayiInstrumentAzureOpenAiConfig]) -> None:
+        if azure_openai_config:
+            model_mappings = azure_openai_config.get("model_mappings", [])
+            OpenAiInstrumentor._azure_openai_deployments = _PayiInstrumentor._model_mapping_to_context_dict(model_mappings)
 
-            wrap_function_wrapper(
-                "openai.resources.responses",
-                "AsyncResponses.create",
-                aresponses_wrapper(instrumentor),
-            )
-
-        except Exception as e:
-            instrumentor._logger.debug(f"Error instrumenting openai: {e}")
+    @staticmethod
+    def instrument(instrumentor: _PayiInstrumentor) -> None:
+        OpenAiInstrumentor._module_version = get_version_helper(OpenAiInstrumentor._module_name)
+
+        wrappers = [
+            ("openai._base_client", "AsyncAPIClient._process_response", _ProviderRequest.aprocess_response_wrapper),
+            ("openai._base_client", "SyncAPIClient._process_response", _ProviderRequest.process_response_wrapper),
+            ("openai.resources.chat.completions", "Completions.create", chat_wrapper(instrumentor)),
+            ("openai.resources.chat.completions", "AsyncCompletions.create", achat_wrapper(instrumentor)),
+            ("openai.resources.embeddings", "Embeddings.create", embeddings_wrapper(instrumentor)),
+            ("openai.resources.embeddings", "AsyncEmbeddings.create", aembeddings_wrapper(instrumentor)),
+            ("openai.resources.responses", "Responses.create", responses_wrapper(instrumentor)),
+            ("openai.resources.responses", "AsyncResponses.create", aresponses_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 embeddings_wrapper(
@@ -201,44 +190,39 @@ class _OpenAiProviderRequest(_ProviderRequest):
         self._input_tokens_details_key = input_tokens_details_key
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool: # type: ignore
-        self._ingest["resource"] = kwargs.get("model", "")
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]',  args: Sequence[Any], kwargs: Any) -> bool: # type: ignore
+        model = kwargs.get("model", "")
 
         if not (instance and hasattr(instance, "_client")) or OpenAiInstrumentor.is_azure(instance) is False:
+            self._ingest["resource"] = model
             return True
 
-        context = self._instrumentor.get_context_safe()
-        price_as_category = extra_headers.get(PayiHeaderNames.price_as_category) or context.get("price_as_category")
-        price_as_resource = extra_headers.get(PayiHeaderNames.price_as_resource) or context.get("price_as_resource")
-        resource_scope = extra_headers.get(PayiHeaderNames.resource_scope) or context.get("resource_scope")
-
-        if PayiHeaderNames.price_as_category in extra_headers:
-            del extra_headers[PayiHeaderNames.price_as_category]
-        if PayiHeaderNames.price_as_resource in extra_headers:
-            del extra_headers[PayiHeaderNames.price_as_resource]
-        if PayiHeaderNames.resource_scope in extra_headers:
-            del extra_headers[PayiHeaderNames.resource_scope]
-            
-        if not price_as_resource and not price_as_category:
+        if not self._price_as.resource and not self._price_as.category and OpenAiInstrumentor._azure_openai_deployments:
+            deployment = OpenAiInstrumentor._azure_openai_deployments.get(model, {})
+            self._price_as.category = deployment.get("price_as_category", None)
+            self._price_as.resource = deployment.get("price_as_resource", None)
+            self._price_as.resource_scope = deployment.get("resource_scope", None)
+
+        if not self._price_as.resource and not self._price_as.category:
             self._instrumentor._logger.error("Azure OpenAI requires price as resource and/or category to be specified, not ingesting")
             return False
 
-        if resource_scope:
-            if not(resource_scope in ["global", "datazone"] or resource_scope.startswith("region")):
+        if self._price_as.resource_scope:
+            if not (self._price_as.resource_scope in ["global", "datazone"] or self._price_as.resource_scope.startswith("region")):
                 self._instrumentor._logger.error("Azure OpenAI invalid resource scope, not ingesting")
                 return False
 
-            self._ingest["resource_scope"] = resource_scope
+            self._ingest["resource_scope"] = self._price_as.resource_scope
 
         self._category = PayiCategories.azure_openai
 
         self._ingest["category"] = self._category
 
-        if price_as_category:
+        if self._price_as.category:
             # price as category overrides default
-            self._ingest["category"] = price_as_category
-        if price_as_resource:
-            self._ingest["resource"] = price_as_resource
+            self._ingest["category"] = self._price_as.category
+        if self._price_as.resource:
+            self._ingest["resource"] = self._price_as.resource
 
         return True
 
@@ -304,7 +288,7 @@ class _OpenAiProviderRequest(_ProviderRequest):
             if input_cache != 0:
                 units["text_cache_read"] = Units(input=input_cache, output=0)
 
-        input = _PayiInstrumentor.update_for_vision(input - input_cache, units, self._estimated_prompt_tokens)
+        input = self.update_for_vision(input - input_cache)
 
         units["text"] = Units(input=input, output=output)
 
@@ -614,4 +598,4 @@ def model_to_dict(model: Any) -> Any:
     elif hasattr(model, "parse"):  # Raw API response
         return model_to_dict(model.parse())
     else:
-        return model        
+        return model

payi/lib/ProviderRequest.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import inspect
+from abc import abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Optional, Sequence
+from dataclasses import dataclass
+
+from payi.types import IngestUnitsParams
+from payi.lib.helpers import PayiPropertyNames
+from payi.types.ingest_units_params import ProviderResponseFunctionCall
+from payi.types.shared.xproxy_error import XproxyError
+from payi.types.shared.xproxy_result import XproxyResult
+from payi.types.shared_params.ingest_units import IngestUnits
+from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
+
+from .helpers import _set_attr_safe
+
+if TYPE_CHECKING:
+    from .instrument import _PayiInstrumentor
+
+class _StreamingType(Enum):
+    generator = 0
+    iterator = 1
+    stream_manager = 2
+
+@dataclass
+class _ChunkResult:
+    send_chunk_to_caller: bool
+    ingest: bool = False
+
+@dataclass
+class PriceAs:
+    category: Optional[str]
+    resource: Optional[str]
+    resource_scope: Optional[str]
+
+class _ProviderRequest:
+    excluded_headers = {
+        "transfer-encoding",
+    }
+
+    _instrumented_response_headers_attr = "_instrumented_response_headers"
+    _xproxy_result_attr = "xproxy_result"
+
+    def __init__(
+            self,
+            instrumentor: _PayiInstrumentor,
+            category: str,
+            streaming_type: _StreamingType,
+            module_name: str,
+            module_version: str,
+            is_aws_client: Optional[bool] = None,
+            is_google_vertex_or_genai_client: Optional[bool] = None,
+            ) -> None:
+        self._instrumentor: _PayiInstrumentor = instrumentor
+        self._module_name: str = module_name
+        self._module_version: str = module_version  
+        self._estimated_prompt_tokens: Optional[int] = None
+        self._category: str = category
+        self._ingest: IngestUnitsParams = { "category": category, "units": {} } # type: ignore
+        self._streaming_type: '_StreamingType' = streaming_type
+        self._is_aws_client: Optional[bool] = is_aws_client
+        self._is_google_vertex_or_genai_client: Optional[bool] = is_google_vertex_or_genai_client
+        self._function_call_builder: Optional[dict[int, ProviderResponseFunctionCall]] = None
+        self._building_function_response: bool = False
+        self._function_calls: Optional[list[ProviderResponseFunctionCall]] = None
+        self._is_large_context: bool = False
+        self._internal_request_properties: dict[str, Optional[str]] = {}
+        self._price_as: PriceAs = PriceAs(category=None, resource=None, resource_scope=None)
+
+    def process_chunk(self, _chunk: Any) -> _ChunkResult:
+        return _ChunkResult(send_chunk_to_caller=True)
+
+    def process_synchronous_response(self, response: Any, log_prompt_and_response: bool, kwargs: Any) -> Optional[object]:  # noqa: ARG002
+        return None
+    
+    @abstractmethod
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
+        ...
+    
+    def process_request_prompt(self, prompt: 'dict[str, Any]', args: Sequence[Any], kwargs: 'dict[str, Any]') -> None:
+        ...
+    
+    def process_initial_stream_response(self, response: Any) -> None:
+        self.add_instrumented_response_headers(response)
+
+    def remove_inline_data(self, prompt: 'dict[str, Any]') -> bool:# noqa: ARG002
+        return False
+
+    @property
+    def is_aws_client(self) -> bool:
+        return self._is_aws_client if self._is_aws_client is not None else False
+
+    @property
+    def is_google_vertex_or_genai_client(self) -> bool:
+        return self._is_google_vertex_or_genai_client if self._is_google_vertex_or_genai_client is not None else False
+
+    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool: # noqa: ARG002
+        self.exception_to_semantic_failure(exception)
+        return True
+    
+    @property
+    def supports_extra_headers(self) -> bool:
+        return not self.is_aws_client and not self.is_google_vertex_or_genai_client
+    
+    @property
+    def streaming_type(self) -> '_StreamingType':
+        return self._streaming_type
+
+    def add_internal_request_property(self, key: str, value: str) -> None:
+        self._internal_request_properties[key] = value
+
+    def exception_to_semantic_failure(self, e: Exception) -> None:
+        exception_str = f"{type(e).__name__}"
+    
+        fields: list[str] = []
+    
+        for attr in dir(e):
+            if not attr.startswith("__"):
+                try:
+                    value = getattr(e, attr)
+                    if value and not inspect.ismethod(value) and not inspect.isfunction(value) and not callable(value):
+                        fields.append(f"{attr}={value}")
+                except Exception as _ex:
+                    pass
+ 
+        self.add_internal_request_property(PayiPropertyNames.failure, exception_str)
+        if fields:
+            failure_description = ",".join(fields)
+            self.add_internal_request_property(PayiPropertyNames.failure_description, failure_description)
+
+        if "http_status_code" not in self._ingest:
+            # use a non existent http status code so when presented to the user, the origin is clear
+            self._ingest["http_status_code"] = 299
+
+    def add_streaming_function_call(self, index: int, name: Optional[str], arguments: Optional[str]) -> None:
+        if not self._function_call_builder:
+            self._function_call_builder = {}
+
+        if not index in self._function_call_builder:
+            self._function_call_builder[index] = ProviderResponseFunctionCall(name=name or "", arguments=arguments or "")
+        else:
+            function = self._function_call_builder[index]
+            if name:
+                function["name"] = function["name"] + name
+            if arguments:
+                function["arguments"] = (function.get("arguments", "") or "") + arguments
+
+    def add_synchronous_function_call(self, name: str, arguments: Optional[str]) -> None:
+        if not self._function_calls:
+            self._function_calls = []
+            self._ingest["provider_response_function_calls"] = self._function_calls
+        self._function_calls.append(ProviderResponseFunctionCall(name=name, arguments=arguments))
+    
+    def add_instrumented_response_headers(self, response: Any) -> None:
+        response_headers  = getattr(response, _ProviderRequest._instrumented_response_headers_attr, {})
+        if response_headers:
+            self.add_response_headers(response_headers)
+
+    def add_response_headers(self, response_headers: 'dict[str, Any]') -> None:
+        self._ingest["provider_response_headers"] = [
+            PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) 
+            for k, v in response_headers.items() 
+            if (k_lower := k.lower()) not in _ProviderRequest.excluded_headers and not k_lower.startswith("content-")
+        ]
+
+    def merge_internal_request_properties(self) -> None:
+        if not self._internal_request_properties:
+            return
+        
+        properties = self._ingest.get("properties") or {}
+        self._ingest["properties"] = properties
+        for key, value in self._internal_request_properties.items():
+            if key not in properties:
+                properties[key] = value
+                
+    def update_for_vision(self, input: int) -> int:
+        if self._estimated_prompt_tokens:
+            vision = input - self._estimated_prompt_tokens
+            if (vision > 0):
+                key = "vision_large_context" if self._is_large_context else "vision"
+                self._ingest["units"][key] = IngestUnits(input=vision, output=0)
+                input = self._estimated_prompt_tokens
+        
+        return input
+
+    @staticmethod
+    def assign_xproxy_result(o: Any, xproxy_result: XproxyResult |  XproxyError| None) -> None:
+        if xproxy_result:
+            _set_attr_safe(o, _ProviderRequest._xproxy_result_attr, xproxy_result)
+
+    @staticmethod
+    def process_response_wrapper(wrapped: Any, _instance: Any, args: Any, kwargs: Any) -> Any:
+        httpResponse = kwargs.get("response", None)
+
+        r =  wrapped(*args, **kwargs)
+
+        if httpResponse:
+            headers = getattr(httpResponse, "headers", None)
+            _set_attr_safe(r, _ProviderRequest._instrumented_response_headers_attr, dict(headers) if headers else {})
+
+        return r
+
+    @staticmethod
+    async def aprocess_response_wrapper(wrapped: Any, _instance: Any, args: Any, kwargs: Any) -> Any:
+        httpResponse = kwargs.get("response", None)
+
+        r = await wrapped(*args, **kwargs)
+
+        if httpResponse:
+            headers = getattr(httpResponse, "headers", None)
+            _set_attr_safe(r, _ProviderRequest._instrumented_response_headers_attr, dict(headers) if headers else {})
+
+        return r
+