Change pydantic root_validator to model_validator

pydantic v1 uses `root_validator` and pydantic v2 uses `model_validator`.

pydantic v2 emits a warning when `root_validator` is used. E.g.:

```
litellm/proxy/_types.py:225
  /Users/abramowi/Code/OpenSource/litellm/litellm/proxy/_types.py:225:
  PydanticDeprecatedSince20: Pydantic V1 style `@root_validator`
  validators are deprecated.
  You should migrate to Pydantic V2 style `@model_validator` validators,
  see the migration guide for more details.
  Deprecated in Pydantic V2.0 to be removed in V3.0.
  See Pydantic V2 Migration Guide at
  https://errors.pydantic.dev/2.7/migration/
    @root_validator(pre=True)
```

This change eliminates those warnings with pydantic v2, while retaining
compatibility with pydantic v1.

pydantic 2.7.1 before

```
$ env -i PATH=$PATH poetry run pytest litellm/tests/test_proxy_server.py
...
litellm/proxy/_types.py:225
  /Users/abramowi/Code/OpenSource/litellm/litellm/proxy/_types.py:225: PydanticDeprecatedSince20: Pydantic V1 style `@root_validator` validators are deprecated. You should migrate to Pydantic V2 style `@model_validator` validators, see the migration guide for more details. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.7/migration/
    @root_validator(pre=True)
...
========================== 10 passed, 2 skipped, 39 warnings in 8.67s ===========================
```

pydantic 2.7.1 after

```
$ env -i PATH=$PATH poetry run pytest litellm/tests/test_proxy_server.py
...
========================== 10 passed, 2 skipped, 27 warnings in 9.85s ===========================
```

pydantic 1.10.5 after

```
$ poetry run pip install 'pydantic<2'
...
Successfully installed pydantic-1.10.15
$ env -i PATH=$PATH poetry run pytest litellm/tests/test_proxy_server.py
...
=========================== 10 passed, 2 skipped, 1 warning in 8.13s ============================
```
This commit is contained in:
Marc Abramowitz 2024-05-13 10:15:37 -07:00
parent 61143c8b45
commit ad8c3ac2c3

View file

@ -6,6 +6,15 @@ from datetime import datetime
import uuid, json, sys, os
from litellm.types.router import UpdateRouterConfig
try:
from pydantic import model_validator # pydantic v2
except ImportError:
from pydantic import root_validator # pydantic v1
def model_validator(mode):
pre = mode == "before"
return root_validator(pre=pre)
def hash_token(token: str):
import hashlib
@ -222,7 +231,7 @@ class LiteLLMPromptInjectionParams(LiteLLMBase):
llm_api_system_prompt: Optional[str] = None
llm_api_fail_call_string: Optional[str] = None
@root_validator(pre=True)
@model_validator(mode="before")
def check_llm_api_params(cls, values):
llm_api_check = values.get("llm_api_check")
if llm_api_check is True:
@ -312,7 +321,7 @@ class ModelInfo(LiteLLMBase):
extra = Extra.allow # Allow extra fields
protected_namespaces = ()
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
if values.get("id") is None:
values.update({"id": str(uuid.uuid4())})
@ -341,7 +350,7 @@ class ModelParams(LiteLLMBase):
class Config:
protected_namespaces = ()
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
if values.get("model_info") is None:
values.update({"model_info": ModelInfo()})
@ -388,7 +397,7 @@ class GenerateKeyResponse(GenerateKeyRequest):
user_id: Optional[str] = None
token_id: Optional[str] = None
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
if values.get("token") is not None:
values.update({"key": values.get("token")})
@ -457,7 +466,7 @@ class UpdateUserRequest(GenerateRequestBase):
user_role: Optional[str] = None
max_budget: Optional[float] = None
@root_validator(pre=True)
@model_validator(mode="before")
def check_user_info(cls, values):
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
@ -477,7 +486,7 @@ class NewEndUserRequest(LiteLLMBase):
None # if no equivalent model in allowed region - default all requests to this model
)
@root_validator(pre=True)
@model_validator(mode="before")
def check_user_info(cls, values):
if values.get("max_budget") is not None and values.get("budget_id") is not None:
raise ValueError("Set either 'max_budget' or 'budget_id', not both.")
@ -490,7 +499,7 @@ class Member(LiteLLMBase):
user_id: Optional[str] = None
user_email: Optional[str] = None
@root_validator(pre=True)
@model_validator(mode="before")
def check_user_info(cls, values):
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
@ -535,7 +544,7 @@ class TeamMemberDeleteRequest(LiteLLMBase):
user_id: Optional[str] = None
user_email: Optional[str] = None
@root_validator(pre=True)
@model_validator(mode="before")
def check_user_info(cls, values):
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
@ -572,7 +581,7 @@ class LiteLLM_TeamTable(TeamBase):
class Config:
protected_namespaces = ()
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
dict_fields = [
"metadata",
@ -867,7 +876,7 @@ class UserAPIKeyAuth(
user_role: Optional[Literal["proxy_admin", "app_owner", "app_user"]] = None
allowed_model_region: Optional[Literal["eu"]] = None
@root_validator(pre=True)
@model_validator(mode="before")
def check_api_key(cls, values):
if values.get("api_key") is not None:
values.update({"token": hash_token(values.get("api_key"))})
@ -894,7 +903,7 @@ class LiteLLM_UserTable(LiteLLMBase):
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
if values.get("spend") is None:
values.update({"spend": 0.0})
@ -915,7 +924,7 @@ class LiteLLM_EndUserTable(LiteLLMBase):
default_model: Optional[str] = None
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
@root_validator(pre=True)
@model_validator(mode="before")
def set_model_info(cls, values):
if values.get("spend") is None:
values.update({"spend": 0.0})