refactor vectorizers to work with new pydantic base model

Author: tylerhutcherson

redisvl/extensions/llmcache/semantic.py

@@ -125,7 +125,7 @@ def __init__(
 
         # Create semantic cache schema and index
         schema = SemanticCacheIndexSchema.from_params(
-            name, prefix, vectorizer.dims, vectorizer.dtype
+            name, prefix, vectorizer.dims, vectorizer.dtype  # type: ignore
         )
         schema = self._modify_schema(schema, filterable_fields)
         self._index = SearchIndex(schema=schema)

redisvl/extensions/router.py

@@ -0,0 +1,17 @@
+from pydantic import root_validator
+
+
+class SemanticRouter(BaseModel):
+    # existing fields...
+    vectorizer: Optional[HFTextVectorizer] = None
+    dtype: str = "float32"
+
+    @root_validator
+    def check_vectorizer_dtype(cls, values):
+        router_dtype = values.get("dtype")
+        vectorizer = values.get("vectorizer")
+        if vectorizer is not None and vectorizer.dtype != router_dtype:
+            raise ValueError(
+                f"Mismatched vectorizer dtype: {vectorizer.dtype} does not match router dtype: {router_dtype}"
+            )
+        return values

redisvl/extensions/router/semantic.py

@@ -108,7 +108,7 @@ def _initialize_index(
     ):
         """Initialize the search index and handle Redis connection."""
         schema = SemanticRouterIndexSchema.from_params(
-            self.name, self.vectorizer.dims, self.vectorizer.dtype
+            self.name, self.vectorizer.dims, self.vectorizer.dtype  # type: ignore
         )
         self._index = SearchIndex(schema=schema)
 

redisvl/extensions/session_manager/semantic_session.py

@@ -94,7 +94,7 @@ def __init__(
         self.set_distance_threshold(distance_threshold)
 
         schema = SemanticSessionIndexSchema.from_params(
-            name, prefix, self._vectorizer.dims, vectorizer.dtype
+            name, prefix, vectorizer.dims, vectorizer.dtype  # type: ignore
         )
 
         self._index = SearchIndex(schema=schema)

redisvl/utils/vectorize/__init__.py

@@ -27,16 +27,18 @@ def vectorizer_from_dict(vectorizer: dict) -> BaseVectorizer:
     vectorizer_type = Vectorizers(vectorizer["type"])
     model = vectorizer["model"]
     if vectorizer_type == Vectorizers.cohere:
-        return CohereTextVectorizer(model)
+        return CohereTextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.openai:
-        return OpenAITextVectorizer(model)
+        return OpenAITextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.azure_openai:
-        return AzureOpenAITextVectorizer(model)
+        return AzureOpenAITextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.hf:
-        return HFTextVectorizer(model)
+        return HFTextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.mistral:
-        return MistralAITextVectorizer(model)
+        return MistralAITextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.vertexai:
-        return VertexAITextVectorizer(model)
+        return VertexAITextVectorizer(model=model)
     elif vectorizer_type == Vectorizers.voyageai:
-        return VoyageAITextVectorizer(model)
+        return VoyageAITextVectorizer(model=model)
+    else:
+        raise ValueError(f"Unsupported vectorizer type: {vectorizer_type}")

redisvl/utils/vectorize/base.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from typing import Callable, List, Optional
 
-from pydantic.v1 import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from redisvl.redis.utils import array_to_buffer
 from redisvl.schema.fields import VectorDataType
@@ -19,16 +19,19 @@ class Vectorizers(Enum):
 
 
 class BaseVectorizer(BaseModel, ABC):
+    """Base vectorizer interface."""
+
     model: str
-    dims: int
-    dtype: str = Field(default="float32")
+    dtype: str = "float32"
+    dims: Optional[int] = None
 
     @property
     def type(self) -> str:
         return "base"
 
-    @validator("dtype")
-    def check_dtype(dtype):
+    @field_validator("dtype")
+    @classmethod
+    def check_dtype(cls, dtype):
         try:
             VectorDataType(dtype.upper())
         except ValueError:
@@ -37,7 +40,7 @@ def check_dtype(dtype):
             )
         return dtype
 
-    @validator("dims")
+    @field_validator("dims")
     @classmethod
     def check_dims(cls, value):
         """Ensures the dims are a positive integer."""

redisvl/utils/vectorize/text/azureopenai.py

@@ -1,7 +1,7 @@
 import os
 from typing import Any, Callable, Dict, List, Optional
 
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 from tenacity.retry import retry_if_not_exception_type
 
@@ -57,6 +57,7 @@ def __init__(
         model: str = "text-embedding-ada-002",
         api_config: Optional[Dict] = None,
         dtype: str = "float32",
+        **kwargs,
     ):
         """Initialize the AzureOpenAI vectorizer.
 
@@ -76,10 +77,13 @@ def __init__(
             ValueError: If the AzureOpenAI API key, version, or endpoint are not provided.
             ValueError: If an invalid dtype is provided.
         """
-        self._initialize_clients(api_config)
-        super().__init__(model=model, dims=self._set_model_dims(model), dtype=dtype)
+        super().__init__(model=model, dtype=dtype)
+        # Init client
+        self._initialize_clients(api_config, **kwargs)
+        # Set model dimensions
+        self.dims = self._set_model_dims()
 
-    def _initialize_clients(self, api_config: Optional[Dict]):
+    def _initialize_clients(self, api_config: Optional[Dict], **kwargs):
         """
         Setup the OpenAI clients using the provided API key or an
         environment variable.
@@ -141,21 +145,19 @@ def _initialize_clients(self, api_config: Optional[Dict]):
             api_version=api_version,
             azure_endpoint=azure_endpoint,
             **api_config,
+            **kwargs,
         )
         self._aclient = AsyncAzureOpenAI(
             api_key=api_key,
             api_version=api_version,
             azure_endpoint=azure_endpoint,
             **api_config,
+            **kwargs,
         )
 
-    def _set_model_dims(self, model) -> int:
+    def _set_model_dims(self) -> int:
         try:
-            embedding = (
-                self._client.embeddings.create(input=["dimension test"], model=model)
-                .data[0]
-                .embedding
-            )
+            embedding = self.embed("dimension check")
         except (KeyError, IndexError) as ke:
             raise ValueError(f"Unexpected response from the AzureOpenAI API: {str(ke)}")
         except Exception as e:  # pylint: disable=broad-except

redisvl/utils/vectorize/text/bedrock.py

@@ -2,7 +2,7 @@
 import os
 from typing import Any, Callable, Dict, List, Optional
 
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 from tenacity.retry import retry_if_not_exception_type
 
@@ -51,6 +51,7 @@ def __init__(
         model: str = "amazon.titan-embed-text-v2:0",
         api_config: Optional[Dict[str, str]] = None,
         dtype: str = "float32",
+        **kwargs,
     ) -> None:
         """Initialize the AWS Bedrock Vectorizer.
 
@@ -68,6 +69,17 @@ def __init__(
             ImportError: If boto3 is not installed.
             ValueError: If an invalid dtype is provided.
         """
+        super().__init__(model=model, dtype=dtype)
+        # Init client
+        self._initialize_client(api_config, **kwargs)
+        # Set model dimensions after init
+        self.dims = self._set_model_dims()
+
+    def _initialize_client(self, api_config: Optional[Dict], **kwargs):
+        """
+        Setup the Bedrock client using the provided API keys or
+        environment variables.
+        """
         try:
             import boto3  # type: ignore
         except ImportError:
@@ -98,21 +110,18 @@ def __init__(
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
             region_name=aws_region,
+            **kwargs,
         )
 
-        super().__init__(model=model, dims=self._set_model_dims(model), dtype=dtype)
-
-    def _set_model_dims(self, model: str) -> int:
-        """Initialize model and determine embedding dimensions."""
+    def _set_model_dims(self) -> int:
         try:
-            response = self._client.invoke_model(
-                modelId=model, body=json.dumps({"inputText": "dimension test"})
-            )
-            response_body = json.loads(response["body"].read())
-            embedding = response_body["embedding"]
-            return len(embedding)
-        except Exception as e:
-            raise ValueError(f"Error initializing Bedrock model: {str(e)}")
+            embedding = self.embed("dimension check")
+        except (KeyError, IndexError) as ke:
+            raise ValueError(f"Unexpected response from the OpenAI API: {str(ke)}")
+        except Exception as e:  # pylint: disable=broad-except
+            # fall back (TODO get more specific)
+            raise ValueError(f"Error setting embedding model dimensions: {str(e)}")
+        return len(embedding)
 
     @retry(
         wait=wait_random_exponential(min=1, max=60),

redisvl/utils/vectorize/text/cohere.py

@@ -1,7 +1,7 @@
 import os
 from typing import Any, Callable, Dict, List, Optional
 
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 from tenacity.retry import retry_if_not_exception_type
 
@@ -52,6 +52,7 @@ def __init__(
         model: str = "embed-english-v3.0",
         api_config: Optional[Dict] = None,
         dtype: str = "float32",
+        **kwargs,
     ):
         """Initialize the Cohere vectorizer.
 
@@ -70,24 +71,29 @@ def __init__(
             ValueError: If the API key is not provided.
             ValueError: If an invalid dtype is provided.
         """
-        self._initialize_client(api_config)
-        super().__init__(model=model, dims=self._set_model_dims(model), dtype=dtype)
+        super().__init__(model=model, dtype=dtype)
+        # Init client
+        self._initialize_client(api_config, **kwargs)
+        # Set model dimensions after init
+        self.dims = self._set_model_dims()
 
-    def _initialize_client(self, api_config: Optional[Dict]):
+    def _initialize_client(self, api_config: Optional[Dict], **kwargs):
         """
         Setup the Cohere clients using the provided API key or an
         environment variable.
         """
+        if api_config is None:
+            api_config = {}
+
         # Dynamic import of the cohere module
         try:
-            from cohere import AsyncClient, Client
+            from cohere import Client
         except ImportError:
             raise ImportError(
                 "Cohere vectorizer requires the cohere library. \
                     Please install with `pip install cohere`"
             )
 
-        # Fetch the API key from api_config or environment variable
         api_key = (
             api_config.get("api_key") if api_config else os.getenv("COHERE_API_KEY")
         )
@@ -96,15 +102,11 @@ def _initialize_client(self, api_config: Optional[Dict]):
                 "Cohere API key is required. "
                 "Provide it in api_config or set the COHERE_API_KEY environment variable."
             )
-        self._client = Client(api_key=api_key, client_name="redisvl")
+        self._client = Client(api_key=api_key, client_name="redisvl", **kwargs)
 
-    def _set_model_dims(self, model) -> int:
+    def _set_model_dims(self) -> int:
         try:
-            embedding = self._client.embed(
-                texts=["dimension test"],
-                model=model,
-                input_type="search_document",
-            ).embeddings[0]
+            embedding = self.embed("dimension check", input_type="search_document")
         except (KeyError, IndexError) as ke:
             raise ValueError(f"Unexpected response from the Cohere API: {str(ke)}")
         except Exception as e:  # pylint: disable=broad-except

redisvl/utils/vectorize/text/custom.py

@@ -1,6 +1,6 @@
 from typing import Any, Callable, List, Optional
 
-from pydantic.v1 import PrivateAttr
+from pydantic import PrivateAttr
 
 from redisvl.utils.utils import deprecated_argument
 from redisvl.utils.vectorize.base import BaseVectorizer
@@ -114,17 +114,16 @@ def __init__(
         Raises:
             ValueError: if embedding validation fails.
         """
+        super().__init__(model=self.type, dtype=dtype)
+
         # Store user-provided callables
         self._embed = embed
         self._embed_many = embed_many
         self._aembed = aembed
         self._aembed_many = aembed_many
 
-        # Manually validate sync methods to discover dimension
-        dims = self._validate_sync_callables()
-
-        # Initialize the base class now that we know the dimension
-        super().__init__(model=self.type, dims=dims, dtype=dtype)
+        # Set dims
+        self.dims = self._validate_sync_callables()
 
     @property
     def type(self) -> str: