Fix Pydantic Models to support V1 and V2 of Pydantic and fix deprecation warnings (mlflow#14535)

Signed-off-by: Nathan Brake <[email protected]>
Signed-off-by: Nathan Brake <[email protected]>
Co-authored-by: Yair Bonastre <[email protected]>

Author: njbrake

docs/docs/llms/deployments/index.mdx

@@ -1469,7 +1469,7 @@ and a config class that inherits from `mlflow.gateway.base_models.ConfigModel`.
 import os
 from typing import AsyncIterable
 
-from pydantic import validator
+from mlflow.utils.pydantic_utils import field_validator
 from mlflow.gateway.base_models import ConfigModel
 from mlflow.gateway.config import RouteConfig
 from mlflow.gateway.providers import BaseProvider
@@ -1480,7 +1480,7 @@ class MyLLMConfig(ConfigModel):
     # This model defines the configuration for the provider such as API keys
     my_llm_api_key: str
 
-    @validator("my_llm_api_key", pre=True)
+    @field_validator("my_llm_api_key", mode="before")
     def validate_my_llm_api_key(cls, value):
         return os.environ[value.lstrip("$")]
 

examples/gateway/plugin/my-llm/my_llm/config.py

@@ -1,14 +1,13 @@
 import os
 
-from pydantic import validator
-
 from mlflow.gateway.base_models import ConfigModel
+from mlflow.utils.pydantic_utils import field_validator
 
 
 class MyLLMConfig(ConfigModel):
     my_llm_api_key: str
 
-    @validator("my_llm_api_key", pre=True)
+    @field_validator("my_llm_api_key", mode="before")
     def validate_my_llm_api_key(cls, value):
         if value.startswith("$"):
             # This resolves the API key from an environment variable

mlflow/gateway/config.py

@@ -9,7 +9,7 @@
 import pydantic
 import yaml
 from packaging.version import Version
-from pydantic import ConfigDict, Field, ValidationError, root_validator, validator
+from pydantic import ConfigDict, Field, ValidationError
 from pydantic.json import pydantic_encoder
 
 from mlflow.exceptions import MlflowException
@@ -26,7 +26,7 @@
     is_valid_endpoint_name,
     is_valid_mosiacml_chat_model,
 )
-from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER, field_validator, model_validator
 
 _logger = logging.getLogger(__name__)
 
@@ -59,7 +59,7 @@ def values(cls):
 class TogetherAIConfig(ConfigModel):
     togetherai_api_key: str
 
-    @validator("togetherai_api_key", pre=True)
+    @field_validator("togetherai_api_key", mode="before")
     def validate_togetherai_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -73,15 +73,15 @@ class RouteType(str, Enum):
 class CohereConfig(ConfigModel):
     cohere_api_key: str
 
-    @validator("cohere_api_key", pre=True)
+    @field_validator("cohere_api_key", mode="before")
     def validate_cohere_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
 
 class AI21LabsConfig(ConfigModel):
     ai21labs_api_key: str
 
-    @validator("ai21labs_api_key", pre=True)
+    @field_validator("ai21labs_api_key", mode="before")
     def validate_ai21labs_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -90,7 +90,7 @@ class MosaicMLConfig(ConfigModel):
     mosaicml_api_key: str
     mosaicml_api_base: Optional[str] = None
 
-    @validator("mosaicml_api_key", pre=True)
+    @field_validator("mosaicml_api_key", mode="before")
     def validate_mosaicml_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -120,7 +120,7 @@ class OpenAIConfig(ConfigModel):
     openai_deployment_name: Optional[str] = None
     openai_organization: Optional[str] = None
 
-    @validator("openai_api_key", pre=True)
+    @field_validator("openai_api_key", mode="before")
     def validate_openai_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -158,34 +158,24 @@ def _validate_field_compatibility(cls, info: dict[str, Any]):
 
         return info
 
-    if IS_PYDANTIC_V2_OR_NEWER:
-        from pydantic import model_validator as _model_validator
-
-        @_model_validator(mode="before")
-        def validate_field_compatibility(cls, info: dict[str, Any]):
-            return cls._validate_field_compatibility(info)
-
-    else:
-        from pydantic import root_validator as _root_validator
-
-        @_root_validator(pre=False)
-        def validate_field_compatibility(cls, config: dict[str, Any]):
-            return cls._validate_field_compatibility(config)
+    @model_validator(mode="before")
+    def validate_field_compatibility(cls, info: dict[str, Any]):
+        return cls._validate_field_compatibility(info)
 
 
 class AnthropicConfig(ConfigModel):
     anthropic_api_key: str
     anthropic_version: str = "2023-06-01"
 
-    @validator("anthropic_api_key", pre=True)
+    @field_validator("anthropic_api_key", mode="before")
     def validate_anthropic_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
 
 class PaLMConfig(ConfigModel):
     palm_api_key: str
 
-    @validator("palm_api_key", pre=True)
+    @field_validator("palm_api_key", mode="before")
     def validate_palm_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -225,7 +215,7 @@ class AmazonBedrockConfig(ConfigModel):
 class MistralConfig(ConfigModel):
     mistral_api_key: str
 
-    @validator("mistral_api_key", pre=True)
+    @field_validator("mistral_api_key", mode="before")
     def validate_mistral_api_key(cls, value):
         return _resolve_api_key_from_input(value)
 
@@ -284,7 +274,7 @@ class Model(ConfigModel):
     else:
         config: Optional[ConfigModel] = None
 
-    @validator("provider", pre=True)
+    @field_validator("provider", mode="before")
     def validate_provider(cls, value):
         from mlflow.gateway.provider_registry import provider_registry
 
@@ -298,28 +288,26 @@ def validate_provider(cls, value):
         raise MlflowException.invalid_parameter_value(f"The provider '{value}' is not supported.")
 
     @classmethod
-    def _validate_config(cls, info, values):
+    def _validate_config(cls, val, context):
         from mlflow.gateway.provider_registry import provider_registry
 
-        if provider := values.get("provider"):
-            config_type = provider_registry.get(provider).CONFIG_TYPE
-            return config_type(**info)
+        # For Pydantic v2: 'context' is a ValidationInfo object with a 'data' attribute.
+        # For Pydantic v1: 'context' is dict-like 'values'.
+        if IS_PYDANTIC_V2_OR_NEWER:
+            provider = context.data.get("provider")
+        else:
+            provider = context.get("provider") if context else None
 
+        if provider:
+            config_type = provider_registry.get(provider).CONFIG_TYPE
+            return config_type(**val) if isinstance(val, dict) else val
         raise MlflowException.invalid_parameter_value(
             "A provider must be provided for each gateway route."
         )
 
-    if IS_PYDANTIC_V2_OR_NEWER:
-
-        @validator("config", pre=True)
-        def validate_config(cls, info, values):
-            return cls._validate_config(info, values)
-
-    else:
-
-        @validator("config", pre=True)
-        def validate_config(cls, config, values):
-            return cls._validate_config(config, values)
+    @field_validator("config", mode="before")
+    def validate_config(cls, info, values):
+        return cls._validate_config(info, values)
 
 
 class AliasedConfigModel(ConfigModel):
@@ -351,7 +339,7 @@ class RouteConfig(AliasedConfigModel):
     model: Model
     limit: Optional[Limit] = None
 
-    @validator("name")
+    @field_validator("name")
     def validate_endpoint_name(cls, route_name):
         if not is_valid_endpoint_name(route_name):
             raise MlflowException.invalid_parameter_value(
@@ -361,7 +349,7 @@ def validate_endpoint_name(cls, route_name):
             )
         return route_name
 
-    @validator("model", pre=True)
+    @field_validator("model", mode="before")
     def validate_model(cls, model):
         if model:
             model_instance = Model(**model)
@@ -372,10 +360,14 @@ def validate_model(cls, model):
                 )
         return model
 
-    @root_validator(skip_on_failure=True)
+    @model_validator(mode="after", skip_on_failure=True)
     def validate_route_type_and_model_name(cls, values):
-        route_type = values.get("route_type")
-        model = values.get("model")
+        if IS_PYDANTIC_V2_OR_NEWER:
+            route_type = values.route_type
+            model = values.model
+        else:
+            route_type = values.get("route_type")
+            model = values.get("model")
         if (
             model
             and model.provider == "mosaicml"
@@ -394,13 +386,13 @@ def validate_route_type_and_model_name(cls, values):
             )
         return values
 
-    @validator("route_type", pre=True)
+    @field_validator("route_type", mode="before")
     def validate_route_type(cls, value):
         if value in RouteType._value2member_map_:
             return value
         raise MlflowException.invalid_parameter_value(f"The route_type '{value}' is not supported.")
 
-    @validator("limit", pre=True)
+    @field_validator("limit", mode="before")
     def validate_limit(cls, value):
         from limits import parse
 

mlflow/gateway/providers/mlflow.py

@@ -1,19 +1,20 @@
 import time
 
-from pydantic import BaseModel, StrictFloat, StrictStr, ValidationError, validator
+from pydantic import BaseModel, StrictFloat, StrictStr, ValidationError
 
 from mlflow.gateway.config import MlflowModelServingConfig, RouteConfig
 from mlflow.gateway.constants import MLFLOW_SERVING_RESPONSE_KEY
 from mlflow.gateway.exceptions import AIGatewayException
 from mlflow.gateway.providers.base import BaseProvider
 from mlflow.gateway.providers.utils import send_request
 from mlflow.gateway.schemas import chat, completions, embeddings
+from mlflow.utils.pydantic_utils import field_validator
 
 
 class ServingTextResponse(BaseModel):
     predictions: list[StrictStr]
 
-    @validator("predictions", pre=True)
+    @field_validator("predictions", mode="before")
     def extract_choices(cls, predictions):
         if isinstance(predictions, list) and not predictions:
             raise ValueError("The input list is empty")
@@ -35,7 +36,7 @@ def extract_choices(cls, predictions):
 class EmbeddingsResponse(BaseModel):
     predictions: list[list[StrictFloat]]
 
-    @validator("predictions", pre=True)
+    @field_validator("predictions", mode="before")
     def validate_predictions(cls, predictions):
         if isinstance(predictions, list) and not predictions:
             raise ValueError("The input list is empty")

mlflow/types/agent.py

@@ -18,12 +18,7 @@
     Property,
     Schema,
 )
-from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
-
-if IS_PYDANTIC_V2_OR_NEWER:
-    from pydantic import model_validator
-else:
-    from pydantic import root_validator
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER, model_validator
 
 
 class ChatAgentMessage(BaseModel):
@@ -53,28 +48,21 @@ class ChatAgentMessage(BaseModel):
     # TODO make this a pydantic class with subtypes once we have more details on usage
     attachments: Optional[dict[str, str]] = None
 
-    if IS_PYDANTIC_V2_OR_NEWER:
-
-        @model_validator(mode="after")
-        def check_content_and_tool_calls(cls, chat_agent_msg):
-            """
-            Ensure at least one of 'content' or 'tool_calls' is set.
-            """
-            if not chat_agent_msg.content and not chat_agent_msg.tool_calls:
-                raise ValueError("Either 'content' or 'tool_calls' must be provided.")
-            return chat_agent_msg
-    else:
-
-        @root_validator
-        def check_content_and_tool_calls(cls, values):
-            """
-            Ensure at least one of 'content' or 'tool_calls' is set.
-            """
+    @model_validator(mode="after")
+    def check_content_and_tool_calls(cls, values):
+        """
+        Ensure at least one of 'content' or 'tool_calls' is set.
+        """
+        if IS_PYDANTIC_V2_OR_NEWER:
+            content = values.content
+            tool_calls = values.tool_calls
+        else:
             content = values.get("content")
             tool_calls = values.get("tool_calls")
-            if not content and not tool_calls:
-                raise ValueError("Either 'content' or 'tool_calls' must be provided.")
-            return values
+
+        if not content and not tool_calls:
+            raise ValueError("Either 'content' or 'tool_calls' must be provided.")
+        return values
 
 
 class ChatContext(BaseModel):

mlflow/utils/pydantic_utils.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Callable
 
 import pydantic
 from packaging.version import Version
@@ -7,6 +7,38 @@
 IS_PYDANTIC_V2_OR_NEWER = Version(pydantic.VERSION).major >= 2
 
 
+def field_validator(field: str, mode: str = "before"):
+    def decorator(func: Callable) -> Callable:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            from pydantic import field_validator as pydantic_field_validator
+
+            return pydantic_field_validator(field, mode=mode)(func)
+        else:
+            from pydantic import validator as pydantic_field_validator
+
+            return pydantic_field_validator(field, pre=mode == "before")(func)
+
+    return decorator
+
+
+def model_validator(mode: str, skip_on_failure: bool = False):
+    """A wrapper for Pydantic model validator that is compatible with Pydantic v1 and v2.
+    Note that the `skip_on_failure` argument is only available in Pydantic v1.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            from pydantic import model_validator as pydantic_model_validator
+
+            return pydantic_model_validator(mode=mode)(func)
+        else:
+            from pydantic import root_validator
+
+            return root_validator(pre=mode == "before", skip_on_failure=skip_on_failure)(func)
+
+    return decorator
+
+
 def model_dump_compat(pydantic_model: BaseModel, **kwargs: Any) -> dict[str, Any]:
     """
     Dump the Pydantic model to dictionary, in a compatible way for Pydantic v1 and v2.