payi 0.1.0a41__py3-none-any.whl → 0.1.0a42__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.41"  # x-release-please-version
+__version__ = "0.1.0-alpha.42"  # x-release-please-version

payi/lib/AnthropicInstrumentor.py

@@ -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,19 +50,39 @@ def chat_wrapper(
     instrumentor: PayiInstrumentor,
     wrapped: Any,
     instance: Any,
-    args: Any,
-    kwargs: Any,
+    *args: Any,
+    **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        category="system.anthropic",
-        process_chunk=process_chunk,
-        process_request=process_request,
-        process_synchronous_response=process_synchronous_response,
-        is_streaming=IsStreaming.kwargs,
-        wrapped=wrapped,
-        instance=instance,
-        args=args,
-        kwargs=kwargs,
+        "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,
+        kwargs,
     )
 
 
@@ -116,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

@@ -17,12 +17,6 @@ class BedrockInstrumentor:
         try:
             import boto3  # type: ignore #  noqa: F401  I001
 
-            # wrap_function_wrapper(
-            #     "anthropic.resources.completions",
-            #     "Completions.create",
-            #     chat_wrapper(instrumentor),
-            # )
-
             wrap_function_wrapper(
                 "botocore.client",
                 "ClientCreator.create_client",
@@ -40,7 +34,7 @@ class BedrockInstrumentor:
             return
 
 @PayiInstrumentor.payi_wrapper
-def create_client_wrapper(instrumentor: PayiInstrumentor, wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any: #  noqa: ARG001
+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)
 
@@ -108,15 +102,15 @@ def wrap_invoke(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
 
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
-                category="system.aws.bedrock",
-                process_chunk=None,
-                process_request=process_invoke_request,  
-                process_synchronous_response=process_synchronous_invoke_response,  
-                is_streaming=IsStreaming.false,
-                wrapped=wrapped,
-                instance=None,
-                args=args,
-                kwargs=kwargs,
+                "system.aws.bedrock",
+                None,
+                process_invoke_request,  
+                process_synchronous_invoke_response,  
+                IsStreaming.false,
+                wrapped,
+                None,
+                args,
+                kwargs,
         )
         return wrapped(*args, **kwargs)
     
@@ -129,15 +123,15 @@ def wrap_invoke_stream(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
 
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
-                category="system.aws.bedrock",
-                process_chunk=process_invoke_streaming_anthropic_chunk if modelId.startswith("anthropic.") else process_invoke_streaming_llama_chunk,
-                process_request=process_invoke_request, 
-                process_synchronous_response=None,  
-                is_streaming=IsStreaming.true,
-                wrapped=wrapped,
-                instance=None,
-                args=args,
-                kwargs=kwargs,
+                "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)
 
@@ -150,15 +144,15 @@ def wrap_converse(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
 
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
-                category="system.aws.bedrock",
-                process_chunk=None,
-                process_request=process_converse_request,  
-                process_synchronous_response=process_synchronous_converse_response,  
-                is_streaming=IsStreaming.false,
-                wrapped=wrapped,
-                instance=None,
-                args=args,
-                kwargs=kwargs,
+                "system.aws.bedrock",
+                None,
+                process_converse_request,  
+                process_synchronous_converse_response,  
+                IsStreaming.false,
+                wrapped,
+                None,
+                args,
+                kwargs,
         )
         return wrapped(*args, **kwargs)
     
@@ -171,15 +165,15 @@ def wrap_converse_stream(instrumentor: PayiInstrumentor, wrapped: Any) -> Any:
 
         if modelId.startswith("meta.llama3") or modelId.startswith("anthropic."):
             return instrumentor.chat_wrapper(
-                category="system.aws.bedrock",
-                process_chunk=process_converse_streaming_chunk,
-                process_request=process_converse_request, 
-                process_synchronous_response=None,  
-                is_streaming=IsStreaming.true,
-                wrapped=wrapped,
-                instance=None,
-                args=args,
-                kwargs=kwargs,
+                "system.aws.bedrock",
+                process_converse_streaming_chunk,
+                process_converse_request, 
+                None,  
+                IsStreaming.true,
+                wrapped,
+                None,
+                args,
+                kwargs,
             )
         return wrapped(*args, **kwargs)
 
@@ -242,7 +236,7 @@ def process_synchronous_invoke_response(
 
     return response
 
-def process_invoke_request(ingest: IngestUnitsParams, kwargs: Any) -> None: #  noqa: ARG001
+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:
@@ -284,5 +278,5 @@ def process_synchronous_converse_response(
 
     return None    
 
-def process_converse_request(ingest: IngestUnitsParams, kwargs: Any) -> None: #  noqa: ARG001
+def process_converse_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: #  noqa: ARG001
     return

payi/lib/OpenAIInstrumentor.py

@@ -3,7 +3,7 @@ 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
@@ -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,19 +40,39 @@ def chat_wrapper(
     instrumentor: PayiInstrumentor,
     wrapped: Any,
     instance: Any,
-    args: Any,
-    kwargs: Any,
+    *args: Any,
+    **kwargs: Any,
 ) -> Any:
     return instrumentor.chat_wrapper(
-        category="system.openai",
-        process_chunk=process_chat_chunk,
-        process_request=process_request,
-        process_synchronous_response=process_chat_synchronous_response,
-        is_streaming=IsStreaming.kwargs,
-        wrapped=wrapped,
-        instance=instance,
-        args=args,
-        kwargs=kwargs,
+        "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,
+        kwargs,
     )
 
 
@@ -103,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
@@ -112,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