payi 0.1.0a77__py3-none-any.whl → 0.1.0a78__py3-none-any.whl

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

payi/lib/AnthropicInstrumentor.py

@@ -8,7 +8,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.lib.helpers import PayiCategories
 from payi.types.ingest_units_params import Units
 
-from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
+from .instrument import _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
 
 
 class AnthropicInstrumentor:
@@ -32,7 +32,7 @@ class AnthropicInstrumentor:
             wrap_function_wrapper(
                 "anthropic.resources.messages",
                 "Messages.stream",
-                messages_wrapper(instrumentor),
+                stream_messages_wrapper(instrumentor),
             )
 
             wrap_function_wrapper(
@@ -44,7 +44,7 @@ class AnthropicInstrumentor:
             wrap_function_wrapper(
                 "anthropic.resources.messages",
                 "AsyncMessages.stream",
-                amessages_wrapper(instrumentor),
+                astream_messages_wrapper(instrumentor),
             )
 
         except Exception as e:
@@ -61,7 +61,7 @@ def messages_wrapper(
     **kwargs: Any,
 ) -> Any:
     return instrumentor.invoke_wrapper(
-        _AnthropicProviderRequest(instrumentor, instance),
+        _AnthropicProviderRequest(instrumentor=instrumentor, streaming_type=_StreamingType.iterator, instance=instance),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -69,6 +69,23 @@ def messages_wrapper(
         kwargs,
     )
 
+@_PayiInstrumentor.payi_wrapper
+def stream_messages_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return instrumentor.invoke_wrapper(
+        _AnthropicProviderRequest(instrumentor=instrumentor, streaming_type=_StreamingType.stream_manager, instance=instance),
+        _IsStreaming.true,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
 @_PayiInstrumentor.payi_awrapper
 async def amessages_wrapper(
     instrumentor: _PayiInstrumentor,
@@ -78,7 +95,7 @@ async def amessages_wrapper(
     **kwargs: Any,
 ) -> Any:
     return await instrumentor.async_invoke_wrapper(
-        _AnthropicProviderRequest(instrumentor, instance),
+        _AnthropicProviderRequest(instrumentor=instrumentor, streaming_type=_StreamingType.iterator, instance=instance),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -86,12 +103,30 @@ async def amessages_wrapper(
         kwargs,
     )
 
+@_PayiInstrumentor.payi_awrapper
+async def astream_messages_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return await instrumentor.async_invoke_wrapper(
+        _AnthropicProviderRequest(instrumentor=instrumentor, streaming_type=_StreamingType.stream_manager, instance=instance),
+        _IsStreaming.true,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
 class _AnthropicProviderRequest(_ProviderRequest):
-    def __init__(self, instrumentor: _PayiInstrumentor, instance: Any = None) -> None:
+    def __init__(self, instrumentor: _PayiInstrumentor, streaming_type: _StreamingType, instance: Any = None) -> None:
         self._vertex: bool = AnthropicInstrumentor.is_vertex(instance)
         super().__init__(
             instrumentor=instrumentor,
-            category=PayiCategories.google_vertex if self._vertex else PayiCategories.anthropic
+            category=PayiCategories.google_vertex if self._vertex else PayiCategories.anthropic,
+            streaming_type=streaming_type,
             )
 
     @override
@@ -209,5 +244,4 @@ 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
     
-    return False, 0
-
+    return False, 0

payi/lib/BedrockInstrumentor.py

@@ -11,7 +11,7 @@ from payi.lib.helpers import PayiCategories, PayiHeaderNames, payi_aws_bedrock_u
 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
+from .instrument import _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
 
 _supported_model_prefixes = ["meta.llama3", "anthropic.", "amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro"]
 
@@ -216,7 +216,11 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 
 class _BedrockProviderRequest(_ProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor=instrumentor, category=PayiCategories.aws_bedrock)
+        super().__init__(
+            instrumentor=instrumentor,
+            category=PayiCategories.aws_bedrock,
+            streaming_type=_StreamingType.iterator,
+            )
 
     @override
     def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:

payi/lib/GoogleGenAiInstrumentor.py

@@ -0,0 +1,370 @@
+import json
+import math
+import logging
+from typing import Any, List, Union, Optional, Sequence
+from typing_extensions import override
+
+from wrapt import wrap_function_wrapper  # type: ignore
+
+from payi.lib.helpers import PayiCategories
+from payi.types.ingest_units_params import Units
+
+from .instrument import _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
+
+
+class GoogleGenAiInstrumentor:
+    @staticmethod
+    def instrument(instrumentor: _PayiInstrumentor) -> None:
+        try:
+            wrap_function_wrapper(
+                "google.genai.models",
+                "Models.generate_content",
+                generate_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "google.genai.models",
+                "Models.generate_content_stream",
+                generate_stream_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "google.genai.models",
+                "AsyncModels.generate_content",
+                agenerate_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "google.genai.models",
+                "AsyncModels.generate_content_stream",
+                agenerate_stream_wrapper(instrumentor),
+            )
+
+        except Exception as e:
+            logging.debug(f"Error instrumenting vertex: {e}")
+            return
+
+@_PayiInstrumentor.payi_wrapper
+def generate_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return instrumentor.invoke_wrapper(
+        _GoogleGenAiRequest(instrumentor),
+        _IsStreaming.false,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+@_PayiInstrumentor.payi_wrapper
+def generate_stream_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return instrumentor.invoke_wrapper(
+        _GoogleGenAiRequest(instrumentor),
+        _IsStreaming.true,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+@_PayiInstrumentor.payi_awrapper
+async def agenerate_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return await instrumentor.async_invoke_wrapper(
+        _GoogleGenAiRequest(instrumentor),
+        _IsStreaming.false,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+@_PayiInstrumentor.payi_wrapper
+async def agenerate_stream_wrapper(
+    instrumentor: _PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return await instrumentor.async_invoke_wrapper(
+        _GoogleGenAiRequest(instrumentor),
+        _IsStreaming.true,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+def count_chars_skip_spaces(text: str) -> int:
+    return sum(1 for c in text if not c.isspace())
+
+class _GoogleGenAiRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(
+            instrumentor=instrumentor,
+            category=PayiCategories.google_vertex,
+            streaming_type=_StreamingType.generator,
+            )
+        self._prompt_character_count = 0
+        self._candiates_character_count = 0
+
+    @override
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
+        from google.genai.types import Content, PIL_Image, Part # type: ignore #  noqa: F401  I001
+
+        if not kwargs:
+            return True
+
+        model: str = kwargs.get("model", "")
+        self._ingest["resource"] = "google." + model
+
+        value: Union[ # type: ignore
+            Content,
+            str,
+            PIL_Image,
+            Part,
+            List[Union[str, PIL_Image, Part]],
+        ] = kwargs.get("contents", None)  # type: ignore 
+
+        items: List[Union[str, Image, Part]] = [] # type: ignore #  noqa: F401  I001
+
+        if not value:
+            raise TypeError("value must not be empty")
+
+        if isinstance(value, Content):
+            items = value.parts # type: ignore
+        if isinstance(value, (str, PIL_Image, Part)):
+            items = [value] # type: ignore
+        if isinstance(value, list):
+            items = value # type: ignore
+
+        for item in items: # type: ignore
+            text = ""
+            if isinstance(item, Part):
+                d = item.to_json_dict() # type: ignore
+                if "text" in d:
+                    text = d["text"] # type: ignore
+            elif isinstance(item, str):
+                text = item
+
+            if text != "":
+                self._prompt_character_count += count_chars_skip_spaces(text) # type: ignore
+             
+        return True
+
+    @override
+    def process_request_prompt(self, prompt: 'dict[str, Any]', args: Sequence[Any], kwargs: 'dict[str, Any]') -> None:
+        from google.genai.types import Content, PIL_Image, Part, Tool, GenerateContentConfig, GenerateContentConfigDict, ToolConfig  # type: ignore #  noqa: F401  I001
+
+        key = "contents"
+
+        if not kwargs:
+            return
+        
+        value: Union[ # type: ignore
+            Content,
+            str,
+            PIL_Image,
+            Part,
+            List[Union[str, PIL_Image, Part]],
+        ] = kwargs.get("contents", None)  # type: ignore
+
+        items: List[Union[str, Image, Part]] = [] # type: ignore #  noqa: F401  I001
+
+        if not value:
+            return
+
+        if isinstance(value, str):
+            prompt[key] = Content(parts=[Part.from_text(text=value)]).to_json_dict() # type: ignore
+        elif isinstance(value, (PIL_Image, Part)):
+            prompt[key] = Content(parts=[value]).to_json_dict() # type: ignore
+        elif isinstance(value, Content):
+            prompt[key] = value.to_json_dict() # type: ignore
+        elif isinstance(value, list):
+            items = value # type: ignore
+            parts = []
+
+            for item in items: # type: ignore
+                if isinstance(item, str):
+                    parts.append(Part.from_text(text=item)) # type: ignore
+                elif isinstance(item, Part):
+                    parts.append(item) # type: ignore
+                # elif isinstance(item, PIL_Image): TODO
+                #     parts.append(Part.from_image(item)) # type: ignore
+                        
+            prompt[key] = Content(parts=parts).to_json_dict() # type: ignore
+
+        # tools: Optional[list[Tool]] = kwargs.get("tools", None)  # type: ignore
+        # if tools:
+        #     t: list[dict[Any, Any]] = []
+        #     for tool in tools: # type: ignore
+        #         if isinstance(tool, Tool):
+        #             t.append(tool.text=())  # type: ignore
+        #     if t:
+        #         prompt["tools"] = t
+        config_kwarg = kwargs.get("config", None)  # type: ignore
+        if config_kwarg is None:
+            return
+        
+        config: GenerateContentConfigDict = {}
+        if isinstance(config_kwarg, GenerateContentConfig):
+            config = config_kwarg.to_json_dict()  # type: ignore
+        else:
+            config = config_kwarg
+        
+        tools = config.get("tools", None)  # type: ignore
+        if isinstance(tools, list):
+            t: list[dict[str, object]] = []
+            for tool in tools: # type: ignore
+                if isinstance(tool, Tool):
+                    t.append(tool.to_json_dict())  # type: ignore
+            if t:
+                prompt["tools"] = t
+
+        tool_config = config.get("tool_config", None)  # type: ignore
+        if isinstance(tool_config, ToolConfig):
+            prompt["tool_config"] = tool_config.to_json_dict()  # type: ignore
+        elif isinstance(tool_config, dict):
+            prompt["tool_config"] = tool_config
+
+    @override
+    def process_chunk(self, chunk: Any) -> bool:
+        response_dict: dict[str, Any] = chunk.to_json_dict()
+        if "provider_response_id" not in self._ingest:
+            id = response_dict.get("response_id", None)
+            if id:
+                self._ingest["provider_response_id"] = id
+
+        model: str = response_dict.get("model_version", "")
+
+        self._ingest["resource"] = "google." + model
+
+        for candidate in response_dict.get("candidates", []):
+            parts = candidate.get("content", {}).get("parts", [])
+            for part in parts:
+                self._candiates_character_count += count_chars_skip_spaces(part.get("text", ""))
+
+        usage = response_dict.get("usage_metadata", {})
+        if usage and "prompt_token_count" in usage and "candidates_token_count" in usage:
+            self._compute_usage(response_dict, streaming_candidates_characters=self._candiates_character_count)
+
+        return True
+    
+    @staticmethod
+    def _is_character_billing_model(model: str) -> bool:
+        return model.startswith("gemini-1.")
+    
+    @override
+    def process_synchronous_response(
+        self,
+        response: Any,
+        log_prompt_and_response: bool,
+        kwargs: Any) -> Any:
+        response_dict = response.to_json_dict()
+
+        self._ingest["provider_response_id"] = response_dict["response_id"]
+        self._ingest["resource"] = "google." + response_dict["model_version"]
+
+        self._compute_usage(response_dict)
+        
+        if log_prompt_and_response:
+            self._ingest["provider_response_json"] = [json.dumps(response_dict)]
+
+        return None
+
+    def add_units(self, key: str, input: Optional[int] = None, output: Optional[int] = None) -> None:
+        if key not in self._ingest["units"]:
+            self._ingest["units"][key] = {}
+        if input is not None:
+            self._ingest["units"][key]["input"] = input
+        if output is not None:
+            self._ingest["units"][key]["output"] = output
+
+    def _compute_usage(self, response_dict: 'dict[str, Any]', streaming_candidates_characters: Optional[int] = None) -> None:
+        usage = response_dict.get("usage_metadata", {})
+        input = usage.get("prompt_token_count", 0)
+
+        prompt_tokens_details: list[dict[str, Any]] = usage.get("prompt_tokens_details")
+        candidates_tokens_details: list[dict[str, Any]] = usage.get("candidates_tokens_details")
+
+        model: str = response_dict.get("model_version", "")
+
+        if self._is_character_billing_model(model):
+            # gemini 1.0 and 1.5 units are reported in characters, per second, per image, etc...
+            large_context = "" if input < 128000 else "_large_context"
+        
+            for details in prompt_tokens_details:
+                modality = details.get("modality", "")
+                if not modality:
+                    continue
+
+                modality_token_count = details.get("token_count", 0)
+                if modality == "TEXT":
+                    input = self._prompt_character_count
+                    if input == 0:
+                        # back up calc if nothing was calculated from the prompt
+                        input = response_dict["usage_metadata"]["prompt_token_count"] * 4
+
+                    output = 0
+                    if streaming_candidates_characters is None:
+                        for candidate in response_dict.get("candidates", []):
+                            parts = candidate.get("content", {}).get("parts", [])
+                            for part in parts:
+                                output += count_chars_skip_spaces(part.get("text", ""))
+
+                        if output == 0:
+                            # back up calc if no parts
+                            output = response_dict["usage_metadata"]["candidates_token_count"] * 4
+                    else:
+                        output = streaming_candidates_characters
+
+                    self._ingest["units"]["text"+large_context] = Units(input=input, output=output)
+
+                elif modality == "IMAGE":
+                    num_images = math.ceil(modality_token_count / 258)
+                    self.add_units("vision"+large_context, input=num_images)
+
+                elif modality == "VIDEO":
+                    video_seconds = math.ceil(modality_token_count / 285)
+                    self.add_units("video"+large_context, input=video_seconds)
+
+                elif modality == "AUDIO":
+                    audio_seconds = math.ceil(modality_token_count / 25)
+                    self.add_units("audio"+large_context, input=audio_seconds)
+
+        elif model.startswith("gemini-2.0"):
+            for details in prompt_tokens_details:
+                modality = details.get("modality", "")
+                if not modality:
+                    continue
+
+                modality_token_count = details.get("token_count", 0)
+                if modality == "IMAGE":
+                    self.add_units("vision", input=modality_token_count)
+                elif modality in ("VIDEO", "AUDIO", "TEXT"):
+                    self.add_units(modality.lower(), input=modality_token_count)
+            for details in candidates_tokens_details:
+                modality = details.get("modality", "")
+                if not modality:
+                    continue
+
+                modality_token_count = details.get("token_count", 0)
+                if modality in ("VIDEO", "AUDIO", "TEXT", "IMAGE"):
+                    self.add_units(modality.lower(), output=modality_token_count)

payi/lib/OpenAIInstrumentor.py

@@ -10,7 +10,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.lib.helpers import PayiCategories, PayiHeaderNames
 from payi.types.ingest_units_params import Units
 
-from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
+from .instrument import _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
 
 
 class OpenAiInstrumentor:
@@ -178,7 +178,11 @@ class _OpenAiProviderRequest(_ProviderRequest):
     responses_input_tokens_details_key: str = "input_tokens_details"
 
     def __init__(self, instrumentor: _PayiInstrumentor, input_tokens_key: str, output_tokens_key: str, input_tokens_details_key: str) -> None:
-        super().__init__(instrumentor=instrumentor, category=PayiCategories.openai)
+        super().__init__(
+            instrumentor=instrumentor,
+            category=PayiCategories.openai,
+            streaming_type=_StreamingType.iterator,
+            )
         self._input_tokens_key = input_tokens_key
         self._output_tokens_key = output_tokens_key
         self._input_tokens_details_key = input_tokens_details_key

payi/lib/VertexInstrumentor.py

@@ -9,7 +9,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.lib.helpers import PayiCategories
 from payi.types.ingest_units_params import Units
 
-from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
+from .instrument import _IsStreaming, _StreamingType, _ProviderRequest, _PayiInstrumentor
 
 
 class VertexInstrumentor:
@@ -85,7 +85,11 @@ def count_chars_skip_spaces(text: str) -> int:
 
 class _GoogleVertexRequest(_ProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor=instrumentor, category=PayiCategories.google_vertex)
+        super().__init__(
+            instrumentor=instrumentor,
+            category=PayiCategories.google_vertex,
+            streaming_type=_StreamingType.generator,
+            )
         self._prompt_character_count = 0
         self._candiates_character_count = 0
 
@@ -110,11 +114,10 @@ class _GoogleVertexRequest(_ProviderRequest):
             raise TypeError("value must not be empty")
 
         if isinstance(value, Content):
-            return value.parts # type: ignore
+            items = value.parts # type: ignore
         if isinstance(value, (str, Image, Part)):
             items = [value] # type: ignore
-
-        elif isinstance(value, list):
+        if isinstance(value, list):
             items = value # type: ignore
 
         for item in items: # type: ignore

payi/lib/instrument.py

@@ -27,11 +27,12 @@ from .Stopwatch import Stopwatch
 
 
 class _ProviderRequest:
-    def __init__(self, instrumentor: '_PayiInstrumentor', category: str):
+    def __init__(self, instrumentor: '_PayiInstrumentor', category: str, streaming_type: '_StreamingType'):
         self._instrumentor: '_PayiInstrumentor' = instrumentor
         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
 
     def process_chunk(self, _chunk: Any) -> bool:
         return True
@@ -55,6 +56,10 @@ class _ProviderRequest:
     def process_exception(self, exception: Exception, kwargs: Any, ) -> bool: # noqa: ARG002
         self.exception_to_semantic_failure(exception)
         return True
+    
+    @property
+    def streaming_type(self) -> '_StreamingType':
+        return self._streaming_type
 
     def exception_to_semantic_failure(self, e: Exception) -> None:
         exception_str = f"{type(e).__name__}"
@@ -111,6 +116,11 @@ class _IsStreaming(Enum):
     true = 1 
     kwargs = 2
 
+class _StreamingType(Enum):
+    generator = 0
+    iterator = 1
+    stream_manager = 2
+
 class _TrackContext:
     def __init__(
         self,
@@ -192,6 +202,7 @@ class _PayiInstrumentor:
         self._instrument_anthropic()
         self._instrument_aws_bedrock()
         self._instrument_google_vertex()
+        self._instrument_google_genai()
 
     def _instrument_specific(self, instruments: Set[str]) -> None:
         if PayiCategories.openai in instruments or PayiCategories.azure_openai in instruments:
@@ -202,6 +213,7 @@ class _PayiInstrumentor:
             self._instrument_aws_bedrock()
         if PayiCategories.google_vertex in instruments:
             self._instrument_google_vertex()
+            self._instrument_google_genai()
 
     def _instrument_openai(self) -> None:
         from .OpenAIInstrumentor import OpenAiInstrumentor
@@ -239,6 +251,14 @@ class _PayiInstrumentor:
         except Exception as e:
             logging.error(f"Error instrumenting Google Vertex: {e}")
 
+    def _instrument_google_genai(self) -> None:
+        from .GoogleGenAiInstrumentor import GoogleGenAiInstrumentor
+
+        try:
+            GoogleGenAiInstrumentor.instrument(self)
+
+        except Exception as e:
+            logging.error(f"Error instrumenting Google GenAi: {e}")
 
     def _process_ingest_units(self, ingest_units: IngestUnitsParams, log_data: 'dict[str, str]') -> bool:
         if int(ingest_units.get("http_status_code") or 0) < 400:
@@ -690,25 +710,30 @@ class _PayiInstrumentor:
             raise e
 
         if stream:
-            if request.is_vertex():
+            if request.streaming_type == _StreamingType.generator:
                 return _GeneratorWrapper(
                     generator=response,
                     instance=instance,
                     instrumentor=self,
                     stopwatch=sw,
                     request=request,
-                    log_prompt_and_response=self._log_prompt_and_response)
-
-            stream_result = ChatStreamWrapper(
-                response=response,
-                instance=instance,
-                instrumentor=self,
-                log_prompt_and_response=self._log_prompt_and_response,
-                stopwatch=sw,
-                request=request,
-            )
-
-            return stream_result
+                    )
+            elif request.streaming_type == _StreamingType.stream_manager:
+                return _StreamManagerWrapper(
+                    stream_manager=response,
+                    instance=instance,
+                    instrumentor=self,
+                    stopwatch=sw,
+                    request=request,
+                )
+            else:
+                return _StreamIteratorWrapper(
+                    response=response,
+                    instance=instance,
+                    instrumentor=self,
+                    stopwatch=sw,
+                    request=request,
+                )
 
         sw.stop()
         duration = sw.elapsed_ms_int()
@@ -797,32 +822,40 @@ class _PayiInstrumentor:
             raise e
 
         if stream:
-            if request.is_vertex():
+            if request.streaming_type == _StreamingType.generator:
                 return _GeneratorWrapper(
                     generator=response,
                     instance=instance,
                     instrumentor=self,
                     stopwatch=sw,
                     request=request,
-                    log_prompt_and_response=self._log_prompt_and_response)
-
-            stream_result = ChatStreamWrapper(
-                response=response,
-                instance=instance,
-                instrumentor=self,
-                log_prompt_and_response=self._log_prompt_and_response,
-                stopwatch=sw,
-                request=request,
-            )
+                )
+            elif request.streaming_type == _StreamingType.stream_manager:
+                return _StreamManagerWrapper(
+                    stream_manager=response,
+                    instance=instance,
+                    instrumentor=self,
+                    stopwatch=sw,
+                    request=request,
+                )
+            else:
+                # request.streaming_type == _StreamingType.iterator
+                stream_result = _StreamIteratorWrapper(
+                    response=response,
+                    instance=instance,
+                    instrumentor=self,
+                    stopwatch=sw,
+                    request=request,
+                )
 
-            if request.is_bedrock():
-                if "body" in response:
-                    response["body"] = stream_result
-                else:
-                    response["stream"] = stream_result
-                return response
-            
-            return stream_result
+                if request.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()
@@ -986,7 +1019,7 @@ class _PayiInstrumentor:
 
         return _payi_awrapper
 
-class ChatStreamWrapper(ObjectProxy):  # type: ignore
+class _StreamIteratorWrapper(ObjectProxy):  # type: ignore
     def __init__(
         self,
         response: Any,
@@ -994,7 +1027,6 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         instrumentor: _PayiInstrumentor,
         stopwatch: Stopwatch,
         request: _ProviderRequest,
-        log_prompt_and_response: bool = True,
     ) -> None:
 
         bedrock_from_stream: bool = False
@@ -1016,7 +1048,6 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
 
         self._instrumentor = instrumentor
         self._stopwatch: Stopwatch = stopwatch
-        self._log_prompt_and_response: bool = log_prompt_and_response
         self._responses: list[str] = []
 
         self._request: _ProviderRequest = request
@@ -1091,7 +1122,7 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
             self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
             self._first_token = False
 
-        if self._log_prompt_and_response:
+        if self._instrumentor._log_prompt_and_response:
             self._responses.append(self.chunk_to_json(chunk))
 
         return self._request.process_chunk(chunk)
@@ -1101,7 +1132,7 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
         self._request._ingest["http_status_code"] = 200
 
-        if self._log_prompt_and_response:
+        if self._instrumentor._log_prompt_and_response:
             self._request._ingest["provider_response_json"] = self._responses
 
     async def _astop_iteration(self) -> None:
@@ -1124,6 +1155,35 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
             # assume dict
             return json.dumps(chunk)
 
+class _StreamManagerWrapper(ObjectProxy):  # type: ignore
+    def __init__(
+        self,
+        stream_manager: Any,  # type: ignore
+        instance: Any,
+        instrumentor: _PayiInstrumentor, 
+        stopwatch: Stopwatch,
+        request: _ProviderRequest,
+    ) -> None:
+        super().__init__(stream_manager)  # type: ignore
+
+        self._stream_manager = stream_manager  
+        self._instance = instance
+        self._instrumentor = instrumentor
+        self._stopwatch: Stopwatch = stopwatch
+        self._responses: list[str] = []
+        self._request: _ProviderRequest = request
+        self._first_token: bool = True
+        self._done: bool = False
+
+    def __enter__(self) -> _StreamIteratorWrapper:
+        return _StreamIteratorWrapper(
+            response=self.__wrapped__.__enter__(),  # type: ignore
+            instance=self._instance,
+            instrumentor=self._instrumentor,
+            stopwatch=self._stopwatch,
+            request=self._request,
+        )
+
 class _GeneratorWrapper:  # type: ignore
     def __init__(
         self,
@@ -1132,7 +1192,6 @@ class _GeneratorWrapper:  # type: ignore
         instrumentor: _PayiInstrumentor, 
         stopwatch: Stopwatch,
         request: _ProviderRequest,
-        log_prompt_and_response: bool = True,
     ) -> None:
         super().__init__()  # type: ignore
         
@@ -1140,7 +1199,7 @@ class _GeneratorWrapper:  # type: ignore
         self._instance = instance
         self._instrumentor = instrumentor
         self._stopwatch: Stopwatch = stopwatch
-        self._log_prompt_and_response: bool = log_prompt_and_response
+        self._log_prompt_and_response: bool = instrumentor._log_prompt_and_response
         self._responses: list[str] = []
         self._request: _ProviderRequest = request
         self._first_token: bool = True
@@ -1176,13 +1235,22 @@ class _GeneratorWrapper:  # type: ignore
             await self._process_async_stop_iteration()
             raise stop_exception
 
+    @staticmethod
+    def _chunk_to_dict(chunk: Any) -> 'dict[str, object]':
+        if hasattr(chunk, "to_json"):
+            return chunk.to_json() # type: ignore
+        elif hasattr(chunk, "to_json_dict"):  
+            return chunk.to_json_dict() # type: ignore
+        else:
+            return {}
+
     def _process_chunk(self, chunk: Any) -> Any:
         if self._first_token:
             self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
             self._first_token = False
             
         if self._log_prompt_and_response:
-            dict = chunk.to_dict() # type: ignore
+            dict = self._chunk_to_dict(chunk) 
             self._responses.append(json.dumps(dict))
                 
         self._request.process_chunk(chunk)

{payi-0.1.0a77.dist-info → payi-0.1.0a78.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: payi
-Version: 0.1.0a77
+Version: 0.1.0a78
 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.0a77.dist-info → payi-0.1.0a78.dist-info}/RECORD

@@ -11,7 +11,7 @@ payi/_resource.py,sha256=j2jIkTr8OIC8sU6-05nxSaCyj4MaFlbZrwlyg4_xJos,1088
 payi/_response.py,sha256=rh9oJAvCKcPwQFm4iqH_iVrmK8bNx--YP_A2a4kN1OU,28776
 payi/_streaming.py,sha256=Z_wIyo206T6Jqh2rolFg2VXZgX24PahLmpURp0-NssU,10092
 payi/_types.py,sha256=7jE5MoQQFVoVxw5vVzvZ2Ao0kcjfNOGsBgyJfLBEnMo,6195
-payi/_version.py,sha256=wGt5jNH0E97HeewXFLO6adzkx0gEYgIP8AihAT_HARo,165
+payi/_version.py,sha256=562XNlpNWouu5fR8P9024ptW4UWs2S750PSKhFzqfdY,165
 payi/pagination.py,sha256=k2356QGPOUSjRF2vHpwLBdF6P-2vnQzFfRIJQAHGQ7A,1258
 payi/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 payi/_utils/__init__.py,sha256=PNZ_QJuzZEgyYXqkO1HVhGkj5IU9bglVUcw7H-Knjzw,2062
@@ -25,13 +25,14 @@ payi/_utils/_transform.py,sha256=n7kskEWz6o__aoNvhFoGVyDoalNe6mJwp-g7BWkdj88,156
 payi/_utils/_typing.py,sha256=D0DbbNu8GnYQTSICnTSHDGsYXj8TcAKyhejb0XcnjtY,4602
 payi/_utils/_utils.py,sha256=ts4CiiuNpFiGB6YMdkQRh2SZvYvsl7mAF-JWHCcLDf4,12312
 payi/lib/.keep,sha256=wuNrz-5SXo3jJaJOJgz4vFHM41YH_g20F5cRQo0vLes,224
-payi/lib/AnthropicInstrumentor.py,sha256=pr4spNSyTM_VMtdWwoZPCV2o0ERGppoMhHo5rVHdmI8,7609
-payi/lib/BedrockInstrumentor.py,sha256=4IfOwq_MSGlvlSMxAocVu0G-Ctoir6nSL9KtpGfga8I,13969
-payi/lib/OpenAIInstrumentor.py,sha256=xmm2jxuVtuBZUNgORwgCZqiAGE_e6E28IwMEy9i7-nc,17361
+payi/lib/AnthropicInstrumentor.py,sha256=q1Iq2NG06OHUwcwTgYjHa1_kwhbSxE3vrGk_fm5IGSM,8749
+payi/lib/BedrockInstrumentor.py,sha256=cIjHlPQmVEAixabUoT3mV8h50lfzdc1u8iLj0gQNukE,14076
+payi/lib/GoogleGenAiInstrumentor.py,sha256=y9FuEtuITaIac_N-WhmypAr-UF1lzRCc9YBRVwwwydQ,13411
+payi/lib/OpenAIInstrumentor.py,sha256=iMOriAm-2zcRNjQDkEAG-NfZhBJUnkBDrNNbO920ajc,17468
 payi/lib/Stopwatch.py,sha256=7OJlxvr2Jyb6Zr1LYCYKczRB7rDVKkIR7gc4YoleNdE,764
-payi/lib/VertexInstrumentor.py,sha256=R0s3aEW7haj4aC5t1lMeA9siOgYvxlKQShLJ-B5iiNM,11642
+payi/lib/VertexInstrumentor.py,sha256=Xmp5kRI8G0u2OieGmlNxIduwnqp8cAXO-fROpkdXbgs,11748
 payi/lib/helpers.py,sha256=K1KAfWrpPT1UUGNxspLe1lHzQjP3XV5Pkh9IU4pKMok,4624
-payi/lib/instrument.py,sha256=nUUEnyDJPNMgVvT1gRgSVX1T9D8IXvBG8UzUcmfM9ws,55743
+payi/lib/instrument.py,sha256=UpZ6SGg3YI9lSxmwH5ziwe1xt_ca6FS35CcNJpG9ONM,58214
 payi/resources/__init__.py,sha256=1rtrPLWbNt8oJGOp6nwPumKLJ-ftez0B6qwLFyfcoP4,2972
 payi/resources/ingest.py,sha256=8HNHEyfgIyJNqCh0rOhO9msoc61-8IyifJ6AbxjCrDg,22612
 payi/resources/categories/__init__.py,sha256=w5gMiPdBSzJA_qfoVtFBElaoe8wGf_O63R7R1Spr6Gk,1093
@@ -141,7 +142,7 @@ payi/types/use_cases/definitions/kpi_retrieve_response.py,sha256=uQXliSvS3k-yDYw
 payi/types/use_cases/definitions/kpi_update_params.py,sha256=jbawdWAdMnsTWVH0qfQGb8W7_TXe3lq4zjSRu44d8p8,373
 payi/types/use_cases/definitions/kpi_update_response.py,sha256=zLyEoT0S8d7XHsnXZYT8tM7yDw0Aze0Mk-_Z6QeMtc8,459
 payi/types/use_cases/definitions/limit_config_create_params.py,sha256=pzQza_16N3z8cFNEKr6gPbFvuGFrwNuGxAYb--Kbo2M,449
-payi-0.1.0a77.dist-info/METADATA,sha256=RXzXLw8b3r97mwWsST8Zt8u8bwyU331k36T83w2AaO8,15180
-payi-0.1.0a77.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-payi-0.1.0a77.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
-payi-0.1.0a77.dist-info/RECORD,,
+payi-0.1.0a78.dist-info/METADATA,sha256=kKY9frrRc7_WbcIUIphMDmRBHQVZgdSGH1eLkF26Pzk,15180
+payi-0.1.0a78.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+payi-0.1.0a78.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
+payi-0.1.0a78.dist-info/RECORD,,