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) <noreply@anthropic.com>
This commit is contained in:
utkarsh240799
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
224fd6c7d8
commit
6d8875c196
@@ -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)
|
||||
Reference in New Issue
Block a user