payi 0.1.0a40__py3-none-any.whl → 0.1.0a41__py3-none-any.whl

payi/_constants.py

@@ -6,7 +6,7 @@ RAW_RESPONSE_HEADER = "X-Stainless-Raw-Response"
 OVERRIDE_CAST_TO_HEADER = "____stainless_override_cast_to"
 
 # default timeout is 1 minute
-DEFAULT_TIMEOUT = httpx.Timeout(timeout=60.0, connect=5.0)
+DEFAULT_TIMEOUT = httpx.Timeout(timeout=60, connect=5.0)
 DEFAULT_MAX_RETRIES = 2
 DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20)
 

payi/_models.py

@@ -172,7 +172,7 @@ class BaseModel(pydantic.BaseModel):
     @override
     def __str__(self) -> str:
         # mypy complains about an invalid self arg
-        return f'{self.__repr_name__()}({self.__repr_str__(", ")})'  # type: ignore[misc]
+        return f"{self.__repr_name__()}({self.__repr_str__(', ')})"  # type: ignore[misc]
 
     # Override the 'construct' method in a way that supports recursive parsing without validation.
     # Based on https://github.com/samuelcolvin/pydantic/issues/1168#issuecomment-817742836.

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.40"  # x-release-please-version
+__version__ = "0.1.0-alpha.41"  # x-release-please-version

payi/lib/AnthropicInstrumentor.py

@@ -7,7 +7,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.types import IngestUnitsParams
 from payi.types.ingest_units_params import Units
 
-from .instrument import PayiInstrumentor
+from .instrument import IsStreaming, PayiInstrumentor
 
 
 class AnthropicIntrumentor:
@@ -48,14 +48,15 @@ def chat_wrapper(
     kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        "system.anthropic",
-        process_chunk,
-        process_request,
-        process_synchronous_response,
-        wrapped,
-        instance,
-        args,
-        kwargs,
+        category="system.anthropic",
+        process_chunk=process_chunk,
+        process_request=process_request,
+        process_synchronous_response=process_synchronous_response,
+        is_streaming=IsStreaming.kwargs,
+        wrapped=wrapped,
+        instance=instance,
+        args=args,
+        kwargs=kwargs,
     )
 
 
@@ -81,10 +82,10 @@ def process_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
         ingest["units"]["text"]["output"] = usage.output_tokens
 
 
-def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+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
-    ouptut = usage.output_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:
@@ -97,11 +98,13 @@ def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_p
 
     input = PayiInstrumentor.update_for_vision(input, units)
 
-    units["text"] = Units(input=input, output=ouptut)
+    units["text"] = Units(input=input, output=output)
 
     if log_prompt_and_response:
         ingest["provider_response_json"] = response.to_json()
 
+    return None
+
 def has_image_and_get_texts(encoding: tiktoken.Encoding, content: Union[str, 'list[Any]']) -> 'tuple[bool, int]':
     if isinstance(content, str):
         return False, 0

payi/lib/BedrockInstrumentor.py

@@ -0,0 +1,288 @@
+import json
+import logging
+from typing import Any
+from functools import wraps
+
+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
+
+
+class BedrockInstrumentor:
+    @staticmethod
+    def instrument(instrumentor: PayiInstrumentor) -> None:
+        try:
+            import boto3  # type: ignore #  noqa: F401  I001
+
+            # wrap_function_wrapper(
+            #     "anthropic.resources.completions",
+            #     "Completions.create",
+            #     chat_wrapper(instrumentor),
+            # )
+
+            wrap_function_wrapper(
+                "botocore.client",
+                "ClientCreator.create_client",
+                create_client_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "botocore.session",
+                "Session.create_client",
+                create_client_wrapper(instrumentor),
+            )
+
+        except Exception as e:
+            logging.debug(f"Error instrumenting bedrock: {e}")
+            return
+
+@PayiInstrumentor.payi_wrapper
+def create_client_wrapper(instrumentor: PayiInstrumentor, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: #  noqa: ARG001
+    if kwargs.get("service_name") != "bedrock-runtime":
+        return wrapped(*args, **kwargs)
+
+    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)
+
+        return client
+    except Exception as e:
+        logging.debug(f"Error instrumenting bedrock client: {e}")
+    
+    return wrapped(*args, **kwargs)
+    
+class InvokeResponseWrapper(ObjectProxy): # type: ignore
+    def __init__(
+        self,
+        response: Any,
+        instrumentor: PayiInstrumentor,
+        ingest: IngestUnitsParams,
+        log_prompt_and_response: bool
+        ) -> None:
+
+        super().__init__(response) # type: ignore
+        self._response = response
+        self._instrumentor = instrumentor
+        self._ingest = ingest
+        self._log_prompt_and_response = log_prompt_and_response
+
+    def read(self, amt: Any =None): # type: ignore
+        # data is array of bytes
+        data: Any = self.__wrapped__.read(amt) # type: ignore
+        response = json.loads(data)
+
+        resource = self._ingest["resource"]
+        if not resource:
+            return
+        
+        input: int = 0
+        output: int = 0
+        units: dict[str, Units] = self._ingest["units"]
+
+        if resource.startswith("meta.llama3"):
+            input = response['prompt_token_count']
+            output = response['generation_token_count']
+        elif resource.startswith("anthropic."):
+            usage = response['usage']
+            input = usage['input_tokens']
+            output = usage['output_tokens']
+        units["text"] = Units(input=input, output=output)
+
+        if self._log_prompt_and_response:
+            self._ingest["provider_response_json"] = data.decode('utf-8')
+            
+        self._instrumentor._ingest_units(self._ingest)
+
+        return data
+
+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."):
+            return instrumentor.chat_wrapper(
+                category="system.aws.bedrock",
+                process_chunk=None,
+                process_request=process_invoke_request,  
+                process_synchronous_response=process_synchronous_invoke_response,  
+                is_streaming=IsStreaming.false,
+                wrapped=wrapped,
+                instance=None,
+                args=args,
+                kwargs=kwargs,
+        )
+        return wrapped(*args, **kwargs)
+    
+    return invoke_wrapper
+
+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
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                category="system.aws.bedrock",
+                process_chunk=process_invoke_streaming_anthropic_chunk if modelId.startswith("anthropic.") else process_invoke_streaming_llama_chunk,
+                process_request=process_invoke_request, 
+                process_synchronous_response=None,  
+                is_streaming=IsStreaming.true,
+                wrapped=wrapped,
+                instance=None,
+                args=args,
+                kwargs=kwargs,
+            )
+        return wrapped(*args, **kwargs)
+
+    return invoke_wrapper
+
+def wrap_converse(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."):
+            return instrumentor.chat_wrapper(
+                category="system.aws.bedrock",
+                process_chunk=None,
+                process_request=process_converse_request,  
+                process_synchronous_response=process_synchronous_converse_response,  
+                is_streaming=IsStreaming.false,
+                wrapped=wrapped,
+                instance=None,
+                args=args,
+                kwargs=kwargs,
+        )
+        return wrapped(*args, **kwargs)
+    
+    return invoke_wrapper
+
+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
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                category="system.aws.bedrock",
+                process_chunk=process_converse_streaming_chunk,
+                process_request=process_converse_request, 
+                process_synchronous_response=None,  
+                is_streaming=IsStreaming.true,
+                wrapped=wrapped,
+                instance=None,
+                args=args,
+                kwargs=kwargs,
+            )
+        return wrapped(*args, **kwargs)
+
+    return invoke_wrapper
+
+def process_invoke_streaming_anthropic_chunk(chunk: str, ingest: IngestUnitsParams) -> None:
+    chunk_dict =  json.loads(chunk)
+    type = chunk_dict.get("type", "")
+
+    if type == "message_start":
+        usage = chunk_dict['message']['usage']
+        units = ingest["units"]
+
+        input = PayiInstrumentor.update_for_vision(usage['input_tokens'], units)
+
+        units["text"] = Units(input=input, output=0)
+
+        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)
+
+        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']
+        ingest["units"]["text"]["output"] = usage['output_tokens']
+
+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(
+        response: Any,
+        ingest: IngestUnitsParams,
+        log_prompt_and_response: bool,
+        instrumentor: PayiInstrumentor,
+        **kargs: Any) -> Any: #  noqa: ARG001
+
+    metadata = response.get("ResponseMetadata", {})
+
+    # request_id = metadata.get("RequestId", "")
+    # if request_id:
+    #     ingest["provider_request_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["body"] = InvokeResponseWrapper(
+        response=response["body"],
+        instrumentor=instrumentor,
+        ingest=ingest,
+        log_prompt_and_response=log_prompt_and_response)
+
+    return response
+
+def process_invoke_request(ingest: IngestUnitsParams, kwargs: Any) -> None: #  noqa: ARG001
+    return
+
+def process_converse_streaming_chunk(chunk: 'dict[str, Any]', ingest: IngestUnitsParams) -> None:
+    metadata = chunk.get("metadata", {})
+
+    if metadata:
+        usage = metadata['usage']
+        input = usage["inputTokens"]
+        output = usage["outputTokens"]
+        ingest["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
+
+    usage = response["usage"]
+    input = usage["inputTokens"]
+    output = usage["outputTokens"]
+    
+    units: dict[str, Units] = ingest["units"]
+    units["text"] = Units(input=input, output=output)
+
+    metadata = response.get("ResponseMetadata", {})
+
+    # request_id = metadata.get("RequestId", "")
+    # if request_id:
+    #     ingest["provider_request_id"] = request_id
+
+    response_headers = metadata.get("HTTPHeaders", {})
+    if response_headers:
+        ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+
+    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)
+
+    return None    
+
+def process_converse_request(ingest: IngestUnitsParams, kwargs: Any) -> None: #  noqa: ARG001
+    return

payi/lib/Instruments.py

@@ -5,3 +5,4 @@ class Instruments(Enum):
     ALL = "all"
     OPENAI = "openai"
     ANTHROPIC = "anthropic"
+    AWS_BEDROCK = "aws.bedrock"

payi/lib/OpenAIInstrumentor.py

@@ -9,7 +9,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.types import IngestUnitsParams
 from payi.types.ingest_units_params import Units
 
-from .instrument import PayiInstrumentor
+from .instrument import IsStreaming, PayiInstrumentor
 
 
 class OpenAiInstrumentor:
@@ -37,18 +37,19 @@ def chat_wrapper(
     kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        "system.openai",
-        process_chat_chunk,
-        process_request,
-        process_chat_synchronous_response,
-        wrapped,
-        instance,
-        args,
-        kwargs,
+        category="system.openai",
+        process_chunk=process_chat_chunk,
+        process_request=process_request,
+        process_synchronous_response=process_chat_synchronous_response,
+        is_streaming=IsStreaming.kwargs,
+        wrapped=wrapped,
+        instance=instance,
+        args=args,
+        kwargs=kwargs,
     )
 
 
-def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool, **kwargs: Any) -> Any: #  noqa: ARG001
     response_dict = model_to_dict(response)
 
     add_usage_units(response_dict["usage"], ingest["units"])
@@ -56,6 +57,7 @@ def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams,
     if log_prompt_and_response:
         ingest["provider_response_json"] = [json.dumps(response_dict)]
 
+    return None
 
 def process_chat_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
     model = model_to_dict(chunk)

payi/lib/instrument.py

@@ -4,6 +4,7 @@ import asyncio
 import inspect
 import logging
 import traceback
+from enum import Enum
 from typing import Any, Set, Union, Callable, Optional
 
 from wrapt import ObjectProxy  # type: ignore
@@ -11,11 +12,17 @@ from wrapt import ObjectProxy  # type: ignore
 from payi import Payi, AsyncPayi
 from payi.types import IngestUnitsParams
 from payi.types.ingest_units_params import Units
+from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
 from .Stopwatch import Stopwatch
 from .Instruments import Instruments
 
 
+class IsStreaming(Enum):
+    false = 0
+    true = 1 
+    kwargs = 2
+
 class PayiInstrumentor:
     estimated_prompt_tokens: str = "estimated_prompt_tokens"
 
@@ -44,12 +51,15 @@ class PayiInstrumentor:
     def _instrument_all(self) -> None:
         self._instrument_openai()
         self._instrument_anthropic()
+        self._instrument_aws_bedrock()
 
     def _instrument_specific(self, instruments: Set[Instruments]) -> None:
         if Instruments.OPENAI in instruments:
             self._instrument_openai()
         if Instruments.ANTHROPIC in instruments:
             self._instrument_anthropic()
+        if Instruments.AWS_BEDROCK in instruments:
+            self._instrument_aws_bedrock()
 
     def _instrument_openai(self) -> None:
         from .OpenAIInstrumentor import OpenAiInstrumentor
@@ -69,6 +79,15 @@ class PayiInstrumentor:
         except Exception as e:
             logging.error(f"Error instrumenting Anthropic: {e}")
 
+    def _instrument_aws_bedrock(self) -> None:
+        from .BedrockInstrumentor import BedrockInstrumentor
+
+        try:
+            BedrockInstrumentor.instrument(self)
+
+        except Exception as e:
+            logging.error(f"Error instrumenting AWS bedrock: {e}")
+
     def _ingest_units(self, ingest_units: IngestUnitsParams) -> None:
         # return early if there are no units to ingest and on a successul ingest request
         if int(ingest_units.get("http_status_code") or 0) < 400:
@@ -206,18 +225,25 @@ class PayiInstrumentor:
     def chat_wrapper(
         self,
         category: str,
-        process_chunk: Callable[[Any, IngestUnitsParams], None],
+        process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]],
         process_request: Optional[Callable[[IngestUnitsParams, Any], None]],
-        process_synchronous_response: Optional[Callable[[Any, IngestUnitsParams, bool], None]],
+        process_synchronous_response: Any,
+        is_streaming: IsStreaming,
         wrapped: Any,
         instance: Any,
         args: Any,
-        kwargs: 'dict[str, Any]',
+        kwargs: Any,
     ) -> Any:
         context = self.get_context()
 
+        is_bedrock:bool = category == "system.aws.bedrock"
+
         if not context:
-            # should not happen
+            if is_bedrock:
+                # boto3 doesn't allow extra_headers
+                kwargs.pop("extra_headers", None)
+    
+            # wrapped function invoked outside of decorator scope
             return wrapped(*args, **kwargs)
 
         # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
@@ -230,11 +256,14 @@ class PayiInstrumentor:
 
             return wrapped(*args, **kwargs)
 
-        ingest: IngestUnitsParams = {"category": category, "resource": kwargs.get("model"), "units": {}} # type: ignore
+        ingest: IngestUnitsParams = {"category": category, "units": {}} # type: ignore
+        if is_bedrock:
+            # boto3 doesn't allow extra_headers
+            kwargs.pop("extra_headers", None)
+            ingest["resource"] = kwargs.get("modelId", "")
+        else:
+            ingest["resource"] = kwargs.get("model", "")
 
-        # blocked_limit = next((limit for limit in (context.get('limit_ids') or []) if limit in self._blocked_limits), None)
-        # if blocked_limit:
-        #      raise Exception(f"Limit {blocked_limit} is blocked")
         current_frame = inspect.currentframe()
         # f_back excludes the current frame, strip() cleans up whitespace and newlines
         stack = [frame.strip() for frame in traceback.format_stack(current_frame.f_back)]  # type: ignore
@@ -245,7 +274,14 @@ class PayiInstrumentor:
             process_request(ingest, kwargs)
 
         sw = Stopwatch()
-        stream = kwargs.get("stream", False)
+        stream: bool = False
+        
+        if is_streaming == IsStreaming.kwargs:
+            stream = kwargs.get("stream", False)
+        elif is_streaming == IsStreaming.true:
+            stream = True
+        else:
+            stream = False
 
         try:
             limit_ids = extra_headers.pop("xProxy-Limit-IDs", None)
@@ -266,7 +302,7 @@ class PayiInstrumentor:
                 ingest["user_id"] = user_id
 
             if len(extra_headers) > 0:
-                ingest["provider_request_headers"] = {k: [v] for k, v in extra_headers.items()}  # type: ignore
+                ingest["provider_request_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in extra_headers.items()]
 
             provider_prompt = {}
             for k, v in kwargs.items():
@@ -292,7 +328,7 @@ class PayiInstrumentor:
             raise e
 
         if stream:
-            return ChatStreamWrapper(
+            stream_result = ChatStreamWrapper(
                 response=response,
                 instance=instance,
                 instrumentor=self,
@@ -300,15 +336,31 @@ class PayiInstrumentor:
                 ingest=ingest,
                 stopwatch=sw,
                 process_chunk=process_chunk,
+                is_bedrock=is_bedrock,
             )
 
+            if is_bedrock:
+                if "body" in response:
+                    response["body"] = stream_result
+                else:
+                    response["stream"] = stream_result
+                return response
+            
+            return stream_result
+
         sw.stop()
         duration = sw.elapsed_ms_int()
         ingest["end_to_end_latency_ms"] = duration
         ingest["http_status_code"] = 200
 
         if process_synchronous_response:
-            process_synchronous_response(response, ingest, self._log_prompt_and_response)
+            return_result: Any = process_synchronous_response(
+                response=response,
+                ingest=ingest,
+                log_prompt_and_response=self._log_prompt_and_response,
+                instrumentor=self)
+            if return_result:
+                return return_result
 
         self._ingest_units(ingest)
 
@@ -387,7 +439,6 @@ class PayiInstrumentor:
 
         return _payi_wrapper
 
-
 class ChatStreamWrapper(ObjectProxy):  # type: ignore
     def __init__(
         self,
@@ -398,7 +449,19 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         stopwatch: Stopwatch,
         process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]] = None,
         log_prompt_and_response: bool = True,
+        is_bedrock: bool = False,
     ) -> None:
+
+        bedrock_from_stream: bool = False
+        if is_bedrock:
+            stream = response.get("stream", None)
+            if stream:
+                response = stream
+                bedrock_from_stream = True
+            else:
+                response = response.get("body")
+                bedrock_from_stream = False
+
         super().__init__(response)  # type: ignore
 
         self._response = response
@@ -413,6 +476,8 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         self._process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]] = process_chunk
 
         self._first_token: bool = True
+        self._is_bedrock: bool = is_bedrock
+        self._bedrock_from_stream: bool = bedrock_from_stream
 
     def __enter__(self) -> Any:
         return self
@@ -426,9 +491,26 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.__wrapped__.__aexit__(exc_type, exc_val, exc_tb)  # type: ignore
 
-    def __iter__(self) -> Any:
+    def __iter__(self) -> Any:  
+        if self._is_bedrock:
+            # MUST be reside in a separate function so that the yield statement doesn't implicitly return its own iterator and overriding self
+            return self._iter_bedrock()
         return self
 
+    def _iter_bedrock(self) -> Any:
+        # botocore EventStream doesn't have a __next__ method so iterate over the wrapped object in place
+        for event in self.__wrapped__: # type: ignore
+            if (self._bedrock_from_stream):
+                self._evaluate_chunk(event)
+            else:
+                chunk = event.get('chunk') # type: ignore
+                if chunk:
+                    decode = chunk.get('bytes').decode() # type: ignore
+                    self._evaluate_chunk(decode)
+            yield event
+
+        self._stop_iteration()
+
     def __aiter__(self) -> Any:
         return self
 
@@ -460,7 +542,7 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
             self._first_token = False
 
         if self._log_prompt_and_response:
-            self._responses.append(chunk.to_json())
+            self._responses.append(self.chunk_to_json(chunk))
 
         if self._process_chunk:
             self._process_chunk(chunk, self._ingest)
@@ -475,11 +557,21 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
 
         self._instrumentor._ingest_units(self._ingest)
 
+    @staticmethod
+    def chunk_to_json(chunk: Any) -> str:
+        if hasattr(chunk, "to_json"):
+            return str(chunk.to_json())
+        elif isinstance(chunk, bytes):
+            return chunk.decode()
+        elif isinstance(chunk, str):
+            return chunk
+        else:
+            # assume dict
+            return json.dumps(chunk)
 
 global _instrumentor
 _instrumentor: PayiInstrumentor
 
-
 def payi_instrument(
     payi: Optional[Union[Payi, AsyncPayi]] = None,
     instruments: Optional[Set[Instruments]] = None,
@@ -520,7 +612,6 @@ def ingest(
 
     return _ingest
 
-
 def proxy(
     limit_ids: Optional["list[str]"] = None,
     request_tags: Optional["list[str]"] = None,

{payi-0.1.0a40.dist-info → payi-0.1.0a41.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: payi
-Version: 0.1.0a40
+Version: 0.1.0a41
 Summary: The official Python library for the payi API
 Project-URL: Homepage, https://github.com/Pay-i/pay-i-python
 Project-URL: Repository, https://github.com/Pay-i/pay-i-python

{payi-0.1.0a40.dist-info → payi-0.1.0a41.dist-info}/RECORD

@@ -2,16 +2,16 @@ payi/__init__.py,sha256=_eeZx9fx2Wp81adXh7qrpkmXCso7TiRSvIlLkQ0sQhY,2399
 payi/_base_client.py,sha256=CiHJoJuzSweHB7oidXWIlWdU0vrLCpPc9hMlj_S-anE,68128
 payi/_client.py,sha256=aUtMEmV02nTs3_pYYAR-OchCkofUHeXhhRs43tyDHLE,18760
 payi/_compat.py,sha256=VWemUKbj6DDkQ-O4baSpHVLJafotzeXmCQGJugfVTIw,6580
-payi/_constants.py,sha256=JE8kyZa2Q4NK_i4fO--8siEYTzeHnT0fYbOFDgDP4uk,464
+payi/_constants.py,sha256=S14PFzyN9-I31wiV7SmIlL5Ga0MLHxdvegInGdXH7tM,462
 payi/_exceptions.py,sha256=ItygKNrNXIVY0H6LsGVZvFuAHB3Vtm_VZXmWzCnpHy0,3216
 payi/_files.py,sha256=mf4dOgL4b0ryyZlbqLhggD3GVgDf6XxdGFAgce01ugE,3549
-payi/_models.py,sha256=B6f-C-F-PbDp3jRKCLksaAS9osC2g1xs7DpoZV1dlUE,28659
+payi/_models.py,sha256=uZvPAaaeDwCYwvB-yq7nxZnZ70I2Na_KjSAqaPQWfh0,28659
 payi/_qs.py,sha256=AOkSz4rHtK4YI3ZU_kzea-zpwBUgEY8WniGmTPyEimc,4846
 payi/_resource.py,sha256=j2jIkTr8OIC8sU6-05nxSaCyj4MaFlbZrwlyg4_xJos,1088
 payi/_response.py,sha256=CfrNS_3wbL8o9dRyRVfZQ5E1GUlA4CUIUEK8olmfGqE,28777
 payi/_streaming.py,sha256=Z_wIyo206T6Jqh2rolFg2VXZgX24PahLmpURp0-NssU,10092
 payi/_types.py,sha256=2mbMK86K3W1aMTW7sOGQ-VND6-A2IuXKm8p4sYFztBU,6141
-payi/_version.py,sha256=70Q3C2mH_Y5XHCDAYiBgNkVy3fyBDWf8eUqRwndh9Zc,165
+payi/_version.py,sha256=4Z7oXmyN8CqDXoYFdr9Fz53f3-pCQaVF2JTZFXE0NCo,165
 payi/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 payi/_utils/__init__.py,sha256=PNZ_QJuzZEgyYXqkO1HVhGkj5IU9bglVUcw7H-Knjzw,2062
 payi/_utils/_logs.py,sha256=fmnf5D9TOgkgZKfgYmSa3PiUc3SZgkchn6CzJUeo0SQ,768
@@ -23,12 +23,13 @@ payi/_utils/_transform.py,sha256=Dkkyr7OveGmOolepcvXmVJWE3kqim4b0nM0h7yWbgeY,134
 payi/_utils/_typing.py,sha256=nTJz0jcrQbEgxwy4TtAkNxuU0QHHlmc6mQtA6vIR8tg,4501
 payi/_utils/_utils.py,sha256=8UmbPOy_AAr2uUjjFui-VZSrVBHRj6bfNEKRp5YZP2A,12004
 payi/lib/.keep,sha256=wuNrz-5SXo3jJaJOJgz4vFHM41YH_g20F5cRQo0vLes,224
-payi/lib/AnthropicInstrumentor.py,sha256=XMLr8BIrBo7U3VPBnhyp9Wj-icZmu0MGVvvyDPmlF4c,4549
-payi/lib/Instruments.py,sha256=cyL2jxjpRluP9rN8Vn1nmVXq2NNLdZuFIsHMQWWqat4,115
-payi/lib/OpenAIInstrumentor.py,sha256=G3H8s1bNnq-LECAHjCkQDTG59pwnLPXKQKlQbhqgs4k,3998
+payi/lib/AnthropicInstrumentor.py,sha256=uiobG6N6QOKyhyS_kSZPCrKCnsRnKSLPBPNn3sfcUYg,4771
+payi/lib/BedrockInstrumentor.py,sha256=r5h01WJqx9PGszpwKWwVgPdCqUlWIbf6IiS01drN0qc,10684
+payi/lib/Instruments.py,sha256=bapmVS9jbHtFknXCKDzsFFWvf5XLtzEpdlvI7iEWY-o,147
+payi/lib/OpenAIInstrumentor.py,sha256=qYZJ-m723GCK3nGmI-JmRrNMpcDWmr4t4KYR1zWP-Ug,4195
 payi/lib/Stopwatch.py,sha256=vFyGVRvkppamP7W0IuZyypKLMIaqjhB7fcRG0dNyfnQ,757
 payi/lib/helpers.py,sha256=ZgkY8UE2YRc7ok2Pmxg_T9UMqKI8D8542JY3CP8RZCM,1597
-payi/lib/instrument.py,sha256=rLSwXRQZkTmMpWGpTzgZ42_4BOxJjGrTR3hwpF8XHro,20279
+payi/lib/instrument.py,sha256=FH2bcmcQyFCd6GAc7tiikag9RwPfd5FR5QCCNMomt4E,23462
 payi/resources/__init__.py,sha256=isHGXSl9kOrZDduKrX3UenTwrdTpuKJVBjw6NYSBV20,3592
 payi/resources/billing_models.py,sha256=5w3RfGXtGlyq5vbTw6hQrx1UlzRBtlq8ArcFlf5e3TY,20152
 payi/resources/ingest.py,sha256=SvQspsYled4_ErOZKzVtazBIk0tUC1e34Lw8qw4SNEM,15484
@@ -109,7 +110,7 @@ payi/types/requests/request_result.py,sha256=phYQiqhwNaR9igP-Fhs34Y-__dlT7L4wq-r
 payi/types/shared/__init__.py,sha256=-xz5dxK5LBjLnsi2LpLq5btaGDFp-mSjJ0y2qKy0Yus,264
 payi/types/shared/evaluation_response.py,sha256=ejEToMA57PUu1SldEtJ5z9r4fAO3U0tvdjbsyIoVX1s,214
 payi/types/shared/pay_i_common_models_budget_management_cost_details_base.py,sha256=XmIzJXy4zAi-mfrDvEXiYjO3qF1EvugGUl-Gijj4TA4,268
-payi-0.1.0a40.dist-info/METADATA,sha256=qTLzlfTKFZ8cEWce30AegO7WzqqjejiYc6kJi3SqYN8,12625
-payi-0.1.0a40.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-payi-0.1.0a40.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
-payi-0.1.0a40.dist-info/RECORD,,
+payi-0.1.0a41.dist-info/METADATA,sha256=QTeHU7hACFAIb9eEk3z9jPIbSjQwiRcuCrNtzIkr3FU,12625
+payi-0.1.0a41.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+payi-0.1.0a41.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
+payi-0.1.0a41.dist-info/RECORD,,