payi 0.1.0a40__py3-none-any.whl → 0.1.0a42__py3-none-any.whl

payi/_constants.py

@@ -6,7 +6,7 @@ RAW_RESPONSE_HEADER = "X-Stainless-Raw-Response"
 OVERRIDE_CAST_TO_HEADER = "____stainless_override_cast_to"
 
 # default timeout is 1 minute
-DEFAULT_TIMEOUT = httpx.Timeout(timeout=60.0, connect=5.0)
+DEFAULT_TIMEOUT = httpx.Timeout(timeout=60, connect=5.0)
 DEFAULT_MAX_RETRIES = 2
 DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20)
 

payi/_models.py

@@ -172,7 +172,7 @@ class BaseModel(pydantic.BaseModel):
     @override
     def __str__(self) -> str:
         # mypy complains about an invalid self arg
-        return f'{self.__repr_name__()}({self.__repr_str__(", ")})'  # type: ignore[misc]
+        return f"{self.__repr_name__()}({self.__repr_str__(', ')})"  # type: ignore[misc]
 
     # Override the 'construct' method in a way that supports recursive parsing without validation.
     # Based on https://github.com/samuelcolvin/pydantic/issues/1168#issuecomment-817742836.

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

payi/lib/AnthropicInstrumentor.py

@@ -7,7 +7,7 @@ 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
+from .instrument import IsStreaming, PayiInstrumentor
 
 
 class AnthropicIntrumentor:
@@ -16,12 +16,6 @@ class AnthropicIntrumentor:
         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",
@@ -34,6 +28,18 @@ class AnthropicIntrumentor:
                 chat_wrapper(instrumentor),
             )
 
+            wrap_function_wrapper(
+                "anthropic.resources.messages",
+                "AsyncMessages.create",
+                achat_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "anthropic.resources.messages",
+                "AsyncMessages.stream",
+                achat_wrapper(instrumentor),
+            )
+
         except Exception as e:
             logging.debug(f"Error instrumenting anthropic: {e}")
             return
@@ -44,14 +50,35 @@ def chat_wrapper(
     instrumentor: PayiInstrumentor,
     wrapped: Any,
     instance: Any,
-    args: Any,
-    kwargs: Any,
+    *args: Any,
+    **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
         "system.anthropic",
         process_chunk,
         process_request,
         process_synchronous_response,
+        IsStreaming.kwargs,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+@PayiInstrumentor.payi_awrapper
+async def achat_wrapper(
+    instrumentor: PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return await instrumentor.achat_wrapper(
+        "system.anthropic",
+        process_chunk,
+        process_request,
+        process_synchronous_response,
+        IsStreaming.kwargs,
         wrapped,
         instance,
         args,
@@ -81,10 +108,10 @@ def process_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
         ingest["units"]["text"]["output"] = usage.output_tokens
 
 
-def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_prompt_and_response: bool, *args: Any, **kwargs: 'dict[str, Any]') -> Any: # noqa: ARG001
     usage = response.usage
     input = usage.input_tokens
-    ouptut = usage.output_tokens
+    output = usage.output_tokens
     units: dict[str, Units] = ingest["units"]
 
     if hasattr(usage, "cache_creation_input_tokens") and usage.cache_creation_input_tokens > 0:
@@ -97,11 +124,13 @@ def process_synchronous_response(response: Any, ingest: IngestUnitsParams, log_p
 
     input = PayiInstrumentor.update_for_vision(input, units)
 
-    units["text"] = Units(input=input, output=ouptut)
+    units["text"] = Units(input=input, output=output)
 
     if log_prompt_and_response:
         ingest["provider_response_json"] = response.to_json()
 
+    return None
+
 def has_image_and_get_texts(encoding: tiktoken.Encoding, content: Union[str, 'list[Any]']) -> 'tuple[bool, int]':
     if isinstance(content, str):
         return False, 0
@@ -113,7 +142,7 @@ 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
 
-def process_request(ingest: IngestUnitsParams, kwargs: Any) -> None:
+def process_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: # noqa: ARG001
     messages = kwargs.get("messages")
     if not messages or len(messages) == 0:
         return

payi/lib/BedrockInstrumentor.py

@@ -0,0 +1,282 @@
+import json
+import logging
+from typing import Any
+from functools import wraps
+
+from wrapt import ObjectProxy, wrap_function_wrapper  # type: ignore
+
+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, PayiInstrumentor
+
+
+class BedrockInstrumentor:
+    @staticmethod
+    def instrument(instrumentor: PayiInstrumentor) -> None:
+        try:
+            import boto3  # type: ignore #  noqa: F401  I001
+
+            wrap_function_wrapper(
+                "botocore.client",
+                "ClientCreator.create_client",
+                create_client_wrapper(instrumentor),
+            )
+
+            wrap_function_wrapper(
+                "botocore.session",
+                "Session.create_client",
+                create_client_wrapper(instrumentor),
+            )
+
+        except Exception as e:
+            logging.debug(f"Error instrumenting bedrock: {e}")
+            return
+
+@PayiInstrumentor.payi_wrapper
+def create_client_wrapper(instrumentor: PayiInstrumentor, wrapped: Any, instance: Any, *args: Any, **kwargs: Any) -> Any: #  noqa: ARG001
+    if kwargs.get("service_name") != "bedrock-runtime":
+        return wrapped(*args, **kwargs)
+
+    try:
+        client: Any = wrapped(*args, **kwargs)
+        client.invoke_model = wrap_invoke(instrumentor, client.invoke_model)
+        client.invoke_model_with_response_stream = wrap_invoke_stream(instrumentor, client.invoke_model_with_response_stream)
+        client.converse = wrap_converse(instrumentor, client.converse)
+        client.converse_stream = wrap_converse_stream(instrumentor, client.converse_stream)
+
+        return client
+    except Exception as e:
+        logging.debug(f"Error instrumenting bedrock client: {e}")
+    
+    return wrapped(*args, **kwargs)
+    
+class InvokeResponseWrapper(ObjectProxy): # type: ignore
+    def __init__(
+        self,
+        response: Any,
+        instrumentor: PayiInstrumentor,
+        ingest: IngestUnitsParams,
+        log_prompt_and_response: bool
+        ) -> None:
+
+        super().__init__(response) # type: ignore
+        self._response = response
+        self._instrumentor = instrumentor
+        self._ingest = ingest
+        self._log_prompt_and_response = log_prompt_and_response
+
+    def read(self, amt: Any =None): # type: ignore
+        # data is array of bytes
+        data: Any = self.__wrapped__.read(amt) # type: ignore
+        response = json.loads(data)
+
+        resource = self._ingest["resource"]
+        if not resource:
+            return
+        
+        input: int = 0
+        output: int = 0
+        units: dict[str, Units] = self._ingest["units"]
+
+        if resource.startswith("meta.llama3"):
+            input = response['prompt_token_count']
+            output = response['generation_token_count']
+        elif resource.startswith("anthropic."):
+            usage = response['usage']
+            input = usage['input_tokens']
+            output = usage['output_tokens']
+        units["text"] = Units(input=input, output=output)
+
+        if self._log_prompt_and_response:
+            self._ingest["provider_response_json"] = data.decode('utf-8')
+            
+        self._instrumentor._ingest_units(self._ingest)
+
+        return data
+
+def wrap_invoke(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
+    @wraps(wrapped)
+    def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
+        modelId:str = kwargs.get("modelId", "") # type: ignore
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                "system.aws.bedrock",
+                None,
+                process_invoke_request,  
+                process_synchronous_invoke_response,  
+                IsStreaming.false,
+                wrapped,
+                None,
+                args,
+                kwargs,
+        )
+        return wrapped(*args, **kwargs)
+    
+    return invoke_wrapper
+
+def wrap_invoke_stream(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
+    @wraps(wrapped)
+    def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
+        modelId:str = kwargs.get("modelId", "") # type: ignore
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                "system.aws.bedrock",
+                process_invoke_streaming_anthropic_chunk if modelId.startswith("anthropic.") else process_invoke_streaming_llama_chunk,
+                process_invoke_request, 
+                None,  
+                IsStreaming.true,
+                wrapped,
+                None,
+                args,
+                kwargs,
+            )
+        return wrapped(*args, **kwargs)
+
+    return invoke_wrapper
+
+def wrap_converse(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
+    @wraps(wrapped)
+    def invoke_wrapper(*args: Any, **kwargs: 'dict[str, Any]') -> Any:
+        modelId:str = kwargs.get("modelId", "") # type: ignore
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                "system.aws.bedrock",
+                None,
+                process_converse_request,  
+                process_synchronous_converse_response,  
+                IsStreaming.false,
+                wrapped,
+                None,
+                args,
+                kwargs,
+        )
+        return wrapped(*args, **kwargs)
+    
+    return invoke_wrapper
+
+def wrap_converse_stream(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
+    @wraps(wrapped)
+    def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
+        modelId:str = kwargs.get("modelId", "") # type: ignore
+
+        if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
+            return instrumentor.chat_wrapper(
+                "system.aws.bedrock",
+                process_converse_streaming_chunk,
+                process_converse_request, 
+                None,  
+                IsStreaming.true,
+                wrapped,
+                None,
+                args,
+                kwargs,
+            )
+        return wrapped(*args, **kwargs)
+
+    return invoke_wrapper
+
+def process_invoke_streaming_anthropic_chunk(chunk: str, ingest: IngestUnitsParams) -> None:
+    chunk_dict =  json.loads(chunk)
+    type = chunk_dict.get("type", "")
+
+    if type == "message_start":
+        usage = chunk_dict['message']['usage']
+        units = ingest["units"]
+
+        input = PayiInstrumentor.update_for_vision(usage['input_tokens'], units)
+
+        units["text"] = Units(input=input, output=0)
+
+        text_cache_write: int = usage.get("cache_creation_input_tokens", 0)
+        if text_cache_write > 0:
+            units["text_cache_write"] = Units(input=text_cache_write, output=0)
+
+        text_cache_read: int = usage.get("cache_read_input_tokens", 0)
+        if text_cache_read > 0:
+            units["text_cache_read"] = Units(input=text_cache_read, output=0)
+
+    elif type == "message_delta":
+        usage = chunk_dict['usage']
+        ingest["units"]["text"]["output"] = usage['output_tokens']
+
+def process_invoke_streaming_llama_chunk(chunk: str, ingest: IngestUnitsParams) -> None:
+    chunk_dict =  json.loads(chunk)
+    metrics = chunk_dict.get("amazon-bedrock-invocationMetrics", {})
+    if metrics:
+        input = metrics.get("inputTokenCount", 0)
+        output = metrics.get("outputTokenCount", 0)
+        ingest["units"]["text"] = Units(input=input, output=output)
+    
+def process_synchronous_invoke_response(
+        response: Any,
+        ingest: IngestUnitsParams,
+        log_prompt_and_response: bool,
+        instrumentor: PayiInstrumentor,
+        **kargs: Any) -> Any: #  noqa: ARG001
+
+    metadata = response.get("ResponseMetadata", {})
+
+    # request_id = metadata.get("RequestId", "")
+    # if request_id:
+    #     ingest["provider_request_id"] = request_id
+
+    response_headers = metadata.get("HTTPHeaders", {}).copy()
+    if response_headers:
+        ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+
+    response["body"] = InvokeResponseWrapper(
+        response=response["body"],
+        instrumentor=instrumentor,
+        ingest=ingest,
+        log_prompt_and_response=log_prompt_and_response)
+
+    return response
+
+def process_invoke_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: #  noqa: ARG001
+    return
+
+def process_converse_streaming_chunk(chunk: 'dict[str, Any]', ingest: IngestUnitsParams) -> None:
+    metadata = chunk.get("metadata", {})
+
+    if metadata:
+        usage = metadata['usage']
+        input = usage["inputTokens"]
+        output = usage["outputTokens"]
+        ingest["units"]["text"] = Units(input=input, output=output)
+
+def process_synchronous_converse_response(
+        response: 'dict[str, Any]',
+        ingest: IngestUnitsParams,
+        log_prompt_and_response: bool,
+        **kargs: Any) -> Any: #  noqa: ARG001
+
+    usage = response["usage"]
+    input = usage["inputTokens"]
+    output = usage["outputTokens"]
+    
+    units: dict[str, Units] = ingest["units"]
+    units["text"] = Units(input=input, output=output)
+
+    metadata = response.get("ResponseMetadata", {})
+
+    # request_id = metadata.get("RequestId", "")
+    # if request_id:
+    #     ingest["provider_request_id"] = request_id
+
+    response_headers = metadata.get("HTTPHeaders", {})
+    if response_headers:
+        ingest["provider_response_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in response_headers.items()]
+
+    if log_prompt_and_response:
+        response_without_metadata = response.copy()
+        response_without_metadata.pop("ResponseMetadata", None)
+        ingest["provider_response_json"] = json.dumps(response_without_metadata)
+
+    return None    
+
+def process_converse_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: #  noqa: ARG001
+    return

payi/lib/Instruments.py

@@ -5,3 +5,4 @@ class Instruments(Enum):
     ALL = "all"
     OPENAI = "openai"
     ANTHROPIC = "anthropic"
+    AWS_BEDROCK = "aws.bedrock"

payi/lib/OpenAIInstrumentor.py

@@ -3,13 +3,13 @@ import logging
 from typing import Any, Union
 from importlib.metadata import version
 
-import tiktoken
+import tiktoken  # type: ignore
 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
+from .instrument import IsStreaming, PayiInstrumentor
 
 
 class OpenAiInstrumentor:
@@ -23,6 +23,13 @@ class OpenAiInstrumentor:
                 "Completions.create",
                 chat_wrapper(instrumentor),
             )
+
+            wrap_function_wrapper(
+                "openai.resources.chat.completions",
+                "AsyncCompletions.create",
+                achat_wrapper(instrumentor),
+            )
+
         except Exception as e:
             logging.debug(f"Error instrumenting openai: {e}")
             return
@@ -33,14 +40,35 @@ def chat_wrapper(
     instrumentor: PayiInstrumentor,
     wrapped: Any,
     instance: Any,
-    args: Any,
-    kwargs: Any,
+    *args: Any,
+    **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
         "system.openai",
         process_chat_chunk,
         process_request,
         process_chat_synchronous_response,
+        IsStreaming.kwargs,
+        wrapped,
+        instance,
+        args,
+        kwargs,
+    )
+
+@PayiInstrumentor.payi_awrapper
+async def achat_wrapper(
+    instrumentor: PayiInstrumentor,
+    wrapped: Any,
+    instance: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return await instrumentor.achat_wrapper(
+        "system.openai",
+        process_chat_chunk,
+        process_request,
+        process_chat_synchronous_response,
+        IsStreaming.kwargs,
         wrapped,
         instance,
         args,
@@ -48,7 +76,7 @@ def chat_wrapper(
     )
 
 
-def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool) -> None:
+def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool, **kwargs: Any) -> Any: #  noqa: ARG001
     response_dict = model_to_dict(response)
 
     add_usage_units(response_dict["usage"], ingest["units"])
@@ -56,6 +84,7 @@ def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams,
     if log_prompt_and_response:
         ingest["provider_response_json"] = [json.dumps(response_dict)]
 
+    return None
 
 def process_chat_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
     model = model_to_dict(chunk)
@@ -101,7 +130,7 @@ 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
 
-def process_request(ingest: IngestUnitsParams, kwargs: Any) -> None:
+def process_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: # noqa: ARG001
     messages = kwargs.get("messages")
     if not messages or len(messages) == 0:
         return
@@ -110,9 +139,9 @@ def process_request(ingest: IngestUnitsParams, kwargs: Any) -> None:
     has_image = False
 
     try: 
-        enc = tiktoken.encoding_for_model(kwargs.get("model"))
+        enc = tiktoken.encoding_for_model(kwargs.get("model")) # type: ignore
     except KeyError:
-        enc = tiktoken.get_encoding("o200k_base")
+        enc = tiktoken.get_encoding("o200k_base") # type: ignore
     
     for message in messages:
         msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))

payi/lib/Stopwatch.py

@@ -15,7 +15,7 @@ class Stopwatch:
 
     def elapsed_s(self) -> float:
         if self.start_time is None:
-            raise ValueError("Stopwatch has not been started")
+            return 0.0 # 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