Upgrade to Pydantic v2: Update Validators and Packaging

- Refactored code to use `@model_validator` instead of `@root_validator`
- Updated `Field` usage to align with Pydantic v2 standards
- Modified validation logic to comply with the new API
- Updated `poetry.lock` and dependencies in `pyproject.toml`

Author: xiaobo

nemoguardrails/eval/models.py

@@ -16,7 +16,7 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, Field, model_validator
 
 from nemoguardrails.eval.utils import load_dict_from_path
 from nemoguardrails.logging.explain import LLMCallInfo
@@ -107,7 +107,7 @@ class InteractionSet(BaseModel):
         description="A list of tags that should be associated with the interactions. Useful for filtering when reporting.",
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
     def instantiate_expected_output(cls, values: Any):
         """Creates the right instance of the expected output."""
         type_mapping = {
@@ -147,11 +147,11 @@ class EvalConfig(BaseModel):
         description="The prompts that should be used for the various LLM tasks.",
     )
 
-    @root_validator(pre=False, skip_on_failure=True)
-    def validate_policy_ids(cls, values: Any):
+    @model_validator(mode="after")
+    def validate_policy_ids(cls, values: "EvalConfig") -> "EvalConfig":
         """Validates the policy ids used in the interactions."""
-        policy_ids = {policy.id for policy in values.get("policies")}
-        for interaction_set in values.get("interactions"):
+        policy_ids = {policy.id for policy in values.policies}
+        for interaction_set in values.interactions:
             for expected_output in interaction_set.expected_output:
                 if expected_output.policy not in policy_ids:
                     raise ValueError(
@@ -180,7 +180,7 @@ def from_path(
         else:
             raise ValueError(f"Invalid config path {config_path}.")
 
-        return cls.parse_obj(config_obj)
+        return cls.model_validate(config_obj)
 
 
 class ComplianceCheckLog(BaseModel):
@@ -361,4 +361,4 @@ def from_path(
         else:
             raise ValueError(f"Invalid config path {output_path}.")
 
-        return cls.parse_obj(output_obj)
+        return cls.model_validate(output_obj)

nemoguardrails/library/factchecking/align_score/requirements.txt

@@ -1,5 +1,5 @@
 # The minimal set of requirements for the AlignScore server to run.
-pydantic>=1.10
+pydantic>=2.10.6
 fastapi>=0.109.1
 starlette>=0.36.2
 typer>=0.7.0

nemoguardrails/library/jailbreak_detection/requirements.txt

@@ -1,5 +1,5 @@
 # The minimal set of requirements for the jailbreak detection server to run.
-pydantic>=1.10.9
+pydantic>=2.10.6
 fastapi>=0.103.1
 starlette>=0.27.0
 typer>=0.7.0

nemoguardrails/llm/providers/_langchain_nvidia_ai_endpoints_patch.py

@@ -22,7 +22,7 @@
 from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatResult
 from langchain_nvidia_ai_endpoints import ChatNVIDIA as ChatNVIDIAOriginal
-from pydantic.v1 import Field
+from pydantic import Field
 
 log = logging.getLogger(__name__)
 

nemoguardrails/llm/providers/nemollm.py

@@ -27,7 +27,7 @@
 from langchain.schema import Generation
 from langchain.schema.output import GenerationChunk, LLMResult
 from langchain_core.language_models.llms import BaseLLM
-from pydantic.v1 import root_validator
+from pydantic import model_validator
 
 log = logging.getLogger(__name__)
 
@@ -52,7 +52,7 @@ class NeMoLLM(BaseLLM):
     streaming: bool = False
     check_api_host_version: bool = True
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def check_env_variables(cls, values):
         for field in ["api_host", "api_key", "organization_id"]:
             # If it's an explicit environment variable, we use that

nemoguardrails/llm/providers/trtllm/llm.py

@@ -22,7 +22,7 @@
 
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.llms import BaseLLM
-from pydantic.v1 import Field, root_validator
+from pydantic import Field, model_validator
 
 from nemoguardrails.llm.providers.trtllm.client import TritonClient
 
@@ -61,12 +61,12 @@ class TRTLLM(BaseLLM):
     client: Any
     streaming: Optional[bool] = True
 
-    @root_validator(allow_reuse=True)
+    @model_validator(mode="after")
     @classmethod
-    def validate_environment(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    def validate_environment(cls, values: "TRTLLM") -> "TRTLLM":
         """Validate that python package exists in environment."""
         try:
-            values["client"] = TritonClient(values["server_url"])
+            values.client = TritonClient(values.server_url)
 
         except ImportError as err:
             raise ImportError(

nemoguardrails/rails/llm/config.py

@@ -22,8 +22,7 @@
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import yaml
-from pydantic import BaseModel, ConfigDict, ValidationError, root_validator
-from pydantic.fields import Field
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
 
 from nemoguardrails import utils
 from nemoguardrails.colang import parse_colang_file, parse_flow_elements
@@ -213,7 +212,7 @@ class TaskPrompt(BaseModel):
         description="The maximum number of tokens that can be generated in the chat completion.",
     )
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def check_fields(cls, values):
         if not values.get("content") and not values.get("messages"):
             raise ValidationError("One of `content` or `messages` must be provided.")
@@ -867,16 +866,15 @@ class RailsConfig(BaseModel):
         description="The list of bot messages that should be used for the rails.",
     )
 
-    # NOTE: the Any below is used to get rid of a warning with pydantic 1.10.x;
-    #   The correct typing should be List[Dict, Flow]. To be updated when
-    #   support for pydantic 1.10.x is dropped.
-    flows: List[Union[Dict, Any]] = Field(
+    flows: List[Union[Dict, Flow]] = Field(
         default_factory=list,
         description="The list of flows that should be used for the rails.",
     )
 
     instructions: Optional[List[Instruction]] = Field(
-        default=[Instruction.parse_obj(obj) for obj in _default_config["instructions"]],
+        default=[
+            Instruction.model_validate(obj) for obj in _default_config["instructions"]
+        ],
         description="List of instructions in natural language that the LLM should use.",
     )
 
@@ -981,7 +979,7 @@ class RailsConfig(BaseModel):
         description="Configuration for tracing.",
     )
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def check_prompt_exist_for_self_check_rails(cls, values):
         rails = values.get("rails", {})
 
@@ -1035,7 +1033,7 @@ def check_prompt_exist_for_self_check_rails(cls, values):
 
         return values
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def check_output_parser_exists(cls, values):
         tasks_requiring_output_parser = [
             "self_check_input",
@@ -1068,7 +1066,7 @@ def check_output_parser_exists(cls, values):
                 )
         return values
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def fill_in_default_values_for_v2_x(cls, values):
         instructions = values.get("instructions", {})
         sample_conversation = values.get("sample_conversation")
@@ -1197,7 +1195,7 @@ def parse_object(cls, obj):
                 ):
                     flow_data["elements"] = parse_flow_elements(flow_data["elements"])
 
-        return cls.parse_obj(obj)
+        return cls.model_validate(obj)
 
     @property
     def streaming_supported(self):

nemoguardrails/rails/llm/options.py

@@ -78,7 +78,7 @@
 """
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, root_validator
+from pydantic import BaseModel, Field, model_validator
 
 from nemoguardrails.logging.explain import LLMCallInfo, LLMCallSummary
 
@@ -168,7 +168,7 @@ class GenerationOptions(BaseModel):
         description="Options about what to include in the log. By default, nothing is included. ",
     )
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")
     def check_fields(cls, values):
         # Translate the `rails` generation option from List[str] to dict.
         if "rails" in values and isinstance(values["rails"], list):

nemoguardrails/server/api.py

@@ -25,7 +25,7 @@
 
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field, root_validator, validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from starlette.responses import StreamingResponse
 from starlette.staticfiles import StaticFiles
 
@@ -140,7 +140,7 @@ class RequestBody(BaseModel):
         description="A state object that should be used to continue the interaction.",
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
     def ensure_config_id(cls, data: Any) -> Any:
         if isinstance(data, dict):
             if data.get("config_id") is not None and data.get("config_ids") is not None:
@@ -155,11 +155,15 @@ def ensure_config_id(cls, data: Any) -> Any:
                 )
         return data
 
-    @validator("config_ids", pre=True, always=True)
+    @field_validator("config_ids", mode="before")
     def ensure_config_ids(cls, v, values):
-        if v is None and values.get("config_id") and values.get("config_ids") is None:
+        if (
+            v is None
+            and values.data.get("config_id")
+            and values.data.get("config_ids") is None
+        ):
             # populate config_ids with config_id if only config_id is provided
-            return [values["config_id"]]
+            return [values.data["config_id"]]
         return v
 
 

poetry.lock

@@ -65,7 +65,7 @@ langchain-community = ">=0.0.16,<0.4.0"
 lark = ">=1.1.7"
 nest-asyncio = ">=1.5.6,"
 prompt-toolkit = ">=3.0"
-pydantic = ">=1.10"
+pydantic = ">=2.0"
 pyyaml = ">=6.0"
 rich = ">=13.5.2"
 simpleeval = ">=0.9.13,"