payi 0.1.0a76__tar.gz → 0.1.0a78__tar.gz

payi-0.1.0a78/.release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+  ".": "0.1.0-alpha.78"
+}

{payi-0.1.0a76 → payi-0.1.0a78}/CHANGELOG.md

@@ -1,5 +1,27 @@
 # Changelog
 
+## 0.1.0-alpha.78 (2025-06-03)
+
+Full Changelog: [v0.1.0-alpha.77...v0.1.0-alpha.78](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.77...v0.1.0-alpha.78)
+
+### Features
+
+* **api:** google genai, anthropic stream() ([#306](https://github.com/Pay-i/pay-i-python/issues/306)) ([341da40](https://github.com/Pay-i/pay-i-python/commit/341da40fb228cb83994bec12e7cda457f5452692))
+
+## 0.1.0-alpha.77 (2025-06-03)
+
+Full Changelog: [v0.1.0-alpha.76...v0.1.0-alpha.77](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.76...v0.1.0-alpha.77)
+
+### Features
+
+* **client:** add follow_redirects request option ([3308085](https://github.com/Pay-i/pay-i-python/commit/3308085121c320d80931d6749f36ecee049be669))
+
+
+### Chores
+
+* **docs:** remove reference to rye shell ([eb25599](https://github.com/Pay-i/pay-i-python/commit/eb255999808df39a9eaaf5b91825c8fe150dbf55))
+* **docs:** remove unnecessary param examples ([4d34a34](https://github.com/Pay-i/pay-i-python/commit/4d34a343cd19fc091fb836d411034d9ad62bf564))
+
 ## 0.1.0-alpha.76 (2025-05-31)
 
 Full Changelog: [v0.1.0-alpha.75...v0.1.0-alpha.76](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.75...v0.1.0-alpha.76)

{payi-0.1.0a76 → payi-0.1.0a78}/CONTRIBUTING.md

@@ -17,8 +17,7 @@ $ rye sync --all-features
 You can then run scripts using `rye run python script.py` or by activating the virtual environment:
 
 ```sh
-$ rye shell
-# or manually activate - https://docs.python.org/3/library/venv.html#how-venvs-work
+# Activate the virtual environment - https://docs.python.org/3/library/venv.html#how-venvs-work
 $ source .venv/bin/activate
 
 # now you can omit the `rye run` prefix

{payi-0.1.0a76 → payi-0.1.0a78}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: payi
-Version: 0.1.0a76
+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
@@ -188,12 +188,7 @@ client = Payi()
 use_case_definition = client.use_cases.definitions.create(
     description="x",
     name="x",
-    limit_config={
-        "max": 0,
-        "limit_tags": ["tag1", "tag2"],
-        "limit_type": "block",
-        "threshold": 0,
-    },
+    limit_config={"max": 0},
 )
 print(use_case_definition.limit_config)
 ```

{payi-0.1.0a76 → payi-0.1.0a78}/README.md

@@ -154,12 +154,7 @@ client = Payi()
 use_case_definition = client.use_cases.definitions.create(
     description="x",
     name="x",
-    limit_config={
-        "max": 0,
-        "limit_tags": ["tag1", "tag2"],
-        "limit_type": "block",
-        "threshold": 0,
-    },
+    limit_config={"max": 0},
 )
 print(use_case_definition.limit_config)
 ```

{payi-0.1.0a76 → payi-0.1.0a78}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "payi"
-version = "0.1.0-alpha.76"
+version = "0.1.0-alpha.78"
 description = "The official Python library for the payi API"
 dynamic = ["readme"]
 license = "Apache-2.0"

{payi-0.1.0a76 → payi-0.1.0a78}/src/payi/_base_client.py

@@ -960,6 +960,9 @@ class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]):
             if self.custom_auth is not None:
                 kwargs["auth"] = self.custom_auth
 
+            if options.follow_redirects is not None:
+                kwargs["follow_redirects"] = options.follow_redirects
+
             log.debug("Sending HTTP Request: %s %s", request.method, request.url)
 
             response = None
@@ -1460,6 +1463,9 @@ class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]):
             if self.custom_auth is not None:
                 kwargs["auth"] = self.custom_auth
 
+            if options.follow_redirects is not None:
+                kwargs["follow_redirects"] = options.follow_redirects
+
             log.debug("Sending HTTP Request: %s %s", request.method, request.url)
 
             response = None

{payi-0.1.0a76 → payi-0.1.0a78}/src/payi/_models.py

@@ -737,6 +737,7 @@ class FinalRequestOptionsInput(TypedDict, total=False):
     idempotency_key: str
     json_data: Body
     extra_json: AnyMapping
+    follow_redirects: bool
 
 
 @final
@@ -750,6 +751,7 @@ class FinalRequestOptions(pydantic.BaseModel):
     files: Union[HttpxRequestFiles, None] = None
     idempotency_key: Union[str, None] = None
     post_parser: Union[Callable[[Any], Any], NotGiven] = NotGiven()
+    follow_redirects: Union[bool, None] = None
 
     # It should be noted that we cannot use `json` here as that would override
     # a BaseModel method in an incompatible fashion.

{payi-0.1.0a76 → payi-0.1.0a78}/src/payi/_types.py

@@ -100,6 +100,7 @@ class RequestOptions(TypedDict, total=False):
     params: Query
     extra_json: AnyMapping
     idempotency_key: str
+    follow_redirects: bool
 
 
 # Sentinel class used until PEP 0661 is accepted
@@ -215,3 +216,4 @@ class _GenericAlias(Protocol):
 
 class HttpxSendArgs(TypedDict, total=False):
     auth: httpx.Auth
+    follow_redirects: bool

{payi-0.1.0a76 → payi-0.1.0a78}/src/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.76"  # x-release-please-version
+__version__ = "0.1.0-alpha.78"  # x-release-please-version

{payi-0.1.0a76 → payi-0.1.0a78}/src/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-0.1.0a76 → payi-0.1.0a78}/src/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-0.1.0a78/src/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-0.1.0a76 → payi-0.1.0a78}/src/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