Expand dtype changes to SemanticSessionManager

Author: abrookins

redisvl/extensions/session_manager/semantic_session.py

@@ -19,12 +19,13 @@
 from redisvl.index import SearchIndex
 from redisvl.query import FilterQuery, RangeQuery
 from redisvl.query.filter import Tag
-from redisvl.utils.utils import validate_vector_dims
+from redisvl.utils.utils import deprecated_argument, validate_vector_dims
 from redisvl.utils.vectorize import BaseVectorizer, HFTextVectorizer
 
 
 class SemanticSessionManager(BaseSessionManager):
 
+    @deprecated_argument("dtype", "vectorizer")
     def __init__(
         self,
         name: str,
@@ -70,16 +71,30 @@ def __init__(
         super().__init__(name, session_tag)
 
         prefix = prefix or name
+        dtype = kwargs.get("dtype")
 
-        self._vectorizer = vectorizer or HFTextVectorizer(
-            model="sentence-transformers/msmarco-distilbert-cos-v5"
-        )
+        # Validate a provided vectorizer or set the default
+        if vectorizer:
+            if not isinstance(vectorizer, BaseVectorizer):
+                raise TypeError("Must provide a valid redisvl.vectorizer class.")
+            if dtype and vectorizer.dtype != dtype:
+                raise ValueError(
+                    f"Provided dtype {dtype} does not match vectorizer dtype {vectorizer.dtype}"
+                )
+        else:
+            vectorizer_kwargs = {"dtype": dtype} if dtype else {}
+
+            vectorizer = HFTextVectorizer(
+                model="sentence-transformers/msmarco-distilbert-cos-v5",
+                **vectorizer_kwargs,
+            )
+
+        self._vectorizer = vectorizer
 
         self.set_distance_threshold(distance_threshold)
 
-        dtype = kwargs.get("dtype", "float32")
         schema = SemanticSessionIndexSchema.from_params(
-            name, prefix, self._vectorizer.dims, dtype
+            name, prefix, self._vectorizer.dims, vectorizer.dtype
         )
 
         self._index = SearchIndex(schema=schema)
@@ -215,7 +230,7 @@ def get_relevant(
             num_results=top_k,
             return_score=True,
             filter_expression=session_filter,
-            dtype=self._index.schema.fields[SESSION_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
+            dtype=self._vectorizer.dtype,
         )
         messages = self._index.query(query)
 
@@ -341,7 +356,7 @@ def add_messages(
             if TOOL_FIELD_NAME in message:
                 chat_message.tool_call_id = message[TOOL_FIELD_NAME]
 
-            chat_messages.append(chat_message.to_dict(dtype=self._index.schema.fields[SESSION_VECTOR_FIELD_NAME].attrs.datatype))  # type: ignore[union-attr]
+            chat_messages.append(chat_message.to_dict(dtype=self._vectorizer.dtype))
 
         self._index.load(data=chat_messages, id_field=ID_FIELD_NAME)
 

tests/integration/test_llmcache.py

@@ -2,6 +2,7 @@
 import os
 from collections import namedtuple
 from time import sleep, time
+import warnings
 
 import pytest
 from pydantic.v1 import ValidationError
@@ -71,6 +72,13 @@ def cache_with_redis_client(vectorizer, client):
     cache_instance._index.delete(True)  # Clean up index
 
 
+@pytest.fixture(autouse=True)
+def disable_deprecation_warnings():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        yield
+
+
 def test_bad_ttl(cache):
     with pytest.raises(ValueError):
         cache.set_ttl(2.5)
@@ -892,6 +900,7 @@ def test_vectorizer_dtype_mismatch():
             name="test_dtype_mismatch",
             dtype="float32",
             vectorizer=HFTextVectorizer(dtype="float16"),
+            overwrite=True,
         )
 
 
@@ -900,15 +909,18 @@ def test_invalid_vectorizer():
         SemanticCache(
             name="test_invalid_vectorizer",
             vectorizer="invalid_vectorizer",  # type: ignore
+            overwrite=True,
         )
 
 
 def test_passes_through_dtype_to_default_vectorizer():
     # The default is float32, so we should see float64 if we pass it in.
-    cache = SemanticCache(name="test_pass_through_dtype)", dtype="float64")
+    cache = SemanticCache(
+        name="test_pass_through_dtype", dtype="float64", overwrite=True
+    )
     assert cache._vectorizer.dtype == "float64"
 
 
 def test_deprecated_dtype_argument():
     with pytest.warns(DeprecationWarning):
-        SemanticCache(name="test_deprecated_dtype", dtype="float32")
+        SemanticCache(name="test_deprecated_dtype", dtype="float32", overwrite=True)

tests/integration/test_semantic_router.py

@@ -1,13 +1,16 @@
 import os
 import pathlib
+import warnings
 
 import pytest
 from redis.exceptions import ConnectionError
 
 from redisvl.exceptions import RedisModuleVersionError
+from redisvl.extensions.llmcache.semantic import SemanticCache
 from redisvl.extensions.router import SemanticRouter
 from redisvl.extensions.router.schema import Route, RoutingConfig
 from redisvl.redis.connection import compare_versions
+from redisvl.utils.vectorize.text.huggingface import HFTextVectorizer
 
 
 def get_base_path():
@@ -45,6 +48,13 @@ def semantic_router(client, routes):
     router.delete()
 
 
+@pytest.fixture(autouse=True)
+def disable_deprecation_warnings():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        yield
+
+
 def test_initialize_router(semantic_router):
     assert semantic_router.name == "test-router"
     assert len(semantic_router.routes) == 2
@@ -303,3 +313,44 @@ def test_bad_dtype_connecting_to_exiting_router(redis_url, routes):
             dtype="float16",
             redis_url=redis_url,
         )
+
+
+def test_vectorizer_dtype_mismatch(routes):
+    with pytest.raises(ValueError):
+        SemanticRouter(
+            name="test_dtype_mismatch",
+            routes=routes,
+            dtype="float32",
+            vectorizer=HFTextVectorizer(dtype="float16"),
+            overwrite=True,
+        )
+
+
+def test_invalid_vectorizer(routes):
+    with pytest.raises(TypeError):
+        SemanticRouter(
+            name="test_invalid_vectorizer",
+            vectorizer="invalid_vectorizer",  # type: ignore
+            overwrite=True,
+        )
+
+
+def test_passes_through_dtype_to_default_vectorizer(routes):
+    # The default is float32, so we should see float64 if we pass it in.
+    router = SemanticRouter(
+        name="test_pass_through_dtype",
+        routes=routes,
+        dtype="float64",
+        overwrite=True,
+    )
+    assert router.vectorizer.dtype == "float64"
+
+
+def test_deprecated_dtype_argument(routes):
+    with pytest.warns(DeprecationWarning):
+        SemanticRouter(
+            name="test_deprecated_dtype",
+            routes=routes,
+            dtype="float32",
+            overwrite=True,
+        )

tests/integration/test_session_manager.py

@@ -1,4 +1,4 @@
-import os
+import warnings
 
 import pytest
 from redis.exceptions import ConnectionError
@@ -9,6 +9,7 @@
     SemanticSessionManager,
     StandardSessionManager,
 )
+from redisvl.utils.vectorize.text.huggingface import HFTextVectorizer
 
 
 @pytest.fixture
@@ -26,6 +27,13 @@ def semantic_session(app_name, client):
     session.delete()
 
 
+@pytest.fixture(autouse=True)
+def disable_deprecation_warnings():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        yield
+
+
 # test standard session manager
 def test_specify_redis_client(client):
     session = StandardSessionManager(name="test_app", redis_client=client)
@@ -579,3 +587,35 @@ def test_bad_dtype_connecting_to_exiting_session(redis_url):
         bad_type = SemanticSessionManager(
             name="float64 session", dtype="float16", redis_url=redis_url
         )
+
+
+def test_vectorizer_dtype_mismatch():
+    with pytest.raises(ValueError):
+        SemanticSessionManager(
+            name="test_dtype_mismatch",
+            dtype="float32",
+            vectorizer=HFTextVectorizer(dtype="float16"),
+            overwrite=True,
+        )
+
+
+def test_invalid_vectorizer():
+    with pytest.raises(TypeError):
+        SemanticSessionManager(
+            name="test_invalid_vectorizer",
+            vectorizer="invalid_vectorizer",  # type: ignore
+            overwrite=True,
+        )
+
+
+def test_passes_through_dtype_to_default_vectorizer():
+    # The default is float32, so we should see float64 if we pass it in.
+    cache = SemanticSessionManager(
+        name="test_pass_through_dtype", dtype="float64", overwrite=True
+    )
+    assert cache._vectorizer.dtype == "float64"
+
+
+def test_deprecated_dtype_argument():
+    with pytest.warns(DeprecationWarning):
+        SemanticSessionManager(name="float64 session", dtype="float64", overwrite=True)