feat: get_headers for assistants and assistants files

portkey_ai/api_resources/apis/assistants.py

@@ -1,6 +1,8 @@
+import json
 from typing import Any
 from portkey_ai.api_resources.apis.api_resource import APIResource, AsyncAPIResource
 from portkey_ai.api_resources.client import AsyncPortkey, Portkey
+from portkey_ai.api_resources.utils import Assistant, AssistantDeleted, AssistantFile, AssistantFileDeleted, AssistantFileList, AssistantList
 
 
 class Assistants(APIResource):
@@ -9,55 +11,86 @@ def __init__(self, client: Portkey) -> None:
         self.openai_client = client.openai_client
         self.files = AssistantFiles(client)
 
-    def create(self, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.create(**kwargs)
-        return response
+    def create(self, **kwargs) -> Assistant:
+        response = self.openai_client.with_raw_response.beta.assistants.create(
+            **kwargs)
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
 
-    def retrieve(self, assistant_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.retrieve(
+        return data
+
+    def retrieve(self, assistant_id, **kwargs) -> Assistant:
+        response = self.openai_client.with_raw_response.beta.assistants.retrieve(
             assistant_id=assistant_id, **kwargs
         )
-        return response
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
-    def update(self, assistant_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.update(
+    def update(self, assistant_id, **kwargs) -> Assistant:
+        response = self.openai_client.with_raw_response.beta.assistants.update(
             assistant_id=assistant_id, **kwargs
         )
-        return response
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    def list(self, **kwargs) -> AssistantList:
+        response = self.openai_client.with_raw_response.beta.assistants.list(
+            **kwargs)
+        data = AssistantList(**json.loads(response.text))
+        data._headers = response.headers
 
-    def delete(self, assistant_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.delete(
+        return data
+
+    def delete(self, assistant_id, **kwargs) -> AssistantDeleted:
+        response = self.openai_client.with_raw_response.beta.assistants.delete(
             assistant_id=assistant_id, **kwargs
         )
-        return response
+        data = AssistantDeleted(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
 
 class AssistantFiles(APIResource):
     def __init__(self, client: Portkey) -> None:
         super().__init__(client)
         self.openai_client = client.openai_client
 
-    def create(self, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.files.create(**kwargs)
-        return response
+    def create(self, assistant_id, file_id, **kwargs) -> AssistantFile:
+        response = self.openai_client.with_raw_response.beta.assistants.files.create(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFile(**json.loads(response.text))
+        data._headers = response.headers
 
-    def list(self, assistant_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.files.list(
-            assistant_id=assistant_id, **kwargs
-        )
-        return response
+        return data
 
-    def retrieve(self, assistant_id, file_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.files.retrieve(
-            assistant_id=assistant_id, file_id=file_id, **kwargs
-        )
-        return response
+    def list(self, assistant_id, **kwargs) -> AssistantFileList:
+        response = self.openai_client.with_raw_response.beta.assistants.files.list(
+            assistant_id=assistant_id, **kwargs)
+        data = AssistantFileList(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
-    def delete(self, assistant_id, file_id, **kwargs) -> Any:
-        response = self.openai_client.beta.assistants.files.delete(
-            assistant_id=assistant_id, file_id=file_id, **kwargs
-        )
-        return response
+    def retrieve(self, assistant_id, file_id, **kwargs) -> AssistantFile:
+        response = self.openai_client.with_raw_response.beta.assistants.files.retrieve(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFile(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    def delete(self, assistant_id, file_id, **kwargs) -> AssistantFileDeleted:
+        response = self.openai_client.with_raw_response.beta.assistants.files.delete(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFileDeleted(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
 
 class AsyncAssistants(AsyncAPIResource):
@@ -66,52 +99,81 @@ def __init__(self, client: AsyncPortkey) -> None:
         self.openai_client = client.openai_client
         self.files = AsyncAssistantFiles(client)
 
-    async def create(self, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.create(**kwargs)
-        return response
+    async def create(self, **kwargs) -> Assistant:
+        response = await self.openai_client.with_raw_response.beta.assistants.create(**kwargs)
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
 
-    async def retrieve(self, assistant_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.retrieve(
+        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
         )
-        return response
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
-    async def update(self, assistant_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.update(
+    async def update(self, assistant_id, **kwargs) -> Assistant:
+        response = await self.openai_client.with_raw_response.beta.assistants.update(
             assistant_id=assistant_id, **kwargs
         )
-        return response
+        data = Assistant(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    async def list(self, **kwargs) -> AssistantList:
+        response = await self.openai_client.with_raw_response.beta.assistants.list(**kwargs)
+        data = AssistantList(**json.loads(response.text))
+        data._headers = response.headers
 
-    async def delete(self, assistant_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.delete(
+        return data
+
+    async def delete(self, assistant_id, **kwargs) -> AssistantDeleted:
+        response = await self.openai_client.with_raw_response.beta.assistants.delete(
             assistant_id=assistant_id, **kwargs
         )
-        return response
+        data = AssistantDeleted(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
 
 
 class AsyncAssistantFiles(AsyncAPIResource):
     def __init__(self, client: AsyncPortkey) -> None:
         super().__init__(client)
         self.openai_client = client.openai_client
 
-    async def create(self, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.files.create(**kwargs)
-        return response
-
-    async def list(self, assistant_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.files.list(
-            assistant_id=assistant_id, **kwargs
-        )
-        return response
-
-    async def retrieve(self, assistant_id, file_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.files.retrieve(
-            assistant_id=assistant_id, file_id=file_id, **kwargs
-        )
-        return response
-
-    async def delete(self, assistant_id, file_id, **kwargs) -> Any:
-        response = await self.openai_client.beta.assistants.files.delete(
-            assistant_id=assistant_id, file_id=file_id, **kwargs
-        )
-        return response
+    async def create(self, assistant_id, file_id, **kwargs) -> AssistantFile:
+        response = await self.openai_client.with_raw_response.beta.assistants.files.create(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFile(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    async def list(self, assistant_id, **kwargs) -> AssistantFileList:
+        response = await self.openai_client.with_raw_response.beta.assistants.files.list(
+            assistant_id=assistant_id, **kwargs)
+        data = AssistantFileList(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    async def retrieve(self, assistant_id, file_id, **kwargs) -> AssistantFile:
+        response = await self.openai_client.with_raw_response.beta.assistants.files.retrieve(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFile(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data
+
+    async def delete(self, assistant_id, file_id, **kwargs) -> AssistantFileDeleted:
+        response = await self.openai_client.with_raw_response.beta.assistants.files.delete(
+            assistant_id=assistant_id, file_id=file_id, **kwargs)
+        data = AssistantFileDeleted(**json.loads(response.text))
+        data._headers = response.headers
+
+        return data

portkey_ai/api_resources/utils.py

@@ -568,6 +568,51 @@ def __str__(self):
     def get_headers(self) -> Optional[Dict[str, str]]:
         return parse_headers(self._headers)
 
+class Assistant(BaseModel):
+    id: Optional[str]
+    created_at: Optional[int]
+    description: Optional[str] = None
+    file_ids: Optional[List[str]]
+    instructions: Optional[str] = None
+    metadata: Optional[object] = None
+    model: Optional[str]
+    name: Optional[str] = None
+    object: Optional[str]
+    tools: Optional[List[Any]]
+    _headers: Optional[httpx.Headers] = None
+
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)
+
+class AssistantList(BaseModel):
+    object: Optional[str]
+    data: Optional[List[Any]]
+    _headers: Optional[httpx.Headers] = None
+
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)
+
+class AssistantDeleted(BaseModel):
+    id: Optional[str]
+    object: Optional[str]
+    deleted: Optional[bool]
+    _headers: Optional[httpx.Headers] = None
+
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)
+
 
 def apikey_from_env(provider: Union[ProviderTypes, ProviderTypesLiteral, str]) -> str:
     env_key = f"{provider.upper().replace('-', '_')}_API_KEY"
@@ -689,23 +734,45 @@ def parse_headers(headers: Optional[httpx.Headers]) -> dict:
     return _headers
 
 
-class FileDeleteResponse(BaseModel):
-    id: str
-
-    deleted: bool
-
-    object: Literal["assistant.file.deleted"]
+class AssistantFileDeleted(BaseModel):
+    id: Optional[str]
+    deleted: Optional[bool]
+    object: Optional[str]
+    _headers: Optional[httpx.Headers] = None
 
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)
+
 
 class AssistantFile(BaseModel):
-    id: str
-    """The identifier, which can be referenced in API endpoints."""
-
-    assistant_id: str
-    """The assistant ID that the file is attached to."""
+    id: Optional[str]
+    assistant_id: Optional[str]
+    created_at: Optional[int]
+    object: Optional[str]
+    _headers: Optional[httpx.Headers] = None
 
-    created_at: int
-    """The Unix timestamp (in seconds) for when the assistant file was created."""
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)
+
+class AssistantFileList(BaseModel):
+    object: Optional[str]
+    data: Optional[List[Any]]
+    first_id: Optional[str]
+    last_id: Optional[str]
+    has_more: Optional[bool]
+    _headers: Optional[httpx.Headers] = None
 
-    object: Literal["assistant.file"]
-    """The object type, which is always `assistant.file`."""
+    def __str__(self):
+        del self._headers
+        return json.dumps(self.dict(), indent=4)
+
+    def get_headers(self) -> Optional[Dict[str, str]]:
+        return parse_headers(self._headers)