Centralize HF vectorizers into a fixture

Author: abrookins

redisvl/extensions/cache/llm/schema.py

@@ -114,7 +114,6 @@ class SemanticCacheIndexSchema(IndexSchema):
 
     @classmethod
     def from_params(cls, name: str, prefix: str, vector_dims: int, dtype: str):
-
         return cls(
             index={"name": name, "prefix": prefix},  # type: ignore
             fields=[  # type: ignore

tests/conftest.py

@@ -80,6 +80,16 @@ def hf_vectorizer():
     )
 
 
+@pytest.fixture(scope="session")
+def hf_vectorizer_float16():
+    return HFTextVectorizer(dtype="float16")
+
+
+@pytest.fixture(scope="session")
+def hf_vectorizer_with_model():
+    return HFTextVectorizer("sentence-transformers/all-mpnet-base-v2")
+
+
 @pytest.fixture
 def sample_datetimes():
     return {

tests/integration/test_aggregation.py

@@ -10,12 +10,16 @@
 
 
 @pytest.fixture
-def index(sample_data, redis_url):
+def index(sample_data, redis_url, request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
     index = SearchIndex.from_dict(
         {
             "index": {
                 "name": "user_index",
-                "prefix": "v1",
+                "prefix": f"v1_{worker_id}",
                 "storage_type": "hash",
             },
             "fields": [

tests/integration/test_async_search_index.py

@@ -32,13 +32,30 @@ def async_index(index_schema, async_client):
 
 
 @pytest.fixture
-def async_index_from_dict():
-    return AsyncSearchIndex.from_dict({"index": {"name": "my_index"}, "fields": fields})
+def async_index_from_dict(request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
+    return AsyncSearchIndex.from_dict(
+        {"index": {"name": "my_index", "prefix": f"rvl_{worker_id}"}, "fields": fields}
+    )
 
 
 @pytest.fixture
-def async_index_from_yaml():
-    return AsyncSearchIndex.from_yaml("schemas/test_json_schema.yaml")
+def async_index_from_yaml(request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
+    # Load the schema from YAML
+    schema = IndexSchema.from_yaml("schemas/test_json_schema.yaml")
+
+    # Modify the prefix to include the worker ID
+    schema.index.prefix = f"{schema.index.prefix}_{worker_id}"
+
+    # Create the AsyncSearchIndex with the modified schema
+    return AsyncSearchIndex(schema=schema)
 
 
 def test_search_index_properties(index_schema, async_index):

tests/integration/test_llmcache.py

@@ -16,8 +16,8 @@
 
 
 @pytest.fixture
-def vectorizer():
-    return HFTextVectorizer("sentence-transformers/all-mpnet-base-v2")
+def vectorizer(hf_vectorizer_with_model):
+    return hf_vectorizer_with_model
 
 
 @pytest.fixture
@@ -929,12 +929,12 @@ def test_bad_dtype_connecting_to_existing_cache(redis_url):
         )
 
 
-def test_vectorizer_dtype_mismatch(redis_url):
+def test_vectorizer_dtype_mismatch(redis_url, hf_vectorizer_float16):
     with pytest.raises(ValueError):
         SemanticCache(
             name="test_dtype_mismatch",
             dtype="float32",
-            vectorizer=HFTextVectorizer(dtype="float16"),
+            vectorizer=hf_vectorizer_float16,
             redis_url=redis_url,
             overwrite=True,
         )

tests/integration/test_message_history.py

@@ -591,12 +591,12 @@ def test_bad_dtype_connecting_to_exiting_history(redis_url):
         )
 
 
-def test_vectorizer_dtype_mismatch(redis_url):
+def test_vectorizer_dtype_mismatch(redis_url, hf_vectorizer_float16):
     with pytest.raises(ValueError):
         SemanticMessageHistory(
             name="test_dtype_mismatch",
             dtype="float32",
-            vectorizer=HFTextVectorizer(dtype="float16"),
+            vectorizer=hf_vectorizer_float16,
             redis_url=redis_url,
             overwrite=True,
         )

tests/integration/test_search_index.py

@@ -42,13 +42,30 @@ def index(index_schema, client):
 
 
 @pytest.fixture
-def index_from_dict():
-    return SearchIndex.from_dict({"index": {"name": "my_index"}, "fields": fields})
+def index_from_dict(request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
+    return SearchIndex.from_dict(
+        {"index": {"name": "my_index", "prefix": f"rvl_{worker_id}"}, "fields": fields}
+    )
 
 
 @pytest.fixture
-def index_from_yaml():
-    return SearchIndex.from_yaml("schemas/test_json_schema.yaml")
+def index_from_yaml(request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
+    # Load the schema from YAML
+    schema = IndexSchema.from_yaml("schemas/test_json_schema.yaml")
+
+    # Modify the prefix to include the worker ID
+    schema.index.prefix = f"{schema.index.prefix}_{worker_id}"
+
+    # Create the SearchIndex with the modified schema
+    return SearchIndex(schema=schema)
 
 
 def test_search_index_properties(index_schema, index):

tests/integration/test_search_results.py

@@ -16,7 +16,11 @@ def filter_query():
 
 
 @pytest.fixture
-def index(sample_data, redis_url):
+def index(sample_data, redis_url, request):
+    # In xdist, the config has "workerid" in workerinput
+    workerinput = getattr(request.config, "workerinput", {})
+    worker_id = workerinput.get("workerid", "master")
+
     fields_spec = [
         {"name": "credit_score", "type": "tag"},
         {"name": "user", "type": "tag"},
@@ -37,7 +41,7 @@ def index(sample_data, redis_url):
     json_schema = {
         "index": {
             "name": "user_index_json",
-            "prefix": "users_json",
+            "prefix": f"users_json_{worker_id}",
             "storage_type": "json",
         },
         "fields": fields_spec,

tests/integration/test_semantic_router.py

@@ -325,13 +325,13 @@ def test_bad_dtype_connecting_to_exiting_router(redis_url, routes):
         )
 
 
-def test_vectorizer_dtype_mismatch(routes, redis_url):
+def test_vectorizer_dtype_mismatch(routes, redis_url, hf_vectorizer_float16):
     with pytest.raises(ValueError):
         SemanticRouter(
             name="test_dtype_mismatch",
             routes=routes,
             dtype="float32",
-            vectorizer=HFTextVectorizer(dtype="float16"),
+            vectorizer=hf_vectorizer_float16,
             redis_url=redis_url,
             overwrite=True,
         )