payi 0.1.0a63__py3-none-any.whl → 0.1.0a65__py3-none-any.whl

payi/_utils/_transform.py

@@ -142,6 +142,10 @@ def _maybe_transform_key(key: str, type_: type) -> str:
     return key
 
 
+def _no_transform_needed(annotation: type) -> bool:
+    return annotation == float or annotation == int
+
+
 def _transform_recursive(
     data: object,
     *,
@@ -184,6 +188,15 @@ def _transform_recursive(
             return cast(object, data)
 
         inner_type = extract_type_arg(stripped_type, 0)
+        if _no_transform_needed(inner_type):
+            # for some types there is no need to transform anything, so we can get a small
+            # perf boost from skipping that work.
+            #
+            # but we still need to convert to a list to ensure the data is json-serializable
+            if is_list(data):
+                return data
+            return list(data)
+
         return [_transform_recursive(d, annotation=annotation, inner_type=inner_type) for d in data]
 
     if is_union_type(stripped_type):
@@ -332,6 +345,15 @@ async def _async_transform_recursive(
             return cast(object, data)
 
         inner_type = extract_type_arg(stripped_type, 0)
+        if _no_transform_needed(inner_type):
+            # for some types there is no need to transform anything, so we can get a small
+            # perf boost from skipping that work.
+            #
+            # but we still need to convert to a list to ensure the data is json-serializable
+            if is_list(data):
+                return data
+            return list(data)
+
         return [await _async_transform_recursive(d, annotation=annotation, inner_type=inner_type) for d in data]
 
     if is_union_type(stripped_type):

payi/_version.py

@@ -1,4 +1,4 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 __title__ = "payi"
-__version__ = "0.1.0-alpha.63"  # x-release-please-version
+__version__ = "0.1.0-alpha.65"  # x-release-please-version

payi/lib/AnthropicInstrumentor.py

@@ -1,13 +1,13 @@
 import logging
 from typing import Any, Union
+from typing_extensions import override
 
 import tiktoken
 from wrapt import wrap_function_wrapper  # type: ignore
 
-from payi.types import IngestUnitsParams
 from payi.types.ingest_units_params import Units
 
-from .instrument import _IsStreaming, _PayiInstrumentor
+from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
 
 
 class AnthropicIntrumentor:
@@ -55,9 +55,7 @@ def chat_wrapper(
 ) -> Any:
     return instrumentor.chat_wrapper(
         "system.anthropic",
-        process_chunk,
-        process_request,
-        process_synchronous_response,
+        _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -75,9 +73,7 @@ async def achat_wrapper(
 ) -> Any:
     return await instrumentor.achat_wrapper(
         "system.anthropic",
-        process_chunk,
-        process_request,
-        process_synchronous_response,
+        _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -85,17 +81,39 @@ async def achat_wrapper(
         kwargs,
     )
 
+class _AnthropicProviderRequest(_ProviderRequest):
+    @override
+    def process_chunk(self, chunk: Any) -> bool:
+        if chunk.type == "message_start":
+            self._ingest["provider_response_id"] = chunk.message.id
 
-def process_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
-    if chunk.type == "message_start":
-        ingest["provider_response_id"] = chunk.message.id
+            usage = chunk.message.usage
+            units = self._ingest["units"]
 
-        usage = chunk.message.usage
-        units = ingest["units"]
+            input = _PayiInstrumentor.update_for_vision(usage.input_tokens, units, self._estimated_prompt_tokens)
 
-        input = _PayiInstrumentor.update_for_vision(usage.input_tokens, units)
+            units["text"] = Units(input=input, output=0)
 
-        units["text"] = Units(input=input, output=0)
+            if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
+                text_cache_write = usage.cache_creation_input_tokens
+                units["text_cache_write"] = Units(input=text_cache_write, output=0)
+
+            if hasattr(usage, "cache_read_input_tokens") and usage.cache_read_input_tokens > 0:
+                text_cache_read = usage.cache_read_input_tokens
+                units["text_cache_read"] = Units(input=text_cache_read, output=0)
+
+        elif chunk.type == "message_delta":
+            usage = chunk.usage
+            self._ingest["units"]["text"]["output"] = usage.output_tokens
+        
+        return True
+
+    @override
+    def process_synchronous_response(self, response: Any, log_prompt_and_response: bool, kwargs: Any) -> Any:
+        usage = response.usage
+        input = usage.input_tokens
+        output = usage.output_tokens
+        units: dict[str, Units] = self._ingest["units"]
 
         if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
             text_cache_write = usage.cache_creation_input_tokens
@@ -105,35 +123,37 @@ def process_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
             text_cache_read = usage.cache_read_input_tokens
             units["text_cache_read"] = Units(input=text_cache_read, output=0)
 
-    elif chunk.type == "message_delta":
-        usage = chunk.usage
-        ingest["units"]["text"]["output"] = usage.output_tokens
-
-
-def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_prompt_and_response: bool, *args: Any, **kwargs: 'dict[str, Any]') -> Any: # noqa: ARG001
-    usage = response.usage
-    input = usage.input_tokens
-    output = usage.output_tokens
-    units: dict[str, Units] = ingest["units"]
-
-    if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
-        text_cache_write = usage.cache_creation_input_tokens
-        units["text_cache_write"] = Units(input=text_cache_write, output=0)
+        input = _PayiInstrumentor.update_for_vision(input, units, self._estimated_prompt_tokens)
 
-    if hasattr(usage, "cache_read_input_tokens") and usage.cache_read_input_tokens > 0:
-        text_cache_read = usage.cache_read_input_tokens
-        units["text_cache_read"] = Units(input=text_cache_read, output=0)
+        units["text"] = Units(input=input, output=output)
 
-    input = _PayiInstrumentor.update_for_vision(input, units)
+        if log_prompt_and_response:
+            self._ingest["provider_response_json"] = response.to_json()
+        
+        self._ingest["provider_response_id"] = response.id
+        
+        return None
 
-    units["text"] = Units(input=input, output=output)
+    @override
+    def process_request(self, kwargs: Any) -> None:
+        messages = kwargs.get("messages")
+        if not messages or len(messages) == 0:
+            return
+        
+        estimated_token_count = 0 
+        has_image = False
 
-    if log_prompt_and_response:
-        ingest["provider_response_json"] = response.to_json()
-    
-    ingest["provider_response_id"] = response.id
-    
-    return None
+        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):
@@ -146,23 +166,3 @@ def has_image_and_get_texts(encoding: tiktoken.Encoding, content: Union[str, 'li
         token_count = sum(len(encoding.encode(item.get("text", ""))) for item in content if item.get("type") == "text")
         return has_image, token_count
 
-def process_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: # noqa: ARG001
-    messages = kwargs.get("messages")
-    if not messages or len(messages) == 0:
-        return
-    
-    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 not has_image or estimated_token_count == 0:
-        return
-
-    ingest["units"][_PayiInstrumentor.estimated_prompt_tokens] = Units(input=estimated_token_count, output=0)

payi/lib/BedrockInstrumentor.py

@@ -2,13 +2,14 @@ import json
 import logging
 from typing import Any
 from functools import wraps
+from typing_extensions import override
 
 from wrapt import ObjectProxy, wrap_function_wrapper  # type: ignore
 
 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, _PayiInstrumentor
+from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
 
 
 class BedrockInstrumentor:
@@ -103,9 +104,7 @@ def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
                 "system.aws.bedrock",
-                None,
-                process_invoke_request,  
-                process_synchronous_invoke_response,  
+                _BedrockInvokeSynchronousProviderRequest(instrumentor),
                 _IsStreaming.false,
                 wrapped,
                 None,
@@ -119,14 +118,12 @@ 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:
-        modelId:str = kwargs.get("modelId", "") # type: ignore
+        model_id: str = kwargs.get("modelId", "") # type: ignore
 
-        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+        if model_id.startswith("meta.llama3") or model_id.startswith("anthropic."):
             return instrumentor.chat_wrapper(
                 "system.aws.bedrock",
-                process_invoke_streaming_anthropic_chunk if modelId.startswith("anthropic.") else process_invoke_streaming_llama_chunk,
-                process_invoke_request, 
-                None,  
+                _BedrockInvokeStreamingProviderRequest(instrumentor, model_id),
                 _IsStreaming.true,
                 wrapped,
                 None,
@@ -145,9 +142,7 @@ def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
                 "system.aws.bedrock",
-                None,
-                process_converse_request,  
-                process_synchronous_converse_response,  
+                _BedrockConverseSynchronousProviderRequest(instrumentor),
                 _IsStreaming.false,
                 wrapped,
                 None,
@@ -161,14 +156,12 @@ 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:
-        modelId:str = kwargs.get("modelId", "") # type: ignore
+        model_id: str = kwargs.get("modelId", "") # type: ignore
 
-        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+        if model_id.startswith("meta.llama3") or model_id.startswith("anthropic."):
             return instrumentor.chat_wrapper(
                 "system.aws.bedrock",
-                process_converse_streaming_chunk,
-                process_converse_request, 
-                None,  
+                _BedrockConverseStreamingProviderRequest(instrumentor),
                 _IsStreaming.true,
                 wrapped,
                 None,
@@ -179,104 +172,121 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 
     return invoke_wrapper
 
-def process_invoke_streaming_anthropic_chunk(chunk: str, ingest: IngestUnitsParams) -> None:
-    chunk_dict =  json.loads(chunk)
-    type = chunk_dict.get("type", "")
+class _BedrockInvokeStreamingProviderRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor, model_id: str):
+        super().__init__(instrumentor)
+        self._is_anthropic: bool = model_id.startswith("anthropic.")
 
-    if type == "message_start":
-        usage = chunk_dict['message']['usage']
-        units = ingest["units"]
+    @override
+    def process_chunk(self, chunk: Any) -> bool:
+        if self._is_anthropic:
+            return self.process_invoke_streaming_anthropic_chunk(chunk)
+        else:
+            return self.process_invoke_streaming_llama_chunk(chunk)
 
-        input = _PayiInstrumentor.update_for_vision(usage['input_tokens'], units)
+    def process_invoke_streaming_anthropic_chunk(self, chunk: str) -> bool:
+        chunk_dict =  json.loads(chunk)
+        type = chunk_dict.get("type", "")
 
-        units["text"] = Units(input=input, output=0)
+        if type == "message_start":
+            usage = chunk_dict['message']['usage']
+            units = self._ingest["units"]
 
-        text_cache_write: int = usage.get("cache_creation_input_tokens", 0)
-        if text_cache_write > 0:
-            units["text_cache_write"] = Units(input=text_cache_write, output=0)
+            input = _PayiInstrumentor.update_for_vision(usage['input_tokens'], units, self._estimated_prompt_tokens)
 
-        text_cache_read: int = usage.get("cache_read_input_tokens", 0)
-        if text_cache_read > 0:
-            units["text_cache_read"] = Units(input=text_cache_read, output=0)
+            units["text"] = Units(input=input, output=0)
 
-    elif type == "message_delta":
-        usage = chunk_dict['usage']
-        ingest["units"]["text"]["output"] = usage['output_tokens']
+            text_cache_write: int = usage.get("cache_creation_input_tokens", 0)
+            if text_cache_write > 0:
+                units["text_cache_write"] = Units(input=text_cache_write, output=0)
 
-def process_invoke_streaming_llama_chunk(chunk: str, ingest: IngestUnitsParams) -> None:
-    chunk_dict =  json.loads(chunk)
-    metrics = chunk_dict.get("amazon-bedrock-invocationMetrics", {})
-    if metrics:
-        input = metrics.get("inputTokenCount", 0)
-        output = metrics.get("outputTokenCount", 0)
-        ingest["units"]["text"] = Units(input=input, output=output)
-    
-def process_synchronous_invoke_response(
+            text_cache_read: int = usage.get("cache_read_input_tokens", 0)
+            if text_cache_read > 0:
+                units["text_cache_read"] = Units(input=text_cache_read, output=0)
+
+        elif type == "message_delta":
+            usage = chunk_dict['usage']
+            self._ingest["units"]["text"]["output"] = usage['output_tokens']
+
+        return True    
+
+    def process_invoke_streaming_llama_chunk(self, chunk: str) -> bool:
+        chunk_dict =  json.loads(chunk)
+        metrics = chunk_dict.get("amazon-bedrock-invocationMetrics", {})
+        if metrics:
+            input = metrics.get("inputTokenCount", 0)
+            output = metrics.get("outputTokenCount", 0)
+            self._ingest["units"]["text"] = Units(input=input, output=output)
+
+        return True
+
+class _BedrockInvokeSynchronousProviderRequest(_ProviderRequest):
+    @override
+    def process_synchronous_response(
+        self,
         response: Any,
-        ingest: IngestUnitsParams,
         log_prompt_and_response: bool,
-        instrumentor: _PayiInstrumentor,
-        **kargs: Any) -> Any: #  noqa: ARG001
+        kwargs: Any) -> Any:
 
-    metadata = response.get("ResponseMetadata", {})
+        metadata = response.get("ResponseMetadata", {})
 
-    request_id = metadata.get("RequestId", "")
-    if request_id:
-        ingest["provider_response_id"] = request_id
+        request_id = metadata.get("RequestId", "")
+        if request_id:
+            self._ingest["provider_response_id"] = request_id
 
-    response_headers = metadata.get("HTTPHeaders", {}).copy()
-    if response_headers:
-        ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+        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()]
 
-    response["body"] = InvokeResponseWrapper(
-        response=response["body"],
-        instrumentor=instrumentor,
-        ingest=ingest,
-        log_prompt_and_response=log_prompt_and_response)
+        response["body"] = InvokeResponseWrapper(
+            response=response["body"],
+            instrumentor=self._instrumentor,
+            ingest=self._ingest,
+            log_prompt_and_response=log_prompt_and_response)
 
-    return response
+        return response
 
-def process_invoke_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: #  noqa: ARG001
-    return
-
-def process_converse_streaming_chunk(chunk: 'dict[str, Any]', ingest: IngestUnitsParams) -> None:
-    metadata = chunk.get("metadata", {})
+class _BedrockConverseSynchronousProviderRequest(_ProviderRequest):
+    @override
+    def process_synchronous_response(
+        self,
+        response: 'dict[str, Any]',
+        log_prompt_and_response: bool,
+        kwargs: Any) -> Any:
 
-    if metadata:
-        usage = metadata['usage']
+        usage = response["usage"]
         input = usage["inputTokens"]
         output = usage["outputTokens"]
-        ingest["units"]["text"] = Units(input=input, output=output)
+        
+        units: dict[str, Units] = self._ingest["units"]
+        units["text"] = Units(input=input, output=output)
 
-def process_synchronous_converse_response(
-        response: 'dict[str, Any]',
-        ingest: IngestUnitsParams,
-        log_prompt_and_response: bool,
-        **kargs: Any) -> Any: #  noqa: ARG001
+        metadata = response.get("ResponseMetadata", {})
 
-    usage = response["usage"]
-    input = usage["inputTokens"]
-    output = usage["outputTokens"]
-    
-    units: dict[str, Units] = ingest["units"]
-    units["text"] = Units(input=input, output=output)
+        request_id = metadata.get("RequestId", "")
+        if request_id:
+            self._ingest["provider_response_id"] = request_id
 
-    metadata = response.get("ResponseMetadata", {})
+        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()]
 
-    request_id = metadata.get("RequestId", "")
-    if request_id:
-        ingest["provider_response_id"] = request_id
+        if log_prompt_and_response:
+            response_without_metadata = response.copy()
+            response_without_metadata.pop("ResponseMetadata", None)
+            self._ingest["provider_response_json"] = json.dumps(response_without_metadata)
 
-    response_headers = metadata.get("HTTPHeaders", {})
-    if response_headers:
-        ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+        return None    
 
-    if log_prompt_and_response:
-        response_without_metadata = response.copy()
-        response_without_metadata.pop("ResponseMetadata", None)
-        ingest["provider_response_json"] = json.dumps(response_without_metadata)
+class _BedrockConverseStreamingProviderRequest(_ProviderRequest):
+    @override
+    def process_chunk(self, chunk: 'dict[str, Any]') -> bool:
+        metadata = chunk.get("metadata", {})
 
-    return None    
+        if metadata:
+            usage = metadata['usage']
+            input = usage["inputTokens"]
+            output = usage["outputTokens"]
+            self._ingest["units"]["text"] = Units(input=input, output=output)
 
-def process_converse_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: #  noqa: ARG001
-    return
+        return True