Feat/async

portkey_ai/__init__.py

@@ -12,16 +12,20 @@
     PortkeyResponse,
     ChatCompletions,
     Completion,
+    AsyncCompletion,
     Params,
     Config,
     RetrySettings,
     ChatCompletion,
+    AsyncChatCompletion,
     ChatCompletionChunk,
     TextCompletion,
     TextCompletionChunk,
     createHeaders,
     Prompts,
+    AsyncPrompts,
     Portkey,
+    AsyncPortkey,
 )
 from portkey_ai.version import VERSION
 from portkey_ai.api_resources.global_constants import (
@@ -49,9 +53,11 @@
     "Message",
     "ChatCompletions",
     "Completion",
+    "AsyncCompletion",
     "Params",
     "RetrySettings",
     "ChatCompletion",
+    "AsyncChatCompletion",
     "ChatCompletionChunk",
     "TextCompletion",
     "TextCompletionChunk",
@@ -61,5 +67,7 @@
     "PORTKEY_GATEWAY_URL",
     "createHeaders",
     "Prompts",
+    "AsyncPrompts",
     "Portkey",
+    "AsyncPortkey",
 ]

portkey_ai/api_resources/__init__.py

@@ -1,10 +1,15 @@
 """"""
 from .apis import (
     Completion,
+    AsyncCompletion,
     ChatCompletion,
+    AsyncChatCompletion,
     Generations,
+    AsyncGenerations,
     Prompts,
+    AsyncPrompts,
     Feedback,
+    AsyncFeedback,
     createHeaders,
 )
 from .utils import (
@@ -25,7 +30,7 @@
     TextCompletion,
     TextCompletionChunk,
 )
-from .client import Portkey
+from .client import Portkey, AsyncPortkey
 
 from portkey_ai.version import VERSION
 
@@ -42,16 +47,22 @@
     "Message",
     "ChatCompletions",
     "Completion",
+    "AsyncCompletion",
     "Params",
     "Config",
     "RetrySettings",
     "ChatCompletion",
+    "AsyncChatCompletion",
     "ChatCompletionChunk",
     "TextCompletion",
     "TextCompletionChunk",
     "Generations",
+    "AsyncGenerations",
     "Prompts",
+    "AsyncPrompts",
     "Feedback",
+    "AsyncFeedback",
     "createHeaders",
     "Portkey",
+    "AsyncPortkey",
 ]

portkey_ai/api_resources/apis/__init__.py

@@ -1,18 +1,25 @@
-from .chat_complete import ChatCompletion
-from .complete import Completion
-from .generation import Generations, Prompts
-from .feedback import Feedback
+from .chat_complete import ChatCompletion, AsyncChatCompletion
+from .complete import Completion, AsyncCompletion
+from .generation import Generations, AsyncGenerations, Prompts, AsyncPrompts
+from .feedback import Feedback, AsyncFeedback
 from .create_headers import createHeaders
-from .post import Post
-from .embeddings import Embeddings
+from .post import Post, AsyncPost
+from .embeddings import Embeddings, AsyncEmbeddings
 
 __all__ = [
     "Completion",
+    "AsyncCompletion",
     "ChatCompletion",
+    "AsyncChatCompletion",
     "Generations",
+    "AsyncGenerations",
     "Feedback",
+    "AsyncFeedback",
     "Prompts",
+    "AsyncPrompts",
     "createHeaders",
     "Post",
+    "AsyncPost",
     "Embeddings",
+    "AsyncEmbeddings",
 ]

portkey_ai/api_resources/apis/api_resource.py

@@ -1,4 +1,5 @@
-from portkey_ai.api_resources.base_client import APIClient
+from portkey_ai.api_resources.base_client import APIClient, AsyncAPIClient
+import asyncio
 
 
 class APIResource:
@@ -17,3 +18,20 @@ def __init__(self, client: APIClient) -> None:
 
     def _post(self, *args, **kwargs):
         return self._client._post(*args, **kwargs)
+
+
+class AsyncAPIResource:
+    _client: AsyncAPIClient
+
+    def __init__(self, client: AsyncAPIClient) -> None:
+        self._client = client
+        # self._get = client.get
+        # self._patch = client.patch
+        # self._put = client.put
+        # self._delete = client.delete
+
+    async def _post(self, *args, **kwargs):
+        return await self._client._post(*args, **kwargs)
+
+    async def _sleep(self, seconds: float) -> None:
+        await asyncio.sleep(seconds)

portkey_ai/api_resources/apis/chat_complete.py

@@ -2,19 +2,19 @@
 
 import json
 from typing import Mapping, Optional, Union, overload, Literal, List
-from portkey_ai.api_resources.base_client import APIClient
+from portkey_ai.api_resources.base_client import APIClient, AsyncAPIClient
 from portkey_ai.api_resources.utils import (
     PortkeyApiPaths,
     Message,
     ChatCompletionChunk,
     ChatCompletions,
 )
 
-from portkey_ai.api_resources.streaming import Stream
-from portkey_ai.api_resources.apis.api_resource import APIResource
+from portkey_ai.api_resources.streaming import AsyncStream, Stream
+from portkey_ai.api_resources.apis.api_resource import APIResource, AsyncAPIResource
 
 
-__all__ = ["ChatCompletion"]
+__all__ = ["ChatCompletion", "AsyncChatCompletion"]
 
 
 class ChatCompletion(APIResource):
@@ -25,6 +25,14 @@ def __init__(self, client: APIClient) -> None:
         self.completions = Completions(client)
 
 
+class AsyncChatCompletion(AsyncAPIResource):
+    completions: AsyncCompletions
+
+    def __init__(self, client: AsyncAPIClient) -> None:
+        super().__init__(client)
+        self.completions = AsyncCompletions(client)
+
+
 class Completions(APIResource):
     def __init__(self, client: APIClient) -> None:
         super().__init__(client)
@@ -107,3 +115,87 @@ def create(
 
     def _get_config_string(self, config: Union[Mapping, str]) -> str:
         return config if isinstance(config, str) else json.dumps(config)
+
+
+class AsyncCompletions(AsyncAPIResource):
+    def __init__(self, client: AsyncAPIClient) -> None:
+        super().__init__(client)
+
+    @overload
+    async def create(
+        self,
+        *,
+        messages: Optional[List[Message]] = None,
+        config: Optional[Union[Mapping, str]] = None,
+        stream: Literal[True],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> AsyncStream[ChatCompletionChunk]:
+        ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        messages: Optional[List[Message]] = None,
+        config: Optional[Union[Mapping, str]] = None,
+        stream: Literal[False] = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> ChatCompletions:
+        ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        messages: Optional[List[Message]] = None,
+        config: Optional[Union[Mapping, str]] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> Union[ChatCompletions, AsyncStream[ChatCompletionChunk]]:
+        ...
+
+    async def create(
+        self,
+        *,
+        messages: Optional[List[Message]] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> Union[ChatCompletions, AsyncStream[ChatCompletionChunk]]:
+        body = dict(
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_k=top_k,
+            top_p=top_p,
+            stream=stream,
+            **kwargs,
+        )
+
+        return await self._post(
+            PortkeyApiPaths.CHAT_COMPLETE_API,
+            body=body,
+            params=None,
+            cast_to=ChatCompletions,
+            stream_cls=AsyncStream[ChatCompletionChunk],
+            stream=stream,
+            headers={},
+        )
+
+    def _get_config_string(self, config: Union[Mapping, str]) -> str:
+        return config if isinstance(config, str) else json.dumps(config)

portkey_ai/api_resources/apis/complete.py

@@ -1,13 +1,13 @@
 from typing import Optional, Union, overload, Literal
-from portkey_ai.api_resources.base_client import APIClient
+from portkey_ai.api_resources.base_client import APIClient, AsyncAPIClient
 from portkey_ai.api_resources.utils import (
     PortkeyApiPaths,
     TextCompletion,
     TextCompletionChunk,
 )
 
-from portkey_ai.api_resources.streaming import Stream
-from portkey_ai.api_resources.apis.api_resource import APIResource
+from portkey_ai.api_resources.streaming import AsyncStream, Stream
+from portkey_ai.api_resources.apis.api_resource import APIResource, AsyncAPIResource
 
 
 class Completion(APIResource):
@@ -85,3 +85,80 @@ def create(
             stream=stream,
             headers={},
         )
+
+
+class AsyncCompletion(AsyncAPIResource):
+    def __init__(self, client: AsyncAPIClient) -> None:
+        super().__init__(client)
+
+    @overload
+    async def create(
+        self,
+        *,
+        prompt: Optional[str] = None,
+        stream: Literal[True],
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> AsyncStream[TextCompletionChunk]:
+        ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        prompt: Optional[str] = None,
+        stream: Literal[False] = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> TextCompletion:
+        ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        prompt: Optional[str] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> Union[TextCompletion, AsyncStream[TextCompletionChunk]]:
+        ...
+
+    async def create(
+        self,
+        *,
+        prompt: Optional[str] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs,
+    ) -> Union[TextCompletion, AsyncStream[TextCompletionChunk]]:
+        body = dict(
+            prompt=prompt,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_k=top_k,
+            top_p=top_p,
+            stream=stream,
+            **kwargs,
+        )
+        return await self._post(
+            PortkeyApiPaths.TEXT_COMPLETE_API,
+            body=body,
+            params=None,
+            cast_to=TextCompletion,
+            stream_cls=AsyncStream[TextCompletionChunk],
+            stream=stream,
+            headers={},
+        )