Merge pull request #256 from Portkey-AI/fix/retrieveCalls

extra_body
Portkey-AI · Dec 11, 2024 · 5c723ec · 5c723ec
2 parents 142d951 + f3eb5d4
commit 5c723ec
Show file tree

Hide file tree

Showing 7 changed files with 255 additions and 95 deletions.
diff --git a/portkey_ai/api_resources/apis/assistants.py b/portkey_ai/api_resources/apis/assistants.py
@@ -48,9 +48,14 @@ def create(
         return data
 
     def retrieve(self, assistant_id, **kwargs) -> Assistant:
-        response = self.openai_client.with_raw_response.beta.assistants.retrieve(
-            assistant_id=assistant_id, **kwargs
-        )
+        if kwargs:
+            response = self.openai_client.with_raw_response.beta.assistants.retrieve(
+                assistant_id=assistant_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.with_raw_response.beta.assistants.retrieve(
+                assistant_id=assistant_id
+            )
         data = Assistant(**json.loads(response.text))
         data._headers = response.headers
 
@@ -156,9 +161,18 @@ async def create(
         return data
 
     async def retrieve(self, assistant_id, **kwargs) -> Assistant:
-        response = await self.openai_client.with_raw_response.beta.assistants.retrieve(
-            assistant_id=assistant_id, **kwargs
-        )
+        if kwargs:
+            response = (
+                await self.openai_client.with_raw_response.beta.assistants.retrieve(
+                    assistant_id=assistant_id, extra_body=kwargs
+                )
+            )
+        else:
+            response = (
+                await self.openai_client.with_raw_response.beta.assistants.retrieve(
+                    assistant_id=assistant_id
+                )
+            )
         data = Assistant(**json.loads(response.text))
         data._headers = response.headers
 

diff --git a/portkey_ai/api_resources/apis/batches.py b/portkey_ai/api_resources/apis/batches.py
@@ -36,9 +36,14 @@ def create(
         return data
 
     def retrieve(self, batch_id, **kwargs) -> Batch:
-        response = self.openai_client.with_raw_response.batches.retrieve(
-            batch_id=batch_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = self.openai_client.with_raw_response.batches.retrieve(
+                batch_id=batch_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.with_raw_response.batches.retrieve(
+                batch_id=batch_id
+            )
         data = Batch(**json.loads(response.text))
         data._headers = response.headers
 
@@ -97,9 +102,14 @@ async def create(
         return data
 
     async def retrieve(self, batch_id, **kwargs) -> Batch:
-        response = await self.openai_client.with_raw_response.batches.retrieve(
-            batch_id=batch_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = await self.openai_client.with_raw_response.batches.retrieve(
+                batch_id=batch_id, extra_body=kwargs
+            )
+        else:
+            response = await self.openai_client.with_raw_response.batches.retrieve(
+                batch_id=batch_id
+            )
         data = Batch(**json.loads(response.text))
         data._headers = response.headers
 

diff --git a/portkey_ai/api_resources/apis/fine_tuning.py b/portkey_ai/api_resources/apis/fine_tuning.py
@@ -56,9 +56,14 @@ def create(
         return data
 
     def retrieve(self, fine_tuning_job_id: str, **kwargs) -> FineTuningJob:
-        response = self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
-            fine_tuning_job_id=fine_tuning_job_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
+                fine_tuning_job_id=fine_tuning_job_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
+                fine_tuning_job_id=fine_tuning_job_id
+            )
         data = FineTuningJob(**json.loads(response.text))
         data._headers = response.headers
 
@@ -179,9 +184,19 @@ async def create(
         return data
 
     async def retrieve(self, fine_tuning_job_id: str, **kwargs) -> FineTuningJob:
-        response = await self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
-            fine_tuning_job_id=fine_tuning_job_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = (
+                await self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
+                    fine_tuning_job_id=fine_tuning_job_id, extra_body=kwargs
+                )
+            )
+        else:
+            response = (
+                await self.openai_client.with_raw_response.fine_tuning.jobs.retrieve(
+                    fine_tuning_job_id=fine_tuning_job_id
+                )
+            )
+
         data = FineTuningJob(**json.loads(response.text))
         data._headers = response.headers
 

diff --git a/portkey_ai/api_resources/apis/main_files.py b/portkey_ai/api_resources/apis/main_files.py
@@ -35,9 +35,14 @@ def list(self, purpose: Union[str, NotGiven] = NOT_GIVEN, **kwargs) -> FileList:
         return data
 
     def retrieve(self, file_id, **kwargs) -> FileObject:
-        response = self.openai_client.with_raw_response.files.retrieve(
-            file_id=file_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = self.openai_client.with_raw_response.files.retrieve(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.with_raw_response.files.retrieve(
+                file_id=file_id
+            )
         data = FileObject(**json.loads(response.text))
         data._headers = response.headers
 
@@ -53,11 +58,21 @@ def delete(self, file_id, **kwargs) -> FileDeleted:
         return data
 
     def content(self, file_id, **kwargs) -> Any:
-        response = self.openai_client.files.content(file_id=file_id, extra_body=kwargs)
+        if kwargs:
+            response = self.openai_client.files.content(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.files.content(file_id=file_id)
         return response
 
     def retrieve_content(self, file_id, **kwargs) -> Any:
-        response = self.openai_client.files.content(file_id=file_id, extra_body=kwargs)
+        if kwargs:
+            response = self.openai_client.files.content(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.files.content(file_id=file_id)
         return response
 
     def wait_for_processing(
@@ -102,9 +117,14 @@ async def list(
         return data
 
     async def retrieve(self, file_id, **kwargs) -> FileObject:
-        response = await self.openai_client.with_raw_response.files.retrieve(
-            file_id=file_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = await self.openai_client.with_raw_response.files.retrieve(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = await self.openai_client.with_raw_response.files.retrieve(
+                file_id=file_id
+            )
         data = FileObject(**json.loads(response.text))
         data._headers = response.headers
 
@@ -120,15 +140,21 @@ async def delete(self, file_id, **kwargs) -> FileDeleted:
         return data
 
     async def content(self, file_id, **kwargs) -> Any:
-        response = await self.openai_client.files.content(
-            file_id=file_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = await self.openai_client.files.content(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = await self.openai_client.files.content(file_id=file_id)
         return response
 
     async def retrieve_content(self, file_id, **kwargs) -> Any:
-        response = await self.openai_client.files.content(
-            file_id=file_id, extra_body=kwargs
-        )
+        if kwargs:
+            response = await self.openai_client.files.content(
+                file_id=file_id, extra_body=kwargs
+            )
+        else:
+            response = await self.openai_client.files.content(file_id=file_id)
         return response
 
     async def wait_for_processing(

diff --git a/portkey_ai/api_resources/apis/models.py b/portkey_ai/api_resources/apis/models.py
@@ -20,9 +20,14 @@ def list(self, **kwargs) -> ModelList:
     def retrieve(
         self, model: str, *, timeout: Union[float, NotGiven] = NOT_GIVEN, **kwargs
     ) -> Model:
-        response = self.openai_client.with_raw_response.models.retrieve(
-            model=model, timeout=timeout, extra_body=kwargs
-        )
+        if kwargs:
+            response = self.openai_client.with_raw_response.models.retrieve(
+                model=model, timeout=timeout, extra_body=kwargs
+            )
+        else:
+            response = self.openai_client.with_raw_response.models.retrieve(
+                model=model, timeout=timeout
+            )
         data = Model(**json.loads(response.text))
         data._headers = response.headers
         return data
@@ -52,9 +57,14 @@ async def list(self, **kwargs) -> ModelList:
     async def retrieve(
         self, model: str, *, timeout: Union[float, NotGiven] = NOT_GIVEN, **kwargs
     ) -> Model:
-        response = await self.openai_client.with_raw_response.models.retrieve(
-            model=model, timeout=timeout, extra_body=kwargs
-        )
+        if kwargs:
+            response = await self.openai_client.with_raw_response.models.retrieve(
+                model=model, timeout=timeout, extra_body=kwargs
+            )
+        else:
+            response = await self.openai_client.with_raw_response.models.retrieve(
+                model=model, timeout=timeout
+            )
         data = Model(**json.loads(response.text))
         data._headers = response.headers
         return data