payi 0.1.0a70__tar.gz → 0.1.0a72__tar.gz

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

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

{payi-0.1.0a70 → payi-0.1.0a72}/CHANGELOG.md

@@ -1,5 +1,32 @@
 # Changelog
 
+## 0.1.0-alpha.72 (2025-05-10)
+
+Full Changelog: [v0.1.0-alpha.71...v0.1.0-alpha.72](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.71...v0.1.0-alpha.72)
+
+### Bug Fixes
+
+* **package:** support direct resource imports ([314f86c](https://github.com/Pay-i/pay-i-python/commit/314f86cb0d46aa35628c9d059c1b1ca25d24d911))
+
+
+### Chores
+
+* **internal:** avoid errors for isinstance checks on proxies ([b935f82](https://github.com/Pay-i/pay-i-python/commit/b935f82c9de01353ec83b9c12c4ed6749e984c7d))
+
+## 0.1.0-alpha.71 (2025-05-03)
+
+Full Changelog: [v0.1.0-alpha.70...v0.1.0-alpha.71](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.70...v0.1.0-alpha.71)
+
+### Features
+
+* **api:** google vertex ([#294](https://github.com/Pay-i/pay-i-python/issues/294)) ([6d9e39b](https://github.com/Pay-i/pay-i-python/commit/6d9e39b2ee522589d4962c73f697f9cce76fdf6a))
+* **api:** manual updates ([7d11112](https://github.com/Pay-i/pay-i-python/commit/7d1111205f30fe78c9e2e833d3f99a62f95419e8))
+
+
+### Bug Fixes
+
+* lint ([e7d5a21](https://github.com/Pay-i/pay-i-python/commit/e7d5a210db38448629511b470a4d1b426368577f))
+
 ## 0.1.0-alpha.70 (2025-04-24)
 
 Full Changelog: [v0.1.0-alpha.69...v0.1.0-alpha.70](https://github.com/Pay-i/pay-i-python/compare/v0.1.0-alpha.69...v0.1.0-alpha.70)

{payi-0.1.0a70 → payi-0.1.0a72}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: payi
-Version: 0.1.0a70
+Version: 0.1.0a72
 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.0a70 → payi-0.1.0a72}/api.md

@@ -2,9 +2,12 @@
 
 ```python
 from payi.types import (
+    IngestUnits,
     PayICommonModelsBudgetManagementCostDetailsBase,
     PayICommonModelsBudgetManagementCreateLimitBase,
     PropertiesResponse,
+    XproxyError,
+    XproxyResult,
 )
 ```
 

{payi-0.1.0a70 → payi-0.1.0a72}/pyproject.toml

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

{payi-0.1.0a70 → payi-0.1.0a72}/src/payi/__init__.py

@@ -1,5 +1,7 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
+import typing as _t
+
 from . import types
 from ._types import NOT_GIVEN, Omit, NoneType, NotGiven, Transport, ProxiesTypes
 from ._utils import file_from_path
@@ -68,6 +70,9 @@ __all__ = [
     "DefaultAsyncHttpxClient",
 ]
 
+if not _t.TYPE_CHECKING:
+    from ._utils._resources_proxy import resources as resources
+
 _setup_logging()
 
 # Update the __module__ attribute for exported symbols so that

{payi-0.1.0a70 → payi-0.1.0a72}/src/payi/_utils/_proxy.py

@@ -46,7 +46,10 @@ class LazyProxy(Generic[T], ABC):
     @property  # type: ignore
     @override
     def __class__(self) -> type:  # pyright: ignore
-        proxied = self.__get_proxied__()
+        try:
+            proxied = self.__get_proxied__()
+        except Exception:
+            return type(self)
         if issubclass(type(proxied), LazyProxy):
             return type(proxied)
         return proxied.__class__

payi-0.1.0a72/src/payi/_utils/_resources_proxy.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Any
+from typing_extensions import override
+
+from ._proxy import LazyProxy
+
+
+class ResourcesProxy(LazyProxy[Any]):
+    """A proxy for the `payi.resources` module.
+
+    This is used so that we can lazily import `payi.resources` only when
+    needed *and* so that users can just import `payi` and reference `payi.resources`
+    """
+
+    @override
+    def __load__(self) -> Any:
+        import importlib
+
+        mod = importlib.import_module("payi.resources")
+        return mod
+
+
+resources = ResourcesProxy().__as_proxied__()

{payi-0.1.0a70 → payi-0.1.0a72}/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.70"  # x-release-please-version
+__version__ = "0.1.0-alpha.72"  # x-release-please-version

{payi-0.1.0a70 → payi-0.1.0a72}/src/payi/lib/AnthropicInstrumentor.py

@@ -1,10 +1,11 @@
 import logging
-from typing import Any, Union, Optional
+from typing import Any, Union, Optional, Sequence
 from typing_extensions import override
 
 import tiktoken
 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
@@ -53,7 +54,7 @@ def chat_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return instrumentor.chat_wrapper(
+    return instrumentor.invoke_wrapper(
         _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -70,7 +71,7 @@ async def achat_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return await instrumentor.achat_wrapper(
+    return await instrumentor.async_invoke_wrapper(
         _AnthropicProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -81,7 +82,7 @@ async def achat_wrapper(
 
 class _AnthropicProviderRequest(_ProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor=instrumentor, category="system.anthropic")
+        super().__init__(instrumentor=instrumentor, category=PayiCategories.anthropic)
 
     @override
     def process_chunk(self, chunk: Any) -> bool:
@@ -136,7 +137,7 @@ class _AnthropicProviderRequest(_ProviderRequest):
         return None
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
         self._ingest["resource"] = kwargs.get("model", "")
         messages = kwargs.get("messages")
         if messages:
@@ -166,7 +167,8 @@ class _AnthropicProviderRequest(_ProviderRequest):
                     self._ingest["http_status_code"] = status_code
 
             if not status_code:
-                return False
+                self.exception_to_semantic_failure(exception,)
+                return True
 
             if hasattr(exception, "request_id"):
                 request_id = getattr(exception, "request_id", None)

{payi-0.1.0a70 → payi-0.1.0a72}/src/payi/lib/BedrockInstrumentor.py

@@ -1,13 +1,13 @@
 import os
 import json
 import logging
-from typing import Any
+from typing import Any, Sequence
 from functools import wraps
 from typing_extensions import override
 
 from wrapt import ObjectProxy, wrap_function_wrapper  # type: ignore
 
-from payi.lib.helpers import PayiHeaderNames, payi_aws_bedrock_url
+from payi.lib.helpers import PayiCategories, PayiHeaderNames, payi_aws_bedrock_url
 from payi.types.ingest_units_params import Units, IngestUnitsParams
 from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
@@ -148,7 +148,7 @@ def wrap_invoke(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
 
         if _is_supported_model(modelId):
-            return instrumentor.chat_wrapper(
+            return instrumentor.invoke_wrapper(
                 _BedrockInvokeSynchronousProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.false,
                 wrapped,
@@ -166,7 +166,7 @@ def wrap_invoke_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         modelId: str = kwargs.get("modelId", "") # type: ignore
 
         if _is_supported_model(modelId):
-            return instrumentor.chat_wrapper(
+            return instrumentor.invoke_wrapper(
                 _BedrockInvokeStreamingProviderRequest(instrumentor=instrumentor, model_id=modelId),
                 _IsStreaming.true,
                 wrapped,
@@ -184,7 +184,7 @@ def wrap_converse(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         modelId:str = kwargs.get("modelId", "") # type: ignore
 
         if _is_supported_model(modelId):
-            return instrumentor.chat_wrapper(
+            return instrumentor.invoke_wrapper(
                 _BedrockConverseSynchronousProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.false,
                 wrapped,
@@ -202,7 +202,7 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
         modelId: str = kwargs.get("modelId", "") # type: ignore
 
         if _is_supported_model(modelId):
-            return instrumentor.chat_wrapper(
+            return instrumentor.invoke_wrapper(
                 _BedrockConverseStreamingProviderRequest(instrumentor=instrumentor),
                 _IsStreaming.true,
                 wrapped,
@@ -216,10 +216,10 @@ def wrap_converse_stream(instrumentor: _PayiInstrumentor, wrapped: Any) -> Any:
 
 class _BedrockProviderRequest(_ProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor=instrumentor, category="system.aws.bedrock")
+        super().__init__(instrumentor=instrumentor, category=PayiCategories.aws_bedrock)
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
         # boto3 doesn't allow extra_headers
         kwargs.pop("extra_headers", None)
         self._ingest["resource"] = kwargs.get("modelId", "")

{payi-0.1.0a70 → payi-0.1.0a72}/src/payi/lib/OpenAIInstrumentor.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Any, Union, Optional
+from typing import Any, Union, Optional, Sequence
 from typing_extensions import override
 from importlib.metadata import version
 
@@ -63,7 +63,7 @@ def embeddings_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return instrumentor.chat_wrapper(
+    return instrumentor.invoke_wrapper(
         _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
@@ -80,7 +80,7 @@ async def aembeddings_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return await instrumentor.achat_wrapper(
+    return await instrumentor.async_invoke_wrapper(
         _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
@@ -97,7 +97,7 @@ def chat_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return instrumentor.chat_wrapper(
+    return instrumentor.invoke_wrapper(
         _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -114,7 +114,7 @@ async def achat_wrapper(
     *args: Any,
     **kwargs: Any,
 ) -> Any:
-    return await instrumentor.achat_wrapper(
+    return await instrumentor.async_invoke_wrapper(
         _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
@@ -125,10 +125,10 @@ async def achat_wrapper(
 
 class _OpenAiProviderRequest(_ProviderRequest):
     def __init__(self, instrumentor: _PayiInstrumentor):
-        super().__init__(instrumentor=instrumentor, category="system.openai")
+        super().__init__(instrumentor=instrumentor, category=PayiCategories.openai)
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool: # type: ignore
         self._ingest["resource"] = kwargs.get("model", "")
 
         if not (instance and hasattr(instance, "_client")) or OpenAiInstrumentor.is_azure(instance) is False:
@@ -172,7 +172,8 @@ class _OpenAiProviderRequest(_ProviderRequest):
                     self._ingest["http_status_code"] = status_code
 
             if not status_code:
-                return False
+                self.exception_to_semantic_failure(exception,)
+                return True
 
             if hasattr(exception, "request_id"):
                 request_id = getattr(exception, "request_id", None)
@@ -232,8 +233,8 @@ class _OpenAiChatProviderRequest(_OpenAiProviderRequest):
         return send_chunk_to_client
 
     @override
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', kwargs: Any) -> bool:
-        result = super().process_request(instance, extra_headers, kwargs)
+    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
+        result = super().process_request(instance, extra_headers, args, kwargs)
         if result is False:
             return result
         

payi-0.1.0a72/src/payi/lib/VertexInstrumentor.py

@@ -0,0 +1,313 @@
+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, _ProviderRequest, _PayiInstrumentor
+
+
+class VertexInstrumentor:
+    @staticmethod
+    def instrument(instrumentor: _PayiInstrumentor) -> None:
+        try:
+            import vertexai  # type: ignore #  noqa: F401  I001
+
+            wrap_function_wrapper(
+                "vertexai.generative_models",
+                "GenerativeModel.generate_content",
+                generate_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "vertexai.preview.generative_models",
+                "GenerativeModel.generate_content",
+                generate_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "vertexai.generative_models",
+                "GenerativeModel.generate_content_async",
+                agenerate_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "vertexai.preview.generative_models",
+                "GenerativeModel.generate_content_async",
+                agenerate_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(
+        _GoogleVertexRequest(instrumentor),
+        _IsStreaming.kwargs,
+        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(
+        _GoogleVertexRequest(instrumentor),
+        _IsStreaming.kwargs,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+def count_chars_skip_spaces(text: str) -> int:
+    return sum(1 for c in text if not c.isspace())
+
+class _GoogleVertexRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor=instrumentor, category=PayiCategories.google_vertex)
+        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 vertexai.generative_models import Content, Image, Part # type: ignore #  noqa: F401  I001
+
+        if not args:
+            return True
+        
+        value: Union[ # type: ignore
+            Content,
+            str,
+            Image,
+            Part,
+            List[Union[str, Image, Part]],
+        ] = args[0] # 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):
+            return value.parts # type: ignore
+        if isinstance(value, (str, Image, Part)):
+            items = [value] # type: ignore
+
+        elif isinstance(value, list):
+            items = value # type: ignore
+
+        for item in items: # type: ignore
+            text = ""
+            if isinstance(item, Part):
+                d = item.to_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 vertexai.generative_models import Content, Image, Part, Tool # type: ignore #  noqa: F401  I001
+
+        key = "contents"
+
+        if not args:
+            return
+        
+        value: Union[ # type: ignore
+            Content,
+            str,
+            Image,
+            Part,
+            List[Union[str, Image, Part]],
+        ] = args[0] # 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(value)]).to_dict() # type: ignore
+        elif isinstance(value, (Image, Part)):
+            prompt[key] = Content(parts=[value]).to_dict() # type: ignore
+        elif isinstance(value, Content):
+            prompt[key] = value.to_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(item)) # type: ignore
+                elif isinstance(item, Part):
+                    parts.append(item) # type: ignore
+                elif isinstance(item, Image):
+                    parts.append(Part.from_image(item)) # type: ignore
+                        
+            prompt[key] = Content(parts=parts).to_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.to_dict())  # type: ignore
+            if t:
+                prompt["tools"] = t
+
+        tool_config = kwargs.get("tool_config", None)  # type: ignore
+        if tool_config:
+            # tool_config does not have to_dict or any other serializable object
+            prompt["tool_config"] = str(tool_config)  # type: ignore
+
+    @override
+    def process_chunk(self, chunk: Any) -> bool:
+        response_dict: dict[str, Any] = chunk.to_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_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.0a70 → payi-0.1.0a72}/src/payi/lib/helpers.py

@@ -22,6 +22,7 @@ class PayiCategories:
     openai:str = "system.openai"
     azure_openai:str = "system.azureopenai"
     aws_bedrock:str = "system.aws.bedrock"
+    google_vertex:str = "system.google.vertex"
 
 def create_limit_header_from_ids(limit_ids: List[str]) -> Dict[str, str]:
     if not isinstance(limit_ids, list):  # type: ignore
@@ -40,7 +41,6 @@ def create_request_header_from_tags(request_tags: List[str]) -> Dict[str, str]:
 
     return { PayiHeaderNames.request_tags: ",".join(valid_tags) } if valid_tags else {}
 
-
 def create_headers(
     limit_ids: Union[List[str], None] = None,
     request_tags: Union[List[str], None] = None,