payi 0.1.0a96__py3-none-any.whl → 0.1.0a97__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.96"  # x-release-please-version
+__version__ = "0.1.0-alpha.97"  # x-release-please-version

payi/lib/VertexRequest.py

@@ -10,6 +10,8 @@ from .instrument import _ChunkResult, _StreamingType, _ProviderRequest, _PayiIns
 
 
 class _VertexRequest(_ProviderRequest): # type: ignore
+    KNOWN_MODALITIES = ("VIDEO", "AUDIO", "TEXT", "VISION", "IMAGE")
+
     def __init__(self, instrumentor: _PayiInstrumentor):
         super().__init__(
             instrumentor=instrumentor,
@@ -154,6 +156,7 @@ class _VertexRequest(_ProviderRequest): # type: ignore
 
         prompt_tokens_details: list[dict[str, Any]] = usage.get("prompt_tokens_details", [])
         candidates_tokens_details: list[dict[str, Any]] = usage.get("candidates_tokens_details", [])
+        cache_tokens_details: list[dict[str, Any]] = usage.get("cache_tokens_details", [])
 
         if not model:
             model = ""
@@ -214,27 +217,50 @@ class _VertexRequest(_ProviderRequest): # type: ignore
             if is_large_context_token_model(model, input):
                 large_context = "_large_context"
 
+            cache_details: dict[str, int] = {}
+
+            for details in cache_tokens_details:
+                modality = details.get("modality", "")
+                if not modality:
+                    continue
+
+                modality_token_count = details.get("token_count", 0)
+                
+                if modality == "IMAGE":
+                    modality = "VISION"
+
+                if modality in _VertexRequest.KNOWN_MODALITIES:
+                    cache_details[modality] = modality_token_count
+                    add_units(self, modality.lower() + "_cache_read" + large_context, input=modality_token_count)
+
             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":
-                    add_units(self, "vision"+large_context, input=modality_token_count)
-                elif modality in ("VIDEO", "AUDIO", "TEXT"):
-                    add_units(self, modality.lower()+large_context, input=modality_token_count)
+                    modality = "VISION"
+
+                if modality in _VertexRequest.KNOWN_MODALITIES:
+                    # Subtract cache_details value if modality is present, floor at zero
+                    if modality in cache_details:
+                        modality_token_count = max(0, modality_token_count - cache_details[modality])
+
+                    add_units(self, modality.lower() + large_context, 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"):
-                    add_units(self, modality.lower()+large_context, output=modality_token_count)
+                if modality in _VertexRequest.KNOWN_MODALITIES:
+                    add_units(self, modality.lower() + large_context, output=modality_token_count)
 
             if thinking_token_count > 0:
-                add_units(self, "reasoning"+large_context, output=thinking_token_count)
+                add_units(self, "reasoning" + large_context, output=thinking_token_count)
 
         if not self._ingest["units"]:
             input = usage.get("prompt_token_count", 0)

payi/lib/instrument.py

@@ -15,11 +15,13 @@ from dataclasses import dataclass
 import nest_asyncio  # type: ignore
 from wrapt import ObjectProxy  # type: ignore
 
-from payi import Payi, AsyncPayi, APIConnectionError, __version__ as _payi_version
+from payi import Payi, AsyncPayi, APIStatusError, APIConnectionError, __version__ as _payi_version
 from payi.types import IngestUnitsParams
 from payi.lib.helpers import PayiHeaderNames
+from payi.types.shared import XproxyResult
 from payi.types.ingest_response import IngestResponse
 from payi.types.ingest_units_params import Units, ProviderResponseFunctionCall
+from payi.types.shared.xproxy_error import XproxyError
 from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
 from .helpers import PayiCategories
@@ -146,6 +148,19 @@ class PayiInstrumentConfig(TypedDict, total=False):
     user_id: Optional[str]
     request_tags: Optional["list[str]"]
 
+class PayiContext(TypedDict, total=False):
+    use_case_name: Optional[str]
+    use_case_id: Optional[str]
+    use_case_version: Optional[int]
+    use_case_step: Optional[str]
+    limit_ids: Optional['list[str]']
+    user_id: Optional[str]
+    request_tags: Optional["list[str]"]
+    price_as_category: Optional[str]
+    price_as_resource: Optional[str]
+    resource_scope: Optional[str]
+    last_result: Optional[Union[XproxyResult, XproxyError]]
+
 class _Context(TypedDict, total=False):
     proxy: Optional[bool]
     use_case_name: Optional[str]
@@ -169,7 +184,7 @@ class _StreamingType(Enum):
     iterator = 1
     stream_manager = 2
 
-class _TrackContext:
+class _InternalTrackContext:
     def __init__(
         self,
         context: _Context,
@@ -235,6 +250,8 @@ class _PayiInstrumentor:
 
         self._instrument_inline_data: bool = global_config.get("instrument_inline_data", False)
 
+        self._last_result: Optional[Union[XproxyResult, XproxyError]] = None
+
         global_instrumentation = global_config.pop("global_instrumentation", True)
 
         if instruments is None or "*" in instruments:
@@ -354,7 +371,7 @@ class _PayiInstrumentor:
 
         return log_ingest_units
         
-    def _process_ingest_units(self, request: _ProviderRequest, log_data: 'dict[str, str]') -> bool:
+    def _process_ingest_units(self, request: _ProviderRequest, log_data: 'dict[str, str]') -> None:
         ingest_units = request._ingest
 
         if request._function_call_builder:
@@ -394,8 +411,6 @@ class _PayiInstrumentor:
             if stack_trace is not None:
                 log_data["stack_trace"] = stack_trace
 
-        return True
-
     def _process_ingest_units_response(self, ingest_response: IngestResponse) -> None:
         if ingest_response.xproxy_result.limits:
             for limit_id, state in ingest_response.xproxy_result.limits.items():
@@ -413,7 +428,7 @@ class _PayiInstrumentor:
                 if removeBlockedId:
                     self._blocked_limits.discard(limit_id)
 
-    def _process_ingest_connection_error(self, e: APIConnectionError, ingest_units: IngestUnitsParams) -> None:
+    def _process_ingest_connection_error(self, e: APIConnectionError, ingest_units: IngestUnitsParams) -> XproxyError:
         now = time.time()
 
         if (now - self._api_connection_error_last_log_time) > self._api_connection_error_window:
@@ -430,7 +445,9 @@ class _PayiInstrumentor:
             # Suppress and count
             self._api_connection_error_count += 1
 
-    async def _aingest_units(self, request: _ProviderRequest) -> Optional[IngestResponse]:
+        return XproxyError(code="api_connection_error", message=str(e))
+
+    async def _aingest_units_worker(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
         ingest_response: Optional[IngestResponse] = None
         ingest_units = request._ingest
 
@@ -438,9 +455,8 @@ class _PayiInstrumentor:
 
         # return early if there are no units to ingest and on a successul ingest request
         log_data: 'dict[str,str]' = {}
-        if not self._process_ingest_units(request, log_data):
-            self._logger.debug(f"_aingest_units: exit early")
-            return None
+
+        self._process_ingest_units(request, log_data)
 
         try:
             if self._logger.isEnabledFor(logging.DEBUG):
@@ -452,7 +468,7 @@ class _PayiInstrumentor:
                 ingest_response = self._payi.ingest.units(**ingest_units)
             else:
                 self._logger.error("No payi instance to ingest units")
-                return None         
+                return XproxyError(code="configuration_error", message="No Payi or AsyncPayi instance configured for ingesting units")
 
             self._logger.debug(f"_aingest_units: success ({ingest_response})")
 
@@ -463,15 +479,21 @@ class _PayiInstrumentor:
                     request_id = ingest_response.xproxy_result.request_id
                     self._prompt_and_response_logger(request_id, log_data)  # type: ignore
 
-            return ingest_response
-        except Exception as e:
-            if isinstance(e, APIConnectionError) and self._api_connection_error_window > 0:
-                self._process_ingest_connection_error(e, ingest_units)
-            else:
-                self._logger.error(f"Error Pay-i async ingesting: exception {e}, request {ingest_units}")
-    
-        return None         
+            return ingest_response.xproxy_result
+
+        except APIConnectionError as api_ex:
+            return self._process_ingest_connection_error(api_ex, ingest_units)
+
+        except APIStatusError as api_status_ex:
+            return self._process_api_status_error(api_status_ex)      
+
+        except Exception as ex:
+            self._logger.error(f"Error Pay-i async ingesting: exception {ex}, request {ingest_units}")
+            return XproxyError(code="unknown_error", message=str(ex))
     
+    async def _aingest_units(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
+        return self.set_xproxy_result(await self._aingest_units_worker(request))
+
     def _call_async_use_case_definition_create(self, use_case_name: str, use_case_description: str) -> None:
         if not self._apayi:
             return
@@ -491,7 +513,7 @@ class _PayiInstrumentor:
         except Exception as e:
             self._logger.error(f"Error calling async use_cases.definitions.create synchronously: {e}")
 
-    def _call_aingest_sync(self, request: _ProviderRequest) -> Optional[IngestResponse]:
+    def _call_aingest_sync(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
         try:
             loop = asyncio.get_running_loop()
         except RuntimeError:
@@ -500,7 +522,7 @@ class _PayiInstrumentor:
         try:
             if loop and loop.is_running():
                 nest_asyncio.apply(loop) # type: ignore
-                return asyncio.run(self._aingest_units(request))                
+                return asyncio.run(self._aingest_units(request))
             else:
                 # When there's no running loop, create a new one
                 return asyncio.run(self._aingest_units(request))
@@ -508,7 +530,41 @@ class _PayiInstrumentor:
             self._logger.error(f"Error calling aingest_units synchronously: {e}")
         return None
         
-    def _ingest_units(self, request: _ProviderRequest) -> Optional[IngestResponse]:
+    def _process_api_status_error(self, e: APIStatusError) -> Optional[XproxyError]:
+        try:
+            body_dict: dict[str, Any] = {}
+
+            # Try to get the response body as JSON
+            body = e.body
+            if body is None:
+                self._logger.error("APIStatusError response has no body attribute")
+                return XproxyError(code="unknown_error", message=str(e))
+
+            # If body is bytes, decode to string
+            if isinstance(body, bytes):
+                body = body.decode("utf-8")
+            if isinstance(body, dict):
+                body_dict = body # type: ignore
+            else:
+                body = str(body)
+
+            if not body_dict:
+                try:
+                    body_dict = json.loads(body)  # type: ignore
+                except Exception as json_ex:
+                    self._logger.error(f"Failed to parse response body as JSON: {json_ex}")
+                    return XproxyError(code="invalid_json", message=str(e))
+
+            xproxy_error = body_dict.get("xproxy_error", {})
+            code = xproxy_error.get("code", "unknown_error")
+            message = xproxy_error.get("message", str(e))
+            return XproxyError(code=code, message=message)
+
+        except Exception as ex:
+            self._logger.error(f"Exception in _process_api_status_error: {ex}")
+            return XproxyError(code="exception", message=str(ex))
+
+    def _ingest_units_worker(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
         ingest_response: Optional[IngestResponse] = None
         ingest_units = request._ingest
 
@@ -516,9 +572,8 @@ class _PayiInstrumentor:
 
         # return early if there are no units to ingest and on a successul ingest request
         log_data: 'dict[str,str]' = {}
-        if not self._process_ingest_units(request, log_data):
-            self._logger.debug(f"_ingest_units: exit early")
-            return None
+
+        self._process_ingest_units(request, log_data)
 
         try:
             if self._payi:
@@ -534,22 +589,28 @@ class _PayiInstrumentor:
                     request_id = ingest_response.xproxy_result.request_id
                     self._prompt_and_response_logger(request_id, log_data)  # type: ignore
 
-                return ingest_response
+                return ingest_response.xproxy_result
             elif self._apayi:
                 # task runs async. aingest_units will invoke the callback and post process
-                ingest_response = self._call_aingest_sync(request)
-                self._logger.debug(f"_ingest_units: apayi success ({ingest_response})")
-                return ingest_response
+                sync_response = self._call_aingest_sync(request)
+                self._logger.debug(f"_ingest_units: apayi success ({sync_response})")
+                return sync_response
             else:
                 self._logger.error("No payi instance to ingest units")
+                return XproxyError(code="configuration_error", message="No Payi or AsyncPayi instance configured for ingesting units")
 
-        except Exception as e:
-            if isinstance(e, APIConnectionError) and self._api_connection_error_window > 0:
-                self._process_ingest_connection_error(e, ingest_units)
-            else:
-                self._logger.error(f"Error Pay-i ingesting: exception {e}, request {ingest_units}")
+        except APIConnectionError as api_ex:
+            return self._process_ingest_connection_error(api_ex, ingest_units)
+
+        except APIStatusError as api_status_ex:
+            return self._process_api_status_error(api_status_ex)      
+
+        except Exception as ex:
+            self._logger.error(f"Error Pay-i async ingesting: exception {ex}, request {ingest_units}")
+            return XproxyError(code="unknown_error", message=str(ex))
         
-        return None
+    def _ingest_units(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
+        return self.set_xproxy_result(self._ingest_units_worker(request))
 
     def _setup_call_func(
         self
@@ -1040,6 +1101,10 @@ class _PayiInstrumentor:
 
         return extra_headers
 
+    def set_xproxy_result(self, response: Optional[Union[XproxyResult, XproxyError]]) -> Optional[Union[XproxyResult, XproxyError]]:
+        self._last_result = response
+        return response
+
     @staticmethod
     def _update_extra_headers(
         context: _Context,
@@ -1617,7 +1682,7 @@ def track_context(
     price_as_resource: Optional[str] = None,
     resource_scope: Optional[str] = None,
     proxy: Optional[bool] = None,
-) -> _TrackContext:
+) -> _InternalTrackContext:
     # Create a new context for tracking
     context: _Context = {}
 
@@ -1637,4 +1702,22 @@ def track_context(
     context["price_as_resource"] = price_as_resource
     context["resource_scope"] = resource_scope
 
-    return _TrackContext(context)
+    return _InternalTrackContext(context)
+
+def get_context() -> PayiContext:
+    """
+    Returns the current tracking context from calls to @track and with track_context().
+    If no context is active, returns an empty context.
+    """
+    if not _instrumentor:
+        return PayiContext()
+    internal_context = _instrumentor.get_context() or {}
+
+    context_dict = {
+        key: value
+        for key, value in internal_context.items()
+        if key in PayiContext.__annotations__ and value is not None
+    }
+    if _instrumentor._last_result:
+        context_dict["last_result"] = _instrumentor._last_result
+    return PayiContext(**dict(context_dict))  # type: ignore

{payi-0.1.0a96.dist-info → payi-0.1.0a97.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: payi
-Version: 0.1.0a96
+Version: 0.1.0a97
 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.0a96.dist-info → payi-0.1.0a97.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=HrUZ97eXJDQhDfJm78BH6NvNVuDnKXoia2Jj2xHrnm0,165
+payi/_version.py,sha256=ox8PPHdQJ81G2cAApozsYaqcHPCKCgrcbVajDxPBCJU,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
@@ -31,9 +31,9 @@ payi/lib/GoogleGenAiInstrumentor.py,sha256=DLbmlwmyOwooe7FuBkEAO_Z3xjPKpvQyO4VwL
 payi/lib/OpenAIInstrumentor.py,sha256=6HENrdN6kLMUwC00wNSdeGeK4fUDBWo2dCbVC9koi0c,22581
 payi/lib/Stopwatch.py,sha256=7OJlxvr2Jyb6Zr1LYCYKczRB7rDVKkIR7gc4YoleNdE,764
 payi/lib/VertexInstrumentor.py,sha256=q-oYr3YuwPc95xYn6KhjvKs5ZweGRg6mOIkxSTz-lic,7081
-payi/lib/VertexRequest.py,sha256=NNH6S86nCHaU54Rb1rxai-b3UiJvatmzKkDKBmqMeKY,10530
+payi/lib/VertexRequest.py,sha256=dmL0J2UcYFTEX1m9yoxAbF4cMzHcFiKKKZz7-ltSC88,11521
 payi/lib/helpers.py,sha256=FPzNSSHGf9bgD6CanB7yVx_U8t4lm2c0jlZKrsziYlc,4242
-payi/lib/instrument.py,sha256=VEyEyjM0GGIdFqv9ogwBwyWQ8kRJgJWUxlWcQc6798Q,64801
+payi/lib/instrument.py,sha256=SJNh4IPZHMsHr7W9SzJvajYtDz36TR6W1srlAim4IHU,68486
 payi/resources/__init__.py,sha256=B2bn1ZfCf6TbHlzZvy5TpFPtALnFcBRPYVKQH3S5qfQ,2457
 payi/resources/ingest.py,sha256=awE7xDdKOUL2Yhcrhcc11hSNycBlj8QRlahptjana5Q,23040
 payi/resources/categories/__init__.py,sha256=WeotN_d-0Ri8ohsrNPbve7RyViD9_N0NA9DrV3WYg3w,1701
@@ -132,7 +132,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.0a96.dist-info/METADATA,sha256=VVawI7dTiqYjzxZQOhcwSagnNaBhaseMbtBybj4HUeY,16333
-payi-0.1.0a96.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-payi-0.1.0a96.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
-payi-0.1.0a96.dist-info/RECORD,,
+payi-0.1.0a97.dist-info/METADATA,sha256=FWpOf7rI_UTcstFDNyRlr6xoYwwVJ7GpWsFkJH9A0dY,16333
+payi-0.1.0a97.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+payi-0.1.0a97.dist-info/licenses/LICENSE,sha256=CQt03aM-P4a3Yg5qBg3JSLVoQS3smMyvx7tYg_6V7Gk,11334
+payi-0.1.0a97.dist-info/RECORD,,