fix(types): fix typing

This commit is contained in:
Krrish Dholakia 2024-05-14 17:09:36 -07:00
parent 888c53e774
commit 1db1af1154
2 changed files with 20 additions and 5 deletions

View file

@ -687,9 +687,7 @@ class KeyManagementSettings(LiteLLMBase):
class TeamDefaultSettings(LiteLLMBase):
team_id: str
model_config = ConfigDict(
extra="allow", # allow params not defined here, these fall in litellm.completion(**kwargs)
)
model_config = get_model_config()
class DynamoDBArgs(LiteLLMBase):

View file

@ -1,6 +1,23 @@
from typing import List, Optional, Union
from pydantic import ConfigDict, BaseModel, validator
from pydantic import ConfigDict, BaseModel, validator, VERSION
# Function to get Pydantic version
def is_pydantic_v2() -> int:
return int(VERSION.split(".")[0])
def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict:
# Version-specific configuration
if is_pydantic_v2() >= 2:
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore
else:
from pydantic import Extra
model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore
return model_config
class EmbeddingRequest(BaseModel):
@ -17,4 +34,4 @@ class EmbeddingRequest(BaseModel):
litellm_call_id: Optional[str] = None
litellm_logging_obj: Optional[dict] = None
logger_fn: Optional[str] = None
model_config = ConfigDict(extra="allow")
model_config = get_model_config()