Revert litellm/types/router.py changes

These don't work in pydantic v1.
This commit is contained in:
Marc Abramowitz
2024-05-13 10:34:05 -07:00
parent abd90e6296
commit eacafa7cd7
+29 -10
View File
@@ -1,6 +1,6 @@
from typing import List, Optional, Union, Dict, Tuple, Literal
import httpx
from pydantic import ConfigDict, BaseModel, validator, Field
from pydantic import BaseModel, validator, Field
from .completion import CompletionRequest
from .embedding import EmbeddingRequest
import uuid, enum
@@ -11,7 +11,9 @@ class ModelConfig(BaseModel):
litellm_params: Union[CompletionRequest, EmbeddingRequest]
tpm: int
rpm: int
model_config = ConfigDict(protected_namespaces=())
class Config:
protected_namespaces = ()
class RouterConfig(BaseModel):
@@ -41,7 +43,9 @@ class RouterConfig(BaseModel):
"usage-based-routing",
"latency-based-routing",
] = "simple-shuffle"
model_config = ConfigDict(protected_namespaces=())
class Config:
protected_namespaces = ()
class UpdateRouterConfig(BaseModel):
@@ -60,13 +64,15 @@ class UpdateRouterConfig(BaseModel):
retry_after: Optional[float] = None
fallbacks: Optional[List[dict]] = None
context_window_fallbacks: Optional[List[dict]] = None
model_config = ConfigDict(protected_namespaces=())
class Config:
protected_namespaces = ()
class ModelInfo(BaseModel):
id: Optional[
str
] = None # Allow id to be optional on input, but it will always be present as a str in the model instance
] # Allow id to be optional on input, but it will always be present as a str in the model instance
db_model: bool = (
False # used for proxy - to separate models which are stored in the db vs. config.
)
@@ -77,7 +83,9 @@ class ModelInfo(BaseModel):
elif isinstance(id, int):
id = str(id)
super().__init__(id=id, **params)
model_config = ConfigDict(extra="allow")
class Config:
extra = "allow"
def __contains__(self, key):
# Define custom behavior for the 'in' operator
@@ -171,7 +179,10 @@ class GenericLiteLLMParams(BaseModel):
if max_retries is not None and isinstance(max_retries, str):
max_retries = int(max_retries) # cast to int
super().__init__(max_retries=max_retries, **args, **params)
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
class Config:
extra = "allow"
arbitrary_types_allowed = True
def __contains__(self, key):
# Define custom behavior for the 'in' operator
@@ -229,7 +240,10 @@ class LiteLLM_Params(GenericLiteLLMParams):
if max_retries is not None and isinstance(max_retries, str):
max_retries = int(max_retries) # cast to int
super().__init__(max_retries=max_retries, **args, **params)
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
class Config:
extra = "allow"
arbitrary_types_allowed = True
def __contains__(self, key):
# Define custom behavior for the 'in' operator
@@ -258,7 +272,9 @@ class updateDeployment(BaseModel):
model_name: Optional[str] = None
litellm_params: Optional[updateLiteLLMParams] = None
model_info: Optional[ModelInfo] = None
model_config = ConfigDict(protected_namespaces=())
class Config:
protected_namespaces = ()
class Deployment(BaseModel):
@@ -290,7 +306,10 @@ class Deployment(BaseModel):
except Exception as e:
# if using pydantic v1
return self.dict(**kwargs)
model_config = ConfigDict(extra="allow", protected_namespaces=())
class Config:
extra = "allow"
protected_namespaces = ()
def __contains__(self, key):
# Define custom behavior for the 'in' operator