make extension classes accept vectorizer kwargs

Author: tylerhutcherson

redisvl/extensions/llmcache/semantic.py

@@ -95,12 +95,8 @@ def __init__(
         }
 
         # Use the index name as the key prefix by default
-        if "prefix" in kwargs:
-            prefix = kwargs["prefix"]
-        else:
-            prefix = name
-
-        dtype = kwargs.get("dtype")
+        prefix = kwargs.pop("prefix", name)
+        dtype = kwargs.pop("dtype", None)
 
         # Validate a provided vectorizer or set the default
         if vectorizer:
@@ -111,7 +107,10 @@ def __init__(
                     f"Provided dtype {dtype} does not match vectorizer dtype {vectorizer.dtype}"
                 )
         else:
-            vectorizer_kwargs = {"dtype": dtype} if dtype else {}
+            vectorizer_kwargs = kwargs
+
+            if dtype:
+                vectorizer_kwargs.update(**{"dtype": dtype})
 
             vectorizer = HFTextVectorizer(
                 model="sentence-transformers/all-mpnet-base-v2",

redisvl/extensions/router/semantic.py

@@ -72,7 +72,7 @@ def __init__(
             connection_kwargs (Dict[str, Any]): The connection arguments
                 for the redis client. Defaults to empty {}.
         """
-        dtype = kwargs.get("dtype")
+        dtype = kwargs.pop("dtype", None)
 
         # Validate a provided vectorizer or set the default
         if vectorizer:
@@ -83,8 +83,15 @@ def __init__(
                     f"Provided dtype {dtype} does not match vectorizer dtype {vectorizer.dtype}"
                 )
         else:
-            vectorizer_kwargs = {"dtype": dtype} if dtype else {}
-            vectorizer = HFTextVectorizer(**vectorizer_kwargs)
+            vectorizer_kwargs = kwargs
+
+            if dtype:
+                vectorizer_kwargs.update(**{"dtype": dtype})
+
+            vectorizer = HFTextVectorizer(
+                model="sentence-transformers/all-mpnet-base-v2",
+                **vectorizer_kwargs,
+            )
 
         if routing_config is None:
             routing_config = RoutingConfig()

redisvl/extensions/session_manager/semantic_session.py

@@ -71,7 +71,7 @@ def __init__(
         super().__init__(name, session_tag)
 
         prefix = prefix or name
-        dtype = kwargs.get("dtype")
+        dtype = kwargs.pop("dtype", None)
 
         # Validate a provided vectorizer or set the default
         if vectorizer:
@@ -82,10 +82,13 @@ def __init__(
                     f"Provided dtype {dtype} does not match vectorizer dtype {vectorizer.dtype}"
                 )
         else:
-            vectorizer_kwargs = {"dtype": dtype} if dtype else {}
+            vectorizer_kwargs = kwargs
+
+            if dtype:
+                vectorizer_kwargs.update(**{"dtype": dtype})
 
             vectorizer = HFTextVectorizer(
-                model="sentence-transformers/msmarco-distilbert-cos-v5",
+                model="sentence-transformers/all-mpnet-base-v2",
                 **vectorizer_kwargs,
             )