payi 0.1.0a35__py3-none-any.whl → 0.1.0a36__py3-none-any.whl

payi/_models.py

@@ -179,14 +179,14 @@ class BaseModel(pydantic.BaseModel):
     @classmethod
     @override
     def construct(  # pyright: ignore[reportIncompatibleMethodOverride]
-        cls: Type[ModelT],
+        __cls: Type[ModelT],
         _fields_set: set[str] | None = None,
         **values: object,
     ) -> ModelT:
-        m = cls.__new__(cls)
+        m = __cls.__new__(__cls)
         fields_values: dict[str, object] = {}
 
-        config = get_model_config(cls)
+        config = get_model_config(__cls)
         populate_by_name = (
             config.allow_population_by_field_name
             if isinstance(config, _ConfigProtocol)
@@ -196,7 +196,7 @@ class BaseModel(pydantic.BaseModel):
         if _fields_set is None:
             _fields_set = set()
 
-        model_fields = get_model_fields(cls)
+        model_fields = get_model_fields(__cls)
         for name, field in model_fields.items():
             key = field.alias
             if key is None or (key not in values and populate_by_name):

payi/_response.py

@@ -136,6 +136,8 @@ class BaseAPIResponse(Generic[R]):
         if cast_to and is_annotated_type(cast_to):
             cast_to = extract_type_arg(cast_to, 0)
 
+        origin = get_origin(cast_to) or cast_to
+
         if self._is_sse_stream:
             if to:
                 if not is_stream_class_type(to):
@@ -195,8 +197,6 @@ class BaseAPIResponse(Generic[R]):
         if cast_to == bool:
             return cast(R, response.text.lower() == "true")
 
-        origin = get_origin(cast_to) or cast_to
-
         if origin == APIResponse:
             raise RuntimeError("Unexpected state - cast_to is `APIResponse`")
 
@@ -210,7 +210,13 @@ class BaseAPIResponse(Generic[R]):
                 raise ValueError(f"Subclasses of httpx.Response cannot be passed to `cast_to`")
             return cast(R, response)
 
-        if inspect.isclass(origin) and not issubclass(origin, BaseModel) and issubclass(origin, pydantic.BaseModel):
+        if (
+            inspect.isclass(
+                origin  # pyright: ignore[reportUnknownArgumentType]
+            )
+            and not issubclass(origin, BaseModel)
+            and issubclass(origin, pydantic.BaseModel)
+        ):
             raise TypeError("Pydantic models must subclass our base model type, e.g. `from payi import BaseModel`")
 
         if (

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

payi/lib/AnthropicInstrumentor.py

@@ -0,0 +1,97 @@
+import logging
+from typing import Any
+
+from wrapt import wrap_function_wrapper  # type: ignore
+
+from payi.types import IngestUnitsParams
+from payi.types.ingest_units_params import Units
+
+from .instrument import PayiInstrumentor
+
+
+class AnthropicIntrumentor:
+    @staticmethod
+    def instrument(instrumentor: PayiInstrumentor) -> None:
+        try:
+            import anthropic  # type: ignore #  noqa: F401  I001
+
+            # wrap_function_wrapper(
+            #     "anthropic.resources.completions",
+            #     "Completions.create",
+            #     chat_wrapper(instrumentor),
+            # )
+
+            wrap_function_wrapper(
+                "anthropic.resources.messages",
+                "Messages.create",
+                chat_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "anthropic.resources.messages",
+                "Messages.stream",
+                chat_wrapper(instrumentor),
+            )
+
+        except Exception as e:
+            logging.debug(f"Error instrumenting anthropic: {e}")
+            return
+
+
+@PayiInstrumentor.payi_wrapper
+def chat_wrapper(
+    instrumentor: PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    args: Any,
+    kwargs: Any,
+) -> Any:
+    return instrumentor.chat_wrapper(
+        "system.anthropic",
+        process_chunk,
+        process_synchronous_response,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+
+def process_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
+    if chunk.type == "message_start":
+        usage = chunk.message.usage
+        units = ingest["units"]
+
+        units["text"] = Units(input=usage.input_tokens, output=0)
+
+        if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
+            text_cache_write = usage.cache_creation_input_tokens
+            units["text_cache_write"] = Units(input=text_cache_write, output=0)
+
+        if hasattr(usage, "cache_read_input_tokens") and usage.cache_read_input_tokens > 0:
+            text_cache_read = usage.cache_read_input_tokens
+            units["text_cache_read"] = Units(input=text_cache_read, output=0)
+
+    elif chunk.type == "message_delta":
+        usage = chunk.usage
+        ingest["units"]["text"]["output"] = usage.output_tokens
+
+
+def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+    usage = response.usage
+    input = usage.input_tokens
+    ouptut = usage.output_tokens
+    units: dict[str, Units] = ingest["units"]
+
+    if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
+        text_cache_write = usage.cache_creation_input_tokens
+        units["text_cache_write"] = Units(input=text_cache_write, output=0)
+
+    if hasattr(usage, "cache_read_input_tokens") and usage.cache_read_input_tokens > 0:
+        text_cache_read = usage.cache_read_input_tokens
+        units["text_cache_read"] = Units(input=text_cache_read, output=0)
+
+    units["text"] = Units(input=input, output=ouptut)
+
+    if log_prompt_and_response:
+        ingest["provider_response_json"] = response.to_json()

payi/lib/Instruments.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class Instruments(Enum):
+    ALL = "all"
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"

payi/lib/OpenAIInstrumentor.py

@@ -0,0 +1,89 @@
+import json
+import logging
+from typing import Any
+from importlib.metadata import version
+
+from wrapt import wrap_function_wrapper  # type: ignore
+
+from payi.types import IngestUnitsParams
+from payi.types.ingest_units_params import Units
+
+from .instrument import PayiInstrumentor
+
+
+class OpenAiInstrumentor:
+    @staticmethod
+    def instrument(instrumentor: PayiInstrumentor) -> None:
+        try:
+            from openai import OpenAI  # type: ignore #  noqa: F401  I001
+
+            wrap_function_wrapper(
+                "openai.resources.chat.completions",
+                "Completions.create",
+                chat_wrapper(instrumentor),
+            )
+        except Exception as e:
+            logging.debug(f"Error instrumenting openai: {e}")
+            return
+
+
+@PayiInstrumentor.payi_wrapper
+def chat_wrapper(
+    instrumentor: PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    args: Any,
+    kwargs: Any,
+) -> Any:
+    return instrumentor.chat_wrapper(
+        "system.openai",
+        process_chat_chunk,
+        process_chat_synchronous_response,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+
+def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+    response_dict = model_to_dict(response)
+
+    add_usage_units(response_dict["usage"], ingest["units"])
+
+    if log_prompt_and_response:
+        ingest["provider_response_json"] = [json.dumps(response_dict)]
+
+
+def process_chat_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
+    model = model_to_dict(chunk)
+    usage = model.get("usage")
+    if usage:
+        add_usage_units(usage, ingest["units"])
+
+
+def model_to_dict(model: Any) -> Any:
+    if version("pydantic") < "2.0.0":
+        return model.dict()
+    if hasattr(model, "model_dump"):
+        return model.model_dump()
+    elif hasattr(model, "parse"):  # Raw API response
+        return model_to_dict(model.parse())
+    else:
+        return model
+
+
+def add_usage_units(usage: "dict[str, Any]", units: "dict[str, Units]") -> None:
+    input = usage["prompt_tokens"] if "prompt_tokens" in usage else 0
+    output = usage["completion_tokens"] if "completion_tokens" in usage else 0
+    input_cache = 0
+
+    prompt_tokens_details = usage.get("prompt_tokens_details")
+    if prompt_tokens_details:
+        input_cache = prompt_tokens_details.get("cached_tokens", 0)
+        if input_cache != 0:
+            units["text_cache_read"] = Units(input=input_cache, output=0)
+
+    input -= input_cache
+
+    units["text"] = Units(input=input, output=output)

payi/lib/Stopwatch.py

@@ -0,0 +1,27 @@
+import time
+from typing import Optional
+
+
+class Stopwatch:
+    def __init__(self) -> None:
+        self.start_time: Optional[float] = None
+        self.end_time: Optional[float] = None
+
+    def start(self) -> None:
+        self.start_time = time.perf_counter()
+
+    def stop(self) -> None:
+        self.end_time = time.perf_counter()
+
+    def elapsed_s(self) -> float:
+        if self.start_time is None:
+            raise ValueError("Stopwatch has not been started")
+        if self.end_time is None:
+            return time.perf_counter() - self.start_time
+        return self.end_time - self.start_time
+
+    def elapsed_ms(self) -> float:
+        return self.elapsed_s() * 1000
+
+    def elapsed_ms_int(self) -> int:
+        return int(self.elapsed_ms())

payi/lib/helpers.py

@@ -3,28 +3,30 @@ from typing import Dict, List, Union
 
 
 def create_limit_header_from_ids(limit_ids: List[str]) -> Dict[str, str]:
-    if not isinstance(limit_ids, list): # type: ignore
+    if not isinstance(limit_ids, list):  # type: ignore
         raise TypeError("limit_ids must be a list")
 
-    valid_ids = [id.strip() for id in limit_ids if isinstance(id, str) and id.strip()] # type: ignore
+    valid_ids = [id.strip() for id in limit_ids if isinstance(id, str) and id.strip()]  # type: ignore
 
     return {"xProxy-Limit-IDs": ",".join(valid_ids)} if valid_ids else {}
 
+
 def create_request_header_from_tags(request_tags: List[str]) -> Dict[str, str]:
-    if not isinstance(request_tags, list): # type: ignore
+    if not isinstance(request_tags, list):  # type: ignore
         raise TypeError("request_tags must be a list")
 
-    valid_tags = [tag.strip() for tag in request_tags if isinstance(tag, str) and tag.strip()] # type: ignore
-    
+    valid_tags = [tag.strip() for tag in request_tags if isinstance(tag, str) and tag.strip()]  # type: ignore
+
     return {"xProxy-Request-Tags": ",".join(valid_tags)} if valid_tags else {}
 
+
 def create_headers(
-    limit_ids: Union[List[str], None] = None, 
+    limit_ids: Union[List[str], None] = None,
     request_tags: Union[List[str], None] = None,
     user_id: Union[str, None] = None,
-    experience_id: Union[str, None] = None, 
-    experience_name: Union[str, None] = None, 
-    ) -> Dict[str, str]:
+    experience_id: Union[str, None] = None,
+    experience_name: Union[str, None] = None,
+) -> Dict[str, str]:
     headers: Dict[str, str] = {}
 
     if limit_ids:
@@ -38,4 +40,4 @@ def create_headers(
     if experience_name:
         headers.update({"xProxy-Experience-Name": experience_name})
 
-    return headers
+    return headers