Merge pull request #248 from Portkey-AI/fix/pydanticv2

Pydantic for v1 and v2

portkey_ai/api_resources/apis/__init__.py

@@ -1,3 +1,10 @@
+"""
+The following change is brought to handle the cases for pydantic v1 and v2
+Including (import sys) and
+sys.modules["openai"] = vendored_openai
+"""
+import sys
+from ..._vendor import openai as vendored_openai
 from .chat_complete import ChatCompletion, AsyncChatCompletion
 from .complete import Completion, AsyncCompletion
 from .generation import Generations, AsyncGenerations, Prompts, AsyncPrompts
@@ -80,6 +87,8 @@
 from .virtual_keys import VirtualKeys, AsyncVirtualKeys
 from .logs import Logs, AsyncLogs
 
+sys.modules["openai"] = vendored_openai  # For pydantic v1 and v2 compatibility
+
 __all__ = [
     "Completion",
     "AsyncCompletion",

portkey_ai/api_resources/apis/beta_chat.py

@@ -1,9 +1,7 @@
 from typing import Any, Union
-from portkey_ai._vendor.openai.types.chat.parsed_chat_completion import (
-    ParsedChatCompletion,
-)
 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.types.beta_chat_type import ParsedChatCompletion
 from ..._vendor.openai._types import NotGiven, NOT_GIVEN
 
 
@@ -27,15 +25,15 @@ def parse(
         response_format: Union[Any, NotGiven] = NOT_GIVEN,
         tools: Union[Any, NotGiven] = NOT_GIVEN,
         **kwargs: Any,
-    ) -> ParsedChatCompletion[Any]:
+    ) -> ParsedChatCompletion:
         response = self.openai_client.beta.chat.completions.parse(
             messages=messages,
             model=model,
             response_format=response_format,
             tools=tools,
             extra_body=kwargs,
         )
-        return response
+        return response  # type: ignore [return-value]
 
     def stream(
         self,
@@ -85,15 +83,15 @@ async def parse(
         response_format: Union[Any, NotGiven] = NOT_GIVEN,
         tools: Union[Any, NotGiven] = NOT_GIVEN,
         **kwargs: Any,
-    ) -> ParsedChatCompletion[Any]:
+    ) -> ParsedChatCompletion:
         response = await self.openai_client.beta.chat.completions.parse(
             messages=messages,
             model=model,
             response_format=response_format,
             tools=tools,
             extra_body=kwargs,
         )
-        return response
+        return response  # type: ignore [return-value]
 
     async def stream(
         self,

portkey_ai/api_resources/types/api_keys_type.py

@@ -5,10 +5,10 @@
 from portkey_ai.api_resources.types.utils import parse_headers
 
 
-class ApiKeyAddResponse(BaseModel):
-    id: Optional[str]
-    key: Optional[str]
-    object: Optional[str]
+class ApiKeyAddResponse(BaseModel, extra="allow"):
+    id: Optional[str] = None
+    key: Optional[str] = None
+    object: Optional[str] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def get_headers(self) -> Optional[Dict[str, str]]:
@@ -24,25 +24,25 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class ApiKeyGetResponse(BaseModel):
-    id: Optional[str]
-    key: Optional[str]
-    name: Optional[str]
-    description: Optional[str]
-    type: Optional[str]
-    organisation_id: Optional[str]
-    workspace_id: Optional[str]
-    user_id: Optional[str]
-    status: Optional[str]
-    created_at: Optional[str]
-    last_updated_at: Optional[str]
-    creation_mode: Optional[str]
-    rate_limits: Optional[List[Dict[str, Any]]]
-    usage_limits: Optional[Dict[str, Any]]
-    reset_usage: Optional[int]
-    scopes: Optional[List[str]]
-    defaults: Optional[Dict[str, Any]]
-    object: Optional[str]
+class ApiKeyGetResponse(BaseModel, extra="allow"):
+    id: Optional[str] = None
+    key: Optional[str] = None
+    name: Optional[str] = None
+    description: Optional[str] = None
+    type: Optional[str] = None
+    organisation_id: Optional[str] = None
+    workspace_id: Optional[str] = None
+    user_id: Optional[str] = None
+    status: Optional[str] = None
+    created_at: Optional[str] = None
+    last_updated_at: Optional[str] = None
+    creation_mode: Optional[str] = None
+    rate_limits: Optional[List[Dict[str, Any]]] = None
+    usage_limits: Optional[Dict[str, Any]] = None
+    reset_usage: Optional[int] = None
+    scopes: Optional[List[str]] = None
+    defaults: Optional[Dict[str, Any]] = None
+    object: Optional[str] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def get_headers(self) -> Optional[Dict[str, str]]:
@@ -58,10 +58,10 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class ApiKeyListResponse(BaseModel):
-    object: Optional[str]
-    total: Optional[int]
-    data: Optional[List[Dict[str, Any]]]
+class ApiKeyListResponse(BaseModel, extra="allow"):
+    object: Optional[str] = None
+    total: Optional[int] = None
+    data: Optional[List[Dict[str, Any]]] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def get_headers(self) -> Optional[Dict[str, str]]:
@@ -77,10 +77,10 @@ def get(self, key: str, default: Optional[Any] = None):
         return getattr(self, key, None) or default
 
 
-class ApiKeyUpdateResponse(BaseModel):
-    object: Optional[str]
-    total: Optional[int]
-    data: Optional[List[Dict[str, Any]]]
+class ApiKeyUpdateResponse(BaseModel, extra="allow"):
+    object: Optional[str] = None
+    total: Optional[int] = None
+    data: Optional[List[Dict[str, Any]]] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def get_headers(self) -> Optional[Dict[str, str]]:

portkey_ai/api_resources/types/assistant_type.py

@@ -12,36 +12,61 @@
     "ToolCodeInterpreter",
     "ToolRetrieval",
     "ToolFunction",
-    "Tool",
+    "AssistantTool",
+    "FunctionDefinition",
+    "FunctionParameters",
+    "FileSearch",
+    "FileSearchTool",
 ]
 
 
-class ToolCodeInterpreter(BaseModel):
-    type: Optional[str]
+class FileSearch(BaseModel, extra="allow"):
+    max_num_results: Optional[int] = None
 
 
-class ToolRetrieval(BaseModel):
-    type: Optional[str]
+class FileSearchTool(BaseModel, extra="allow"):
+    type: Optional[str] = None
+    file_search: Optional[FileSearch] = None
 
 
-class ToolFunction(BaseModel):
-    type: Optional[str]
+FunctionParameters = Dict[str, object]
 
 
-Tool = Union[ToolCodeInterpreter, ToolRetrieval, ToolFunction]
+class FunctionDefinition(BaseModel, extra="allow"):
+    name: Optional[str] = None
+    description: Optional[str] = None
+    parameters: Optional[FunctionParameters] = None
+    strict: Optional[bool] = None
+
+
+class ToolCodeInterpreter(BaseModel, extra="allow"):
+    type: Optional[str] = None
+
+
+class ToolRetrieval(BaseModel, extra="allow"):
+    type: Optional[str] = None
 
 
-class Assistant(BaseModel):
-    id: Optional[str]
-    created_at: Optional[int]
+class ToolFunction(BaseModel, extra="allow"):
+    type: Optional[str] = None
+    function: Optional[FunctionDefinition] = None
+
+
+AssistantTool = Union[ToolCodeInterpreter, ToolRetrieval, ToolFunction, FileSearchTool]
+
+
+class Assistant(BaseModel, extra="allow"):
+    id: Optional[str] = None
+    created_at: Optional[int] = None
     description: Optional[str] = None
     file_ids: Optional[List[str]] = None
     instructions: Optional[str] = None
     metadata: Optional[object] = None
-    model: Optional[str]
+    model: Optional[str] = None
     name: Optional[str] = None
-    object: Optional[str]
-    tools: Optional[List[Tool]]
+    object: Optional[str] = None
+    tools: Optional[List[AssistantTool]] = None
+    response_format: Optional[Any] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def __str__(self):
@@ -59,7 +84,7 @@ def get_headers(self) -> Optional[Dict[str, str]]:
 
 
 class AssistantList(BaseModel, extra="allow"):
-    object: Optional[str]
+    object: Optional[str] = None
     data: Optional[List[Assistant]]
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
@@ -78,9 +103,9 @@ def get_headers(self) -> Optional[Dict[str, str]]:
 
 
 class AssistantDeleted(BaseModel, extra="allow"):
-    id: Optional[str]
-    object: Optional[str]
-    deleted: Optional[bool]
+    id: Optional[str] = None
+    object: Optional[str] = None
+    deleted: Optional[bool] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
     def __str__(self):
@@ -89,48 +114,3 @@ def __str__(self):
 
     def get_headers(self) -> Optional[Dict[str, str]]:
         return parse_headers(self._headers)
-
-
-# class AssistantFile(BaseModel, extra="allow"):
-#     id: Optional[str]
-#     assistant_id: Optional[str]
-#     created_at: Optional[int]
-#     object: Optional[str]
-#     _headers: Optional[httpx.Headers] = PrivateAttr()
-
-#     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, extra="allow"):
-#     object: Optional[str]
-#     data: Optional[List[AssistantFile]]
-#     first_id: Optional[str]
-#     last_id: Optional[str]
-#     has_more: Optional[bool]
-#     _headers: Optional[httpx.Headers] = PrivateAttr()
-
-#     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 AssistantFileDeleted(BaseModel, extra="allow"):
-#     id: Optional[str]
-#     deleted: Optional[bool]
-#     object: Optional[str]
-#     _headers: Optional[httpx.Headers] = PrivateAttr()
-
-#     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)

portkey_ai/api_resources/types/audio_types.py

@@ -8,7 +8,7 @@
 __all__ = ["Transcription", "Translation"]
 
 
-class Transcription(BaseModel):
+class Transcription(BaseModel, extra="allow"):
     text: str
     _headers: Optional[httpx.Headers] = PrivateAttr()
 
@@ -26,7 +26,7 @@ def get_headers(self) -> Optional[Dict[str, str]]:
         return parse_headers(self._headers)
 
 
-class Translation(BaseModel):
+class Translation(BaseModel, extra="allow"):
     text: str
     _headers: Optional[httpx.Headers] = PrivateAttr()
 

portkey_ai/api_resources/types/batches_type.py

@@ -5,20 +5,29 @@
 from .utils import parse_headers
 from typing import List, Any
 from pydantic import BaseModel, PrivateAttr
-from ..._vendor.openai.types.batch_error import BatchError
-from ..._vendor.openai.types.batch_request_counts import BatchRequestCounts
 
 __all__ = ["Batch", "BatchList", "Errors"]
 
 
-class Errors(BaseModel):
-    data: Optional[List[BatchError]] = None
+class BatchError(BaseModel, extra="allow"):
+    code: Optional[str] = None
+    line: Optional[int] = None
+    message: Optional[str] = None
+    param: Optional[str] = None
+
+
+class BatchRequestCounts(BaseModel, extra="allow"):
+    completed: Optional[int] = None
+    failed: Optional[int] = None
+    total: Optional[int] = None
 
+
+class Errors(BaseModel, extra="allow"):
+    data: Optional[List[BatchError]] = None
     object: Optional[str] = None
-    """The object type, which is always `list`."""
 
 
-class Batch(BaseModel):
+class Batch(BaseModel, extra="allow"):
     id: str
     completion_window: str
     created_at: int
@@ -55,7 +64,7 @@ def get_headers(self) -> Optional[Dict[str, str]]:
         return parse_headers(self._headers)
 
 
-class BatchList(BaseModel):
+class BatchList(BaseModel, extra="allow"):
     object: Optional[str] = None
     data: Optional[List[Batch]] = None
     _headers: Optional[httpx.Headers] = PrivateAttr()