From ad8c3ac2c3ab1bdc2ba81e9f141fcce68d817c25 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Mon, 13 May 2024 10:15:37 -0700 Subject: [PATCH] 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 ============================ ``` --- litellm/proxy/_types.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a2776f465e5..988e92f67ed 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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})