pydantic facelift to v2

Author: tylerhutcherson

redisvl/extensions/llmcache/schema.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from pydantic.v1 import BaseModel, Field, root_validator, validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 
 from redisvl.extensions.constants import (
     CACHE_VECTOR_FIELD_NAME,
@@ -34,22 +34,23 @@ class CacheEntry(BaseModel):
     filters: Optional[Dict[str, Any]] = Field(default=None)
     """Optional filter data stored on the cache entry for customizing retrieval"""
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
     @classmethod
     def generate_id(cls, values):
         # Ensure entry_id is set
         if not values.get("entry_id"):
             values["entry_id"] = hashify(values["prompt"], values.get("filters"))
         return values
 
-    @validator("metadata")
+    @field_validator("metadata")
+    @classmethod
     def non_empty_metadata(cls, v):
         if v is not None and not isinstance(v, dict):
             raise TypeError("Metadata must be a dictionary.")
         return v
 
     def to_dict(self, dtype: str) -> Dict:
-        data = self.dict(exclude_none=True)
+        data = self.model_dump(exclude_none=True)
         data["prompt_vector"] = array_to_buffer(self.prompt_vector, dtype)
         if self.metadata is not None:
             data["metadata"] = serialize(self.metadata)
@@ -79,18 +80,18 @@ class CacheHit(BaseModel):
     filters: Optional[Dict[str, Any]] = Field(default=None)
     """Optional filter data stored on the cache entry for customizing retrieval"""
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
     @classmethod
     def validate_cache_hit(cls, values):
         # Deserialize metadata if necessary
         if "metadata" in values and isinstance(values["metadata"], str):
             values["metadata"] = deserialize(values["metadata"])
 
         # Separate filters from other fields
-        known_fields = set(cls.__fields__.keys())
+        known_fields = set(cls.model_fields.keys())
         filters = {k: v for k, v in values.items() if k not in known_fields}
 
-        # Add filters to values
+        # Add filters to valuesgiy s
         if filters:
             values["filters"] = filters
 
@@ -101,7 +102,7 @@ def validate_cache_hit(cls, values):
         return values
 
     def to_dict(self) -> Dict:
-        data = self.dict(exclude_none=True)
+        data = self.model_dump(exclude_none=True)
         if self.filters:
             data.update(self.filters)
             del data["filters"]

redisvl/extensions/router/schema.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Dict, List, Optional
 
-from pydantic.v1 import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from redisvl.extensions.constants import ROUTE_VECTOR_FIELD_NAME
 from redisvl.schema import IndexSchema
@@ -19,21 +19,24 @@ class Route(BaseModel):
     distance_threshold: float = Field(default=0.5)
     """Distance threshold for matching the route."""
 
-    @validator("name")
+    @field_validator("name")
+    @classmethod
     def name_must_not_be_empty(cls, v):
         if not v or not v.strip():
             raise ValueError("Route name must not be empty")
         return v
 
-    @validator("references")
+    @field_validator("references")
+    @classmethod
     def references_must_not_be_empty(cls, v):
         if not v:
             raise ValueError("References must not be empty")
         if any(not ref.strip() for ref in v):
             raise ValueError("All references must be non-empty strings")
         return v
 
-    @validator("distance_threshold")
+    @field_validator("distance_threshold")
+    @classmethod
     def distance_threshold_must_be_positive(cls, v):
         if v is not None and v <= 0:
             raise ValueError("Route distance threshold must be greater than zero")
@@ -79,7 +82,8 @@ class RoutingConfig(BaseModel):
         description="Global distance threshold is deprecated all distance_thresholds now apply at route level.",
     )
 
-    @validator("max_k")
+    @field_validator("max_k")
+    @classmethod
     def max_k_must_be_positive(cls, v):
         if v <= 0:
             raise ValueError("max_k must be a positive integer")

redisvl/extensions/router/semantic.py

@@ -3,7 +3,7 @@
 
 import redis.commands.search.reducers as reducers
 import yaml
-from pydantic.v1 import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 from redis import Redis
 from redis.commands.search.aggregation import AggregateRequest, AggregateResult, Reducer
 from redis.exceptions import ResponseError
@@ -44,8 +44,7 @@ class SemanticRouter(BaseModel):
 
     _index: SearchIndex = PrivateAttr()
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @deprecated_argument("dtype", "vectorizer")
     def __init__(

redisvl/extensions/session_manager/schema.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional
 
-from pydantic.v1 import BaseModel, Field, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from redisvl.extensions.constants import (
     CONTENT_FIELD_NAME,
@@ -33,11 +33,9 @@ class ChatMessage(BaseModel):
     """An optional identifier for a tool call associated with the message."""
     vector_field: Optional[List[float]] = Field(default=None)
     """The vector representation of the message content."""
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    class Config:
-        arbitrary_types_allowed = True
-
-    @root_validator(pre=True)
+    @model_validator(mode="before")
     @classmethod
     def generate_id(cls, values):
         if TIMESTAMP_FIELD_NAME not in values:
@@ -49,7 +47,7 @@ def generate_id(cls, values):
         return values
 
     def to_dict(self, dtype: Optional[str] = None) -> Dict:
-        data = self.dict(exclude_none=True)
+        data = self.model_dump(exclude_none=True)
 
         # handle optional fields
         if SESSION_VECTOR_FIELD_NAME in data:

redisvl/index/storage.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import Any, Callable, Dict, Iterable, List, Optional
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 from redis import Redis
 from redis.asyncio import Redis as AsyncRedis
 from redis.commands.search.indexDefinition import IndexType

redisvl/schema/fields.py

@@ -6,18 +6,16 @@
 """
 
 from enum import Enum
-from typing import Any, Dict, Optional, Tuple, Type, Union
+from typing import Any, Dict, Literal, Optional, Tuple, Type, Union
 
-from pydantic.v1 import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 from redis.commands.search.field import Field as RedisField
 from redis.commands.search.field import GeoField as RedisGeoField
 from redis.commands.search.field import NumericField as RedisNumericField
 from redis.commands.search.field import TagField as RedisTagField
 from redis.commands.search.field import TextField as RedisTextField
 from redis.commands.search.field import VectorField as RedisVectorField
 
-### Attribute Enums ###
-
 
 class VectorDistanceMetric(str, Enum):
     COSINE = "COSINE"
@@ -97,7 +95,7 @@ class BaseVectorFieldAttributes(BaseModel):
     initial_cap: Optional[int] = None
     """Initial vector capacity in the index affecting memory allocation size of the index"""
 
-    @validator("algorithm", "datatype", "distance_metric", pre=True)
+    @field_validator("algorithm", "datatype", "distance_metric", mode="before")
     @classmethod
     def uppercase_strings(cls, v):
         """Validate that provided values are cast to uppercase"""
@@ -119,9 +117,7 @@ def field_data(self) -> Dict[str, Any]:
 class FlatVectorFieldAttributes(BaseVectorFieldAttributes):
     """FLAT vector field attributes"""
 
-    algorithm: VectorIndexAlgorithm = Field(
-        default=VectorIndexAlgorithm.FLAT, const=True
-    )
+    algorithm: Literal[VectorIndexAlgorithm.FLAT] = VectorIndexAlgorithm.FLAT
     """The indexing algorithm for the vector field"""
     block_size: Optional[int] = None
     """Block size to hold amount of vectors in a contiguous array. This is useful when the index is dynamic with respect to addition and deletion"""
@@ -130,9 +126,7 @@ class FlatVectorFieldAttributes(BaseVectorFieldAttributes):
 class HNSWVectorFieldAttributes(BaseVectorFieldAttributes):
     """HNSW vector field attributes"""
 
-    algorithm: VectorIndexAlgorithm = Field(
-        default=VectorIndexAlgorithm.HNSW, const=True
-    )
+    algorithm: Literal[VectorIndexAlgorithm.HNSW] = VectorIndexAlgorithm.HNSW
     """The indexing algorithm for the vector field"""
     m: int = Field(default=16)
     """Number of max outgoing edges for each graph node in each layer"""
@@ -171,7 +165,7 @@ def as_redis_field(self) -> RedisField:
 class TextField(BaseField):
     """Text field supporting a full text search index"""
 
-    type: str = Field(default="text", const=True)
+    type: Literal["text"] = "text"
     attrs: TextFieldAttributes = Field(default_factory=TextFieldAttributes)
 
     def as_redis_field(self) -> RedisField:
@@ -189,7 +183,7 @@ def as_redis_field(self) -> RedisField:
 class TagField(BaseField):
     """Tag field for simple boolean-style filtering"""
 
-    type: str = Field(default="tag", const=True)
+    type: Literal["tag"] = "tag"
     attrs: TagFieldAttributes = Field(default_factory=TagFieldAttributes)
 
     def as_redis_field(self) -> RedisField:
@@ -206,7 +200,7 @@ def as_redis_field(self) -> RedisField:
 class NumericField(BaseField):
     """Numeric field for numeric range filtering"""
 
-    type: str = Field(default="numeric", const=True)
+    type: Literal["numeric"] = "numeric"
     attrs: NumericFieldAttributes = Field(default_factory=NumericFieldAttributes)
 
     def as_redis_field(self) -> RedisField:
@@ -221,7 +215,7 @@ def as_redis_field(self) -> RedisField:
 class GeoField(BaseField):
     """Geo field with a geo-spatial index for location based search"""
 
-    type: str = Field(default="geo", const=True)
+    type: Literal["geo"] = "geo"
     attrs: GeoFieldAttributes = Field(default_factory=GeoFieldAttributes)
 
     def as_redis_field(self) -> RedisField:
@@ -236,7 +230,7 @@ def as_redis_field(self) -> RedisField:
 class FlatVectorField(BaseField):
     "Vector field with a FLAT index (brute force nearest neighbors search)"
 
-    type: str = Field(default="vector", const=True)
+    type: Literal["vector"] = "vector"
     attrs: FlatVectorFieldAttributes
 
     def as_redis_field(self) -> RedisField:
@@ -251,7 +245,7 @@ def as_redis_field(self) -> RedisField:
 class HNSWVectorField(BaseField):
     """Vector field with an HNSW index (approximate nearest neighbors search)"""
 
-    type: str = Field(default="vector", const=True)
+    type: Literal["vector"] = "vector"
     attrs: HNSWVectorFieldAttributes
 
     def as_redis_field(self) -> RedisField:
@@ -269,20 +263,21 @@ def as_redis_field(self) -> RedisField:
         return RedisVectorField(name, self.attrs.algorithm, field_data, as_name=as_name)
 
 
-class FieldFactory:
-    """Factory class to create fields from client data and kwargs."""
+FIELD_TYPE_MAP = {
+    "tag": TagField,
+    "text": TextField,
+    "numeric": NumericField,
+    "geo": GeoField,
+}
 
-    FIELD_TYPE_MAP = {
-        "tag": TagField,
-        "text": TextField,
-        "numeric": NumericField,
-        "geo": GeoField,
-    }
+VECTOR_FIELD_TYPE_MAP = {
+    "flat": FlatVectorField,
+    "hnsw": HNSWVectorField,
+}
 
-    VECTOR_FIELD_TYPE_MAP = {
-        "flat": FlatVectorField,
-        "hnsw": HNSWVectorField,
-    }
+
+class FieldFactory:
+    """Factory class to create fields from client data and kwargs."""
 
     @classmethod
     def pick_vector_field_type(cls, attrs: Dict[str, Any]) -> Type[BaseField]:
@@ -294,10 +289,10 @@ def pick_vector_field_type(cls, attrs: Dict[str, Any]) -> Type[BaseField]:
             raise ValueError("Must provide dims param for the vector field.")
 
         algorithm = attrs["algorithm"].lower()
-        if algorithm not in cls.VECTOR_FIELD_TYPE_MAP:
+        if algorithm not in VECTOR_FIELD_TYPE_MAP:
             raise ValueError(f"Unknown vector field algorithm: {algorithm}")
 
-        return cls.VECTOR_FIELD_TYPE_MAP[algorithm]  # type: ignore
+        return VECTOR_FIELD_TYPE_MAP[algorithm]  # type: ignore
 
     @classmethod
     def create_field(
@@ -312,8 +307,14 @@ def create_field(
         if type == "vector":
             field_class = cls.pick_vector_field_type(attrs)
         else:
-            if type not in cls.FIELD_TYPE_MAP:
+            if type not in FIELD_TYPE_MAP:
                 raise ValueError(f"Unknown field type: {type}")
-            field_class = cls.FIELD_TYPE_MAP[type]  # type: ignore
+            field_class = FIELD_TYPE_MAP[type]  # type: ignore
 
-        return field_class(name=name, path=path, attrs=attrs)  # type: ignore
+        return field_class.model_validate(
+            {
+                "name": name,
+                "path": path,
+                "attrs": attrs,
+            }
+        )