From 6d8875c1962f0d070fa09d26931d4a85c7569b10 Mon Sep 17 00:00:00 2001 From: utkarsh240799 Date: Sat, 28 Mar 2026 19:56:35 +0530 Subject: [PATCH] test: add score normalization integration tests for vector stores Level 1 tests validate each vector store's search() returns similarity scores (higher = better) with correct ordering using known vectors. Level 2 tests validate the full Memory.add() + Memory.search(threshold=...) pipeline end-to-end. Both levels are standalone scripts run via CLI args (e.g., `python test_score_normalization.py faiss`), not collected by pytest. Co-Authored-By: Claude Opus 4.6 (1M context) --- test_e2e_threshold.py | 300 +++++++++++++++++++++++++++++ test_score_normalization.py | 365 ++++++++++++++++++++++++++++++++++++ 2 files changed, 665 insertions(+) create mode 100644 test_e2e_threshold.py create mode 100644 test_score_normalization.py diff --git a/test_e2e_threshold.py b/test_e2e_threshold.py new file mode 100644 index 000000000..e3dbdc1a5 --- /dev/null +++ b/test_e2e_threshold.py @@ -0,0 +1,300 @@ +""" +Level 2: End-to-end test for Memory.search(threshold=...) across vector stores. +Tests the full pipeline: Memory.add() -> Memory.search(threshold=X) -> verify filtering. +Uses Ollama LLM (llama3.2) + Ollama embeddings (nomic-embed-text). +""" +import sys +import os +import shutil +import time +import logging + +logging.basicConfig(level=logging.WARNING) +# Suppress noisy loggers +logging.getLogger("httpx").setLevel(logging.WARNING) +logging.getLogger("httpcore").setLevel(logging.WARNING) + +EMBED_DIMS = 768 # nomic-embed-text + +def get_base_config(vector_store_provider, vector_store_config): + """Build a MemoryConfig dict for the given vector store.""" + return { + "llm": { + "provider": "ollama", + "config": { + "model": "llama3.2", + "ollama_base_url": "http://localhost:11434", + "temperature": 0.0, + }, + }, + "embedder": { + "provider": "ollama", + "config": { + "model": "nomic-embed-text", + "ollama_base_url": "http://localhost:11434", + }, + }, + "vector_store": { + "provider": vector_store_provider, + "config": vector_store_config, + }, + } + + +def run_threshold_test(provider_name, config_dict): + """ + Core e2e test: add memories, search with threshold, verify filtering works. + Returns True if passed, raises on failure. + """ + from mem0 import Memory + + print(f"\n{'='*50}") + print(f"Level 2 E2E: {provider_name}") + print(f"{'='*50}") + + m = Memory.from_config(config_dict) + + user_id = f"test_threshold_{provider_name}_{int(time.time())}" + + # Step 1: Add diverse memories + print(" Adding memories...") + memories_to_add = [ + "I love playing tennis every weekend at the local club", + "My favorite programming language is Python and I use it daily", + "I have a golden retriever named Max who is 3 years old", + "I work as a machine learning engineer at a tech startup", + "My favorite food is sushi, especially salmon nigiri", + ] + for mem_text in memories_to_add: + result = m.add(mem_text, user_id=user_id) + if result.get("results"): + print(f" Added: {result['results'][0].get('memory', mem_text[:40])}") + + # Step 2: Search WITHOUT threshold (baseline) + print("\n Searching without threshold (baseline)...") + results_no_threshold = m.search("What programming language do I use?", user_id=user_id) + baseline_results = results_no_threshold.get("results", []) + print(f" Got {len(baseline_results)} results:") + for r in baseline_results: + print(f" score={r.get('score', 'N/A'):.4f} {r.get('memory', '?')[:60]}") + + if not baseline_results: + raise RuntimeError("No results returned without threshold - something is wrong") + + # Step 3: Search WITH a high threshold - should filter out irrelevant results + scores = [r["score"] for r in baseline_results if r.get("score") is not None] + if not scores: + raise RuntimeError("No scores in results") + + # Verify scores are similarity (higher = better): top result should have highest score + assert scores[0] >= scores[-1], \ + f"Scores should be descending (higher=better), got first={scores[0]}, last={scores[-1]}" + + # Use a threshold between the best and worst score + if len(scores) >= 2: + threshold = (scores[0] + scores[1]) / 2 # between top 2 + else: + threshold = scores[0] * 0.9 + + print(f"\n Searching WITH threshold={threshold:.4f}...") + results_with_threshold = m.search( + "What programming language do I use?", + user_id=user_id, + threshold=threshold, + ) + filtered_results = results_with_threshold.get("results", []) + print(f" Got {len(filtered_results)} results (filtered from {len(baseline_results)}):") + for r in filtered_results: + print(f" score={r.get('score', 'N/A'):.4f} {r.get('memory', '?')[:60]}") + + # Step 4: Verify threshold filtering + # All returned results should have score >= threshold + for r in filtered_results: + s = r.get("score") + if s is not None: + assert s >= threshold, \ + f"Result with score {s:.4f} should not pass threshold {threshold:.4f}" + + # Should have fewer results than baseline (threshold filtered some out) + assert len(filtered_results) <= len(baseline_results), \ + "Threshold filtering should not return MORE results" + + # Step 5: Very high threshold should return very few or no results + print("\n Searching WITH very high threshold=0.99...") + results_high = m.search( + "What programming language do I use?", + user_id=user_id, + threshold=0.99, + ) + high_filtered = results_high.get("results", []) + print(f" Got {len(high_filtered)} results") + for r in high_filtered: + assert r.get("score", 0) >= 0.99, \ + f"Score {r.get('score')} below 0.99 threshold" + + print(f"\n ✓ {provider_name} Level 2 E2E PASSED!") + return True + + +# ============================================================ +# Per-store configs +# ============================================================ + +def test_faiss_e2e(): + path = "/tmp/mem0_e2e_faiss" + shutil.rmtree(path, ignore_errors=True) + config = get_base_config("faiss", { + "collection_name": "e2e_faiss_test", + "path": path, + "distance_strategy": "euclidean", + "embedding_model_dims": EMBED_DIMS, + }) + result = run_threshold_test("FAISS (euclidean)", config) + shutil.rmtree(path, ignore_errors=True) + return result + + +def test_chroma_e2e(): + config = get_base_config("chroma", { + "collection_name": "e2e_chroma_test", + "path": "/tmp/mem0_e2e_chroma", + }) + result = run_threshold_test("ChromaDB", config) + shutil.rmtree("/tmp/mem0_e2e_chroma", ignore_errors=True) + return result + + +def test_pgvector_e2e(): + config = get_base_config("pgvector", { + "collection_name": "e2e_pgvector_test", + "embedding_model_dims": EMBED_DIMS, + "host": "localhost", + "port": 5432, + "user": "mem0", + "password": "mem0test", + "dbname": "mem0_test", + "diskann": False, + "hnsw": True, + }) + return run_threshold_test("PGVector", config) + + +def test_redis_e2e(): + config = get_base_config("redis", { + "collection_name": "e2e_redis_test", + "embedding_model_dims": EMBED_DIMS, + "redis_url": "redis://localhost:6379", + }) + return run_threshold_test("Redis", config) + + +def test_valkey_e2e(): + config = get_base_config("valkey", { + "collection_name": "e2e_valkey_test", + "embedding_model_dims": EMBED_DIMS, + "valkey_url": "valkey://localhost:6380", + }) + return run_threshold_test("Valkey", config) + + +def test_milvus_l2_e2e(): + config = get_base_config("milvus", { + "collection_name": "e2e_milvus_l2_test", + "embedding_model_dims": EMBED_DIMS, + "url": "http://localhost:19530", + "token": "", + "metric_type": "L2", + }) + return run_threshold_test("Milvus (L2)", config) + + +def test_milvus_cosine_e2e(): + config = get_base_config("milvus", { + "collection_name": "e2e_milvus_cos_test", + "embedding_model_dims": EMBED_DIMS, + "url": "http://localhost:19530", + "token": "", + "metric_type": "COSINE", + }) + return run_threshold_test("Milvus (COSINE)", config) + + +def test_cassandra_e2e(): + config = get_base_config("cassandra", { + "collection_name": "e2e_cassandra_test", + "embedding_model_dims": EMBED_DIMS, + "contact_points": ["localhost"], + "port": 9042, + "keyspace": "mem0_e2e", + }) + return run_threshold_test("Cassandra", config) + + +def test_s3_vectors_e2e(): + config = get_base_config("s3_vectors", { + "vector_bucket_name": "mem0-e2e-test", + "collection_name": "e2es3vtest", + "embedding_model_dims": EMBED_DIMS, + "distance_metric": "cosine", + "region_name": "us-east-1", + }) + return run_threshold_test("S3 Vectors", config) + + +def test_supabase_e2e(): + conn_str = os.environ.get("SUPABASE_CONN_STRING", "") + if not conn_str: + raise RuntimeError("SUPABASE_CONN_STRING env var not set") + config = get_base_config("supabase", { + "connection_string": conn_str, + "collection_name": "e2e_supabase_test", + "embedding_model_dims": EMBED_DIMS, + }) + return run_threshold_test("Supabase", config) + + +# ============================================================ +# Runner +# ============================================================ + +if __name__ == "__main__": + test_name = sys.argv[1] if len(sys.argv) > 1 else "faiss" + + tests = { + "faiss": test_faiss_e2e, + "chroma": test_chroma_e2e, + "pgvector": test_pgvector_e2e, + "redis": test_redis_e2e, + "valkey": test_valkey_e2e, + "milvus_l2": test_milvus_l2_e2e, + "milvus_cosine": test_milvus_cosine_e2e, + "cassandra": test_cassandra_e2e, + "s3_vectors": test_s3_vectors_e2e, + "supabase": test_supabase_e2e, + } + + if test_name == "all": + to_run = tests + elif test_name in tests: + to_run = {test_name: tests[test_name]} + else: + print(f"Unknown test: {test_name}") + print(f"Available: {', '.join(tests.keys())}, all") + sys.exit(1) + + passed = 0 + failed = 0 + + for name, test_fn in to_run.items(): + try: + if test_fn(): + passed += 1 + except Exception as e: + print(f"\n ✗ {name} FAILED: {e}") + import traceback + traceback.print_exc() + failed += 1 + + print(f"\n{'='*50}") + print(f"Level 2 E2E Results: {passed} passed, {failed} failed") + sys.exit(1 if failed else 0) diff --git a/test_score_normalization.py b/test_score_normalization.py new file mode 100644 index 000000000..153e219a7 --- /dev/null +++ b/test_score_normalization.py @@ -0,0 +1,365 @@ +""" +End-to-end test for PR #4456: score normalization across vector stores. +Tests that all affected vector stores return similarity scores (higher = better). + +Level 1: Vector store layer - direct insert/search, verify scores +Level 2: Memory pipeline - Memory.add() + Memory.search(threshold=...) +""" +import numpy as np +import sys +import shutil +import os +import logging +import uuid + +logging.basicConfig(level=logging.WARNING) + +DIMS = 128 + +def make_vectors(): + """Create a query vector and 3 document vectors with known similarity ordering.""" + np.random.seed(42) + query = np.random.randn(DIMS).astype(np.float32) + query = query / np.linalg.norm(query) + + close_vec = query + np.random.randn(DIMS).astype(np.float32) * 0.1 + close_vec = close_vec / np.linalg.norm(close_vec) + + mid_vec = query + np.random.randn(DIMS).astype(np.float32) * 0.5 + mid_vec = mid_vec / np.linalg.norm(mid_vec) + + far_vec = np.random.randn(DIMS).astype(np.float32) + far_vec = far_vec / np.linalg.norm(far_vec) + + return query.tolist(), [close_vec.tolist(), mid_vec.tolist(), far_vec.tolist()] + + +def assert_scores_valid(results, test_name): + """Common assertions for all vector store tests.""" + scores = [r.score for r in results] + labels = [r.payload.get("label", "?") for r in results] + + print(f"\n=== {test_name} ===") + for r in results: + print(f" {r.payload.get('label', '?'):>6}: score={r.score:.4f}") + + # 1. All scores must be non-negative (similarity, not raw distance) + assert all(s >= 0 for s in scores), f"All scores must be non-negative, got {scores}" + + # 2. Scores should be <= 1 for normalized similarity + assert all(s <= 1.0 for s in scores), f"Scores should be in [0, 1], got {scores}" + + # 3. Ordering: higher score = more similar + assert scores[0] >= scores[1] >= scores[2], \ + f"Results should be ordered by descending similarity, got {scores}" + + # 4. Closest vector should rank first + assert labels[0] == "close", f"Closest vector should rank first, got {labels[0]}" + + # 5. Threshold filtering test + threshold = scores[1] # mid score + filtered = [r for r in results if r.score >= threshold] + assert len(filtered) >= 2, "Threshold should keep at least close and mid" + + print(" ✓ All assertions passed!") + return True + + +# ============================================================ +# LEVEL 1: Vector Store Layer Tests +# ============================================================ + +def test_faiss_euclidean(): + from mem0.vector_stores.faiss import FAISS + path = "/tmp/faiss_test_norm" + shutil.rmtree(path, ignore_errors=True) + store = FAISS(collection_name="test", path=path, distance_strategy="euclidean", embedding_model_dims=DIMS) + query, docs = make_vectors() + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=["c", "m", "f"]) + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "FAISS (euclidean)") + shutil.rmtree(path, ignore_errors=True) + return result + + +def test_faiss_cosine(): + from mem0.vector_stores.faiss import FAISS + path = "/tmp/faiss_test_cos" + shutil.rmtree(path, ignore_errors=True) + store = FAISS(collection_name="test", path=path, distance_strategy="cosine", embedding_model_dims=DIMS) + query, docs = make_vectors() + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=["c", "m", "f"]) + results = store.search(query="", vectors=query, limit=3) + + # Cosine/IP scores on normalized vectors can be negative for dissimilar vectors + # so we only check ordering and that close > far + print("\n=== FAISS (cosine/IP) ===") + for r in results: + print(f" {r.payload.get('label', '?'):>6}: score={r.score:.4f}") + assert results[0].score >= results[1].score >= results[2].score + assert results[0].payload["label"] == "close" + print(" ✓ All assertions passed!") + return True + + +def test_chroma(): + from mem0.vector_stores.chroma import ChromaDB + store = ChromaDB(collection_name="test_chroma_norm") + query, docs = make_vectors() + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=["c", "m", "f"]) + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "ChromaDB (L2 -> similarity)") + store.delete_col() + return result + + +def test_pgvector(host="localhost", port=5432): + from mem0.vector_stores.pgvector import PGVector + store = PGVector( + collection_name="test_pgvector_norm", + embedding_model_dims=DIMS, + host=host, + port=port, + user="mem0", + password="mem0test", + dbname="mem0_test", + diskann=False, + hnsw=True, + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "PGVector (cosine distance -> similarity)") + # cleanup + store.delete_col() + return result + + +def test_redis(host="localhost", port=6379): + from mem0.vector_stores.redis import RedisDB + from datetime import datetime, timezone + store = RedisDB( + collection_name="test_redis_norm", + embedding_model_dims=DIMS, + redis_url=f"redis://{host}:{port}", + ) + query, docs = make_vectors() + now = datetime.now(timezone.utc).isoformat(timespec="microseconds") + payloads = [ + {"hash": "h1", "data": "close memory", "created_at": now, "user_id": "test", "label": "close"}, + {"hash": "h2", "data": "mid memory", "created_at": now, "user_id": "test", "label": "mid"}, + {"hash": "h3", "data": "far memory", "created_at": now, "user_id": "test", "label": "far"}, + ] + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=payloads, ids=ids) + results = store.search(query="", vectors=query, limit=3, filters={"user_id": "test"}) + + # Redis returns MemoryResult not OutputData, adapt for assertion + print("\n=== Redis (cosine distance -> similarity) ===") + for r in results: + label = r.payload.get("data", "?").split()[0] + print(f" {label:>6}: score={r.score:.4f}") + assert results[0].score >= results[1].score >= results[2].score, \ + f"Descending order expected, got {[r.score for r in results]}" + assert all(r.score >= 0 for r in results), "Scores must be non-negative" + assert all(r.score <= 1.0 for r in results), "Scores must be <= 1.0" + print(" ✓ All assertions passed!") + store.delete_col() + return True + + +def test_valkey(host="localhost", port=6380): + from mem0.vector_stores.valkey import ValkeyDB + from datetime import datetime, timezone + store = ValkeyDB( + collection_name="test_valkey_norm", + embedding_model_dims=DIMS, + valkey_url=f"valkey://{host}:{port}", + ) + query, docs = make_vectors() + now = datetime.now(timezone.utc).isoformat(timespec="microseconds") + payloads = [ + {"hash": "h1", "data": "close memory", "created_at": now, "user_id": "test", "label": "close"}, + {"hash": "h2", "data": "mid memory", "created_at": now, "user_id": "test", "label": "mid"}, + {"hash": "h3", "data": "far memory", "created_at": now, "user_id": "test", "label": "far"}, + ] + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=payloads, ids=ids) + results = store.search(query="", vectors=query, limit=3, filters={"user_id": "test"}) + + print("\n=== Valkey (cosine distance -> similarity) ===") + for r in results: + print(f" score={r.score:.4f}") + assert results[0].score >= results[1].score >= results[2].score, \ + f"Descending order expected, got {[r.score for r in results]}" + assert all(r.score >= 0 for r in results), "Scores must be non-negative" + assert all(r.score <= 1.0 for r in results), "Scores must be <= 1.0" + print(" ✓ All assertions passed!") + store.delete_col() + return True + + +def test_milvus_l2(host="localhost", port=19530): + from mem0.vector_stores.milvus import MilvusDB + store = MilvusDB( + collection_name="test_milvus_l2_norm", + embedding_model_dims=DIMS, + url=f"http://{host}:{port}", + token=None, + db_name="", + metric_type="L2", + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "Milvus (L2 -> similarity)") + store.delete_col() + return result + + +def test_milvus_cosine(host="localhost", port=19530): + from mem0.vector_stores.milvus import MilvusDB + store = MilvusDB( + collection_name="test_milvus_cos_norm", + embedding_model_dims=DIMS, + url=f"http://{host}:{port}", + token=None, + db_name="", + metric_type="COSINE", + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + results = store.search(query="", vectors=query, limit=3) + + # COSINE metric in Milvus already returns similarity (higher = better) + print("\n=== Milvus (COSINE - already similarity) ===") + for r in results: + print(f" {r.payload.get('label', '?'):>6}: score={r.score:.4f}") + assert results[0].score >= results[1].score >= results[2].score + assert results[0].payload["label"] == "close" + print(" ✓ All assertions passed!") + return True + + +def test_cassandra(host="localhost", port=9042): + from mem0.vector_stores.cassandra import CassandraDB + store = CassandraDB( + collection_name="test_cassandra_norm", + embedding_model_dims=DIMS, + contact_points=[host], + port=port, + keyspace="mem0_test", + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + results = store.search(query="", vectors=query, limit=3) + # Cassandra now returns similarity directly (not 1-similarity) + print("\n=== Cassandra (cosine similarity direct) ===") + for r in results: + print(f" {r.payload.get('label', '?'):>6}: score={r.score:.4f}") + assert results[0].score >= results[1].score >= results[2].score + assert results[0].payload["label"] == "close" + # Cosine similarity ranges from [-1, 1]; negative is valid for dissimilar vectors + assert all(r.score >= -1.0 for r in results), "Cosine similarity must be >= -1.0" + assert all(r.score <= 1.0 for r in results), "Cosine similarity must be <= 1.0" + print(" ✓ All assertions passed!") + store.delete_col() + return True + + +def test_s3_vectors(region="us-east-1"): + from mem0.vector_stores.s3_vectors import S3Vectors + store = S3Vectors( + vector_bucket_name="mem0-test-score-norm", + collection_name="tests3vnorm", + embedding_model_dims=DIMS, + distance_metric="cosine", + region_name=region, + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + + # S3 Vectors may need a moment to index + import time + time.sleep(2) + + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "S3 Vectors (cosine distance -> similarity)") + store.delete_col() + return result + + +def test_supabase(): + from mem0.vector_stores.supabase import Supabase + conn_str = os.environ.get("SUPABASE_CONN_STRING", "") + if not conn_str: + raise RuntimeError("SUPABASE_CONN_STRING env var not set") + store = Supabase( + connection_string=conn_str, + collection_name="test_supabase_norm", + embedding_model_dims=DIMS, + ) + query, docs = make_vectors() + ids = [str(uuid.uuid4()) for _ in range(3)] + store.insert(vectors=docs, payloads=[{"label": "close"}, {"label": "mid"}, {"label": "far"}], ids=ids) + results = store.search(query="", vectors=query, limit=3) + result = assert_scores_valid(results, "Supabase (cosine distance -> similarity)") + store.delete_col() + return result + + +# ============================================================ +# Runner +# ============================================================ + +if __name__ == "__main__": + test_name = sys.argv[1] if len(sys.argv) > 1 else "all_local" + + tests = { + "faiss": [test_faiss_euclidean, test_faiss_cosine], + "chroma": [test_chroma], + "pgvector": [test_pgvector], + "redis": [test_redis], + "valkey": [test_valkey], + "milvus": [test_milvus_l2, test_milvus_cosine], + "cassandra": [test_cassandra], + "s3_vectors": [test_s3_vectors], + "supabase": [test_supabase], + } + + # "all_local" runs only faiss + chroma (no Docker) + local_tests = ["faiss", "chroma"] + docker_tests = ["pgvector", "redis", "valkey", "milvus", "cassandra"] + + if test_name == "all_local": + to_run = {k: tests[k] for k in local_tests} + elif test_name == "all": + to_run = tests + elif test_name in tests: + to_run = {test_name: tests[test_name]} + else: + print(f"Unknown test: {test_name}") + print(f"Available: {', '.join(tests.keys())}, all_local, all") + sys.exit(1) + + passed = 0 + failed = 0 + + for name, test_fns in to_run.items(): + for fn in test_fns: + try: + if fn(): + passed += 1 + except Exception as e: + print(f"\n=== {fn.__name__} ===") + print(f" ✗ FAILED: {e}") + failed += 1 + + print(f"\n{'='*40}") + print(f"Results: {passed} passed, {failed} failed") + sys.exit(1 if failed else 0)