Adds support for new vector data types (#222)

Adds support for bfloat16, float16, and float64 data types.

Users can now specify the desired dtype for their schema when creating
semantic caches, semantic routers, and semantic session managers.

When not specified data type still defaults to float32.

Author: justin-cechmanek

docs/api/schema.rst

@@ -88,7 +88,7 @@ Each field type supports specific attributes that customize its behavior. Below
 
 - `dims`: Dimensionality of the vector.
 - `algorithm`: Indexing algorithm (`flat` or `hnsw`).
-- `datatype`: Float datatype of the vector (`float32` or `float64`).
+- `datatype`: Float datatype of the vector (`bfloat16`, `float16`, `float32`, `float64`).
 - `distance_metric`: Metric for measuring query relevance (`COSINE`, `L2`, `IP`).
 
 **HNSW Vector Field Specific Attributes**:

docs/examples/openai_qna.ipynb

@@ -579,7 +579,7 @@
     "api_key = os.getenv(\"OPENAI_API_KEY\") or getpass.getpass(\"Enter your OpenAI API key: \")\n",
     "oaip = OpenAITextVectorizer(EMBEDDINGS_MODEL, api_config={\"api_key\": api_key})\n",
     "\n",
-    "chunked_data[\"embedding\"] = oaip.embed_many(chunked_data[\"content\"].tolist(), as_buffer=True)\n",
+    "chunked_data[\"embedding\"] = oaip.embed_many(chunked_data[\"content\"].tolist(), as_buffer=True, dtype=\"float32\")\n",
     "chunked_data"
    ]
   },
@@ -1073,7 +1073,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.12.2"
   },
   "orig_nbformat": 4
  },

docs/user_guide/hash_vs_json_05.ipynb

@@ -429,7 +429,7 @@
     "json_data = data.copy()\n",
     "\n",
     "for d in json_data:\n",
-    "    d['user_embedding'] = buffer_to_array(d['user_embedding'], dtype=np.float32)"
+    "    d['user_embedding'] = buffer_to_array(d['user_embedding'], dtype='float32')"
    ]
   },
   {

docs/user_guide/vectorizers_04.ipynb

@@ -356,7 +356,7 @@
    "outputs": [],
    "source": [
     "# You can also create many embeddings at once\n",
-    "embeddings = hf.embed_many(sentences, as_buffer=True)\n"
+    "embeddings = hf.embed_many(sentences, as_buffer=True, dtype=\"float32\")\n"
    ]
   },
   {
@@ -569,7 +569,7 @@
    "source": [
     "from redisvl.utils.vectorize import CustomTextVectorizer\n",
     "\n",
-    "def generate_embeddings(text_input):\n",
+    "def generate_embeddings(text_input, **kwargs):\n",
     "    return [0.101] * 768\n",
     "\n",
     "custom_vectorizer = CustomTextVectorizer(generate_embeddings)\n",

poetry.lock

@@ -26,6 +26,7 @@ redis = ">=5.0.0"
 pydantic = { version = ">=2,<3" }
 tenacity = ">=8.2.2"
 tabulate = { version = ">=0.9.0,<1" }
+ml-dtypes = "^0.4.0"
 openai = { version = ">=1.13.0", optional = true }
 sentence-transformers = { version = ">=2.2.2", optional = true }
 google-cloud-aiplatform = { version = ">=1.26", optional = true }

pyproject.toml

@@ -48,9 +48,9 @@ def non_empty_metadata(cls, v):
             raise TypeError("Metadata must be a dictionary.")
         return v
 
-    def to_dict(self) -> Dict:
+    def to_dict(self, dtype: str) -> Dict:
         data = self.dict(exclude_none=True)
-        data["prompt_vector"] = array_to_buffer(self.prompt_vector)
+        data["prompt_vector"] = array_to_buffer(self.prompt_vector, dtype)
         if self.metadata is not None:
             data["metadata"] = serialize(self.metadata)
         if self.filters is not None:
@@ -112,7 +112,7 @@ def to_dict(self) -> Dict:
 class SemanticCacheIndexSchema(IndexSchema):
 
     @classmethod
-    def from_params(cls, name: str, prefix: str, vector_dims: int):
+    def from_params(cls, name: str, prefix: str, vector_dims: int, dtype: str):
 
         return cls(
             index={"name": name, "prefix": prefix},  # type: ignore
@@ -126,7 +126,7 @@ def from_params(cls, name: str, prefix: str, vector_dims: int):
                     "type": "vector",
                     "attrs": {
                         "dims": vector_dims,
-                        "datatype": "float32",
+                        "datatype": dtype,
                         "distance_metric": "cosine",
                         "algorithm": "flat",
                     },

redisvl/extensions/llmcache/schema.py

@@ -104,7 +104,10 @@ def __init__(
         ]
 
         # Create semantic cache schema and index
-        schema = SemanticCacheIndexSchema.from_params(name, prefix, vectorizer.dims)
+        dtype = kwargs.get("dtype", "float32")
+        schema = SemanticCacheIndexSchema.from_params(
+            name, prefix, vectorizer.dims, dtype
+        )
         schema = self._modify_schema(schema, filterable_fields)
         self._index = SearchIndex(schema=schema)
 
@@ -137,6 +140,7 @@ def __init__(
             self._index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.dims,  # type: ignore
         )
         self._vectorizer = vectorizer
+        self._dtype = self.index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.datatype  # type: ignore[union-attr]
 
     def _modify_schema(
         self,
@@ -286,7 +290,7 @@ def _vectorize_prompt(self, prompt: Optional[str]) -> List[float]:
         if not isinstance(prompt, str):
             raise TypeError("Prompt must be a string.")
 
-        return self._vectorizer.embed(prompt)
+        return self._vectorizer.embed(prompt, dtype=self._dtype)
 
     async def _avectorize_prompt(self, prompt: Optional[str]) -> List[float]:
         """Converts a text prompt to its vector representation using the
@@ -368,6 +372,7 @@ def check(
             num_results=num_results,
             return_score=True,
             filter_expression=filter_expression,
+            dtype=self._dtype,
         )
 
         # Search the cache!
@@ -538,7 +543,7 @@ def store(
         # Load cache entry with TTL
         ttl = ttl or self._ttl
         keys = self._index.load(
-            data=[cache_entry.to_dict()],
+            data=[cache_entry.to_dict(self._dtype)],
             ttl=ttl,
             id_field=ENTRY_ID_FIELD_NAME,
         )
@@ -602,7 +607,7 @@ async def astore(
         # Load cache entry with TTL
         ttl = ttl or self._ttl
         keys = await aindex.load(
-            data=[cache_entry.to_dict()],
+            data=[cache_entry.to_dict(self._dtype)],
             ttl=ttl,
             id_field=ENTRY_ID_FIELD_NAME,
         )

redisvl/extensions/llmcache/semantic.py

@@ -4,7 +4,7 @@
 from pydantic.v1 import BaseModel, Field, validator
 
 from redisvl.extensions.constants import ROUTE_VECTOR_FIELD_NAME
-from redisvl.schema import IndexInfo, IndexSchema
+from redisvl.schema import IndexSchema
 
 
 class Route(BaseModel):
@@ -89,7 +89,7 @@ class SemanticRouterIndexSchema(IndexSchema):
     """Customized index schema for SemanticRouter."""
 
     @classmethod
-    def from_params(cls, name: str, vector_dims: int) -> "SemanticRouterIndexSchema":
+    def from_params(cls, name: str, vector_dims: int, dtype: str):
         """Create an index schema based on router name and vector dimensions.
 
         Args:
@@ -100,7 +100,7 @@ def from_params(cls, name: str, vector_dims: int) -> "SemanticRouterIndexSchema"
             SemanticRouterIndexSchema: The constructed index schema.
         """
         return cls(
-            index=IndexInfo(name=name, prefix=name),
+            index={"name": name, "prefix": name},  # type: ignore
             fields=[  # type: ignore
                 {"name": "route_name", "type": "tag"},
                 {"name": "reference", "type": "text"},
@@ -111,7 +111,7 @@ def from_params(cls, name: str, vector_dims: int) -> "SemanticRouterIndexSchema"
                         "algorithm": "flat",
                         "dims": vector_dims,
                         "distance_metric": "cosine",
-                        "datatype": "float32",
+                        "datatype": dtype,
                     },
                 },
             ],

redisvl/extensions/router/schema.py

@@ -85,26 +85,42 @@ def __init__(
             vectorizer=vectorizer,
             routing_config=routing_config,
         )
-        self._initialize_index(redis_client, redis_url, overwrite, **connection_kwargs)
+        dtype = kwargs.get("dtype", "float32")
+        self._initialize_index(
+            redis_client, redis_url, overwrite, dtype, **connection_kwargs
+        )
 
     def _initialize_index(
         self,
         redis_client: Optional[Redis] = None,
         redis_url: str = "redis://localhost:6379",
         overwrite: bool = False,
+        dtype: str = "float32",
         **connection_kwargs,
     ):
         """Initialize the search index and handle Redis connection."""
-        schema = SemanticRouterIndexSchema.from_params(self.name, self.vectorizer.dims)
+        schema = SemanticRouterIndexSchema.from_params(
+            self.name, self.vectorizer.dims, dtype
+        )
         self._index = SearchIndex(schema=schema)
 
         if redis_client:
             self._index.set_client(redis_client)
         elif redis_url:
             self._index.connect(redis_url=redis_url, **connection_kwargs)
 
+        # Check for existing router index
         existed = self._index.exists()
-        self._index.create(overwrite=overwrite)
+        if not overwrite and existed:
+            existing_index = SearchIndex.from_existing(
+                self.name, redis_client=self._index.client
+            )
+            if existing_index.schema != self._index.schema:
+                raise ValueError(
+                    f"Existing index {self.name} schema does not match the user provided schema for the semantic router. "
+                    "If you wish to overwrite the index schema, set overwrite=True during initialization."
+                )
+        self._index.create(overwrite=overwrite, drop=False)
 
         if not existed or overwrite:
             # write the routes to Redis
@@ -153,7 +169,9 @@ def _add_routes(self, routes: List[Route]):
         for route in routes:
             # embed route references as a single batch
             reference_vectors = self.vectorizer.embed_many(
-                [reference for reference in route.references], as_buffer=True
+                [reference for reference in route.references],
+                as_buffer=True,
+                dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
             )
             # set route references
             for i, reference in enumerate(route.references):
@@ -230,6 +248,7 @@ def _classify_route(
             vector_field_name=ROUTE_VECTOR_FIELD_NAME,
             distance_threshold=distance_threshold,
             return_fields=["route_name"],
+            dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
         )
 
         aggregate_request = self._build_aggregate_request(
@@ -282,6 +301,7 @@ def _classify_multi_route(
             vector_field_name=ROUTE_VECTOR_FIELD_NAME,
             distance_threshold=distance_threshold,
             return_fields=["route_name"],
+            dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
         )
         aggregate_request = self._build_aggregate_request(
             vector_range_query, aggregation_method, max_k