From bb5989e675fe54e6abbc37ecb6ce5bb887c9e5c0 Mon Sep 17 00:00:00 2001
From: csgulati09 <chandeepsinghgulati@gmail.com>
Date: Fri, 20 Dec 2024 13:20:02 +0530
Subject: [PATCH] feat: async chat completion request type

---
 .../api_resources/apis/chat_complete.py       | 67 ++++++++++++++++++-
 1 file changed, 65 insertions(+), 2 deletions(-)

diff --git a/portkey_ai/api_resources/apis/chat_complete.py b/portkey_ai/api_resources/apis/chat_complete.py
index de523db..cd89921 100644
--- a/portkey_ai/api_resources/apis/chat_complete.py
+++ b/portkey_ai/api_resources/apis/chat_complete.py
@@ -189,7 +189,21 @@ def __init__(self, client: AsyncPortkey) -> None:
         self.openai_client = client.openai_client
 
     async def stream_create(
-        self, model, messages, stream, temperature, max_tokens, top_p, **kwargs
+        self,
+        model,
+        messages,
+        stream,
+        temperature,
+        max_tokens,
+        top_p,
+        audio,
+        max_completion_tokens,
+        metadata,
+        modalities,
+        prediction,
+        reasoning_effort,
+        store,
+        **kwargs,
     ) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
         async with self.openai_client.with_streaming_response.chat.completions.create(
             model=model,
@@ -198,6 +212,13 @@ async def stream_create(
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
+            audio=audio,
+            max_completion_tokens=max_completion_tokens,
+            metadata=metadata,
+            modalities=modalities,
+            prediction=prediction,
+            reasoning_effort=reasoning_effort,
+            store=store,
             extra_body=kwargs,
         ) as response:
             async for line in response.iter_lines():
@@ -215,7 +236,21 @@ async def stream_create(
                     pass
 
     async def normal_create(
-        self, model, messages, stream, temperature, max_tokens, top_p, **kwargs
+        self,
+        model,
+        messages,
+        stream,
+        temperature,
+        max_tokens,
+        top_p,
+        audio,
+        max_completion_tokens,
+        metadata,
+        modalities,
+        prediction,
+        reasoning_effort,
+        store,
+        **kwargs,
     ) -> ChatCompletions:
         response = await self.openai_client.with_raw_response.chat.completions.create(
             model=model,
@@ -224,6 +259,13 @@ async def normal_create(
             temperature=temperature,
             max_tokens=max_tokens,
             top_p=top_p,
+            audio=audio,
+            max_completion_tokens=max_completion_tokens,
+            metadata=metadata,
+            modalities=modalities,
+            prediction=prediction,
+            reasoning_effort=reasoning_effort,
+            store=store,
             extra_body=kwargs,
         )
         data = ChatCompletions(**json.loads(response.text))
@@ -239,6 +281,13 @@ async def create(
         temperature: Union[float, NotGiven] = NOT_GIVEN,
         max_tokens: Union[int, NotGiven] = NOT_GIVEN,
         top_p: Union[float, NotGiven] = NOT_GIVEN,
+        audio: Optional[Any] = NOT_GIVEN,
+        max_completion_tokens: Union[int, NotGiven] = NOT_GIVEN,
+        metadata: Union[Dict[str, str], NotGiven] = NOT_GIVEN,
+        modalities: Union[List[Any], NotGiven] = NOT_GIVEN,
+        prediction: Union[Any, NotGiven] = NOT_GIVEN,
+        reasoning_effort: Union[Any, NotGiven] = NOT_GIVEN,
+        store: Union[Optional[bool], NotGiven] = NOT_GIVEN,
         **kwargs,
     ) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
         if stream is True:
@@ -249,6 +298,13 @@ async def create(
                 temperature=temperature,
                 max_tokens=max_tokens,
                 top_p=top_p,
+                audio=audio,
+                max_completion_tokens=max_completion_tokens,
+                metadata=metadata,
+                modalities=modalities,
+                prediction=prediction,
+                reasoning_effort=reasoning_effort,
+                store=store,
                 **kwargs,
             )
         else:
@@ -259,6 +315,13 @@ async def create(
                 temperature=temperature,
                 max_tokens=max_tokens,
                 top_p=top_p,
+                audio=audio,
+                max_completion_tokens=max_completion_tokens,
+                metadata=metadata,
+                modalities=modalities,
+                prediction=prediction,
+                reasoning_effort=reasoning_effort,
+                store=store,
                 **kwargs,
             )