Update MMR tests

Author: g-despot

_includes/code/howto/search.similarity.mmr.py

@@ -1,14 +1,41 @@
-# START MMRNearText
 import weaviate
+import random
+import time
+from weaviate.classes.config import Property, DataType, Configure
 from weaviate.classes.query import Diversity
+from weaviate.collections.classes.data import DataObject
 
 client = weaviate.connect_to_local()
 
-collection = client.collections.get("JeopardyQuestion")
+# Setup: create collection with clustered vectors to demonstrate MMR diversity
+client.collections.delete("MMRDemo")
+col = client.collections.create(
+    name="MMRDemo",
+    properties=[Property(name="question", data_type=DataType.TEXT)],
+    vector_config=Configure.Vectors.self_provided(),
+)
+
+random.seed(42)
+base_vec = [random.uniform(-1, 1) for _ in range(128)]
+for i in range(30):
+    if i < 10:
+        vec = [v + random.uniform(-0.05, 0.05) for v in base_vec]
+    elif i < 20:
+        vec = [-v + random.uniform(-0.05, 0.05) for v in base_vec]
+    else:
+        vec = [random.uniform(-1, 1) for _ in range(128)]
+    col.data.insert(properties={"question": f"Question {i}"}, vector=vec)
+
+time.sleep(2)
+
+# START MMRNearText
+from weaviate.classes.query import Diversity
+
+collection = client.collections.get("MMRDemo")
 
 # Retrieve 20 candidates, then rerank to select 5 diverse results
-response = collection.query.near_text(
-    query="animals in movies",
+response = collection.query.near_vector(
+    near_vector=base_vec,
     limit=20,
     selection=Diversity.MMR(
         limit=5,
@@ -21,13 +48,13 @@
 # END MMRNearText
 
 # Test
-assert response.objects[0].collection == "JeopardyQuestion"
+assert response.objects[0].collection == "MMRDemo"
 assert len(response.objects) == 5
 assert "question" in response.objects[0].properties.keys()
 
 # Verify MMR produces different ordering than standard search
-standard_response = collection.query.near_text(
-    query="animals in movies",
+standard_response = collection.query.near_vector(
+    near_vector=base_vec,
     limit=5,
 )
 standard_questions = [o.properties["question"] for o in standard_response.objects]
@@ -39,7 +66,7 @@
 # START MMRNearVector
 from weaviate.classes.query import Diversity
 
-collection = client.collections.get("JeopardyQuestion")
+collection = client.collections.get("MMRDemo")
 
 # Get a vector to use as query
 sample = collection.query.fetch_objects(limit=1, include_vector=True)
@@ -65,25 +92,25 @@
 # START MMRBalanceExamples
 from weaviate.classes.query import Diversity
 
-collection = client.collections.get("JeopardyQuestion")
+collection = client.collections.get("MMRDemo")
 
 # Pure diversity — maximize difference between results
-response_diverse = collection.query.near_text(
-    query="animals in movies",
+response_diverse = collection.query.near_vector(
+    near_vector=base_vec,
     limit=20,
     selection=Diversity.MMR(limit=5, balance=0.0),
 )
 
 # Balanced — equal weight on relevance and diversity
-response_balanced = collection.query.near_text(
-    query="animals in movies",
+response_balanced = collection.query.near_vector(
+    near_vector=base_vec,
     limit=20,
     selection=Diversity.MMR(limit=5, balance=0.5),
 )
 
 # Pure relevance — equivalent to standard vector search
-response_relevant = collection.query.near_text(
-    query="animals in movies",
+response_relevant = collection.query.near_vector(
+    near_vector=base_vec,
     limit=20,
     selection=Diversity.MMR(limit=5, balance=1.0),
 )
@@ -96,8 +123,9 @@
 
 # Different balance values should produce different result orderings
 diverse_questions = [o.properties["question"] for o in response_diverse.objects]
-balanced_questions = [o.properties["question"] for o in response_balanced.objects]
 relevant_questions = [o.properties["question"] for o in response_relevant.objects]
 assert diverse_questions != relevant_questions, "Pure diversity and pure relevance should differ"
 
+# Cleanup
+client.collections.delete("MMRDemo")
 client.close()

tests/test_python.py

@@ -193,6 +193,17 @@ def test_search(empty_weaviates, script_loc):
     run_py_script(script_loc, custom_replace_pairs=utils.edu_readonly_replacements)
 
 
+@pytest.mark.pyv4
+@pytest.mark.parametrize(
+    "script_loc",
+    [
+        "./_includes/code/howto/search.similarity.mmr.py",
+    ],
+)
+def test_search_mmr(empty_weaviates, script_loc):
+    run_py_script(script_loc)
+
+
 # ========== Starter Guides ==========
 
 @pytest.mark.pyv4