Compare commits

...

2 Commits

Author SHA1 Message Date
utkarsh240799 6d8875c196 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>
2026-03-28 21:57:40 +05:30
utkarsh240799 224fd6c7d8 fix: normalize vector store scores to similarity (higher = better)
The threshold parameter in Memory.search() was silently broken for 11
vector stores because they returned raw distance scores (lower = better)
while the threshold check assumed similarity scores (higher = better).
This caused threshold filtering to be inverted — good matches got
dropped and bad matches passed through.

Convert all affected vector stores to return similarity scores:
- Cosine distance stores: score = max(0.0, 1.0 - distance)
- L2 distance stores: score = 1.0 / (1.0 + distance)
- Stores computing similarity then discarding it: return similarity directly

Fixes #4453
Related: #3283

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-28 21:57:40 +05:30
20 changed files with 744 additions and 48 deletions
+5 -6
View File
@@ -286,13 +286,12 @@ class AzureMySQL(VectorStoreBase):
for row in results:
vec = np.array(json.loads(row['vector']))
# Cosine similarity
similarity = np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec))
distance = 1 - similarity
scored_results.append((row['id'], distance, row['payload']))
# Cosine similarity (higher = more similar)
similarity = float(np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec)))
scored_results.append((row['id'], similarity, row['payload']))
# Sort by distance and limit
scored_results.sort(key=lambda x: x[1])
# Sort by similarity (descending) and limit
scored_results.sort(key=lambda x: x[1], reverse=True)
scored_results = scored_results[:limit]
return [
+9 -1
View File
@@ -14,7 +14,15 @@ class VectorStoreBase(ABC):
@abstractmethod
def search(self, query, vectors, limit=5, filters=None):
"""Search for similar vectors."""
"""Search for similar vectors.
All implementations must return similarity scores where higher values
indicate greater similarity (range [0, 1] preferred). Implementations
using distance metrics must convert to similarity before returning:
- Cosine distance: score = max(0.0, 1.0 - distance)
- L2 distance: score = 1.0 / (1.0 + distance)
- Inner product: score = value (already higher = better)
"""
pass
@abstractmethod
+5 -6
View File
@@ -247,9 +247,8 @@ class CassandraDB(VectorStoreBase):
vec = np.array(row.vector)
# Cosine similarity
similarity = np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec))
distance = 1 - similarity
# Cosine similarity (higher = more similar)
similarity = float(np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec)))
# Apply filters if provided
if filters:
@@ -261,10 +260,10 @@ class CassandraDB(VectorStoreBase):
except json.JSONDecodeError:
continue
scored_results.append((row.id, distance, row.payload))
scored_results.append((row.id, similarity, row.payload))
# Sort by distance and limit
scored_results.sort(key=lambda x: x[1])
# Sort by similarity (descending) and limit
scored_results.sort(key=lambda x: x[1], reverse=True)
scored_results = scored_results[:limit]
return [
+4 -1
View File
@@ -97,9 +97,12 @@ class ChromaDB(VectorStoreBase):
result = []
for i in range(max_length):
# Convert L2 distance to similarity score (higher = more similar)
raw_distance = distances[i] if isinstance(distances, list) and distances and i < len(distances) else None
score = 1.0 / (1.0 + raw_distance) if raw_distance is not None else None
entry = OutputData(
id=ids[i] if isinstance(ids, list) and ids and i < len(ids) else None,
score=(distances[i] if isinstance(distances, list) and distances and i < len(distances) else None),
score=score,
payload=(metadatas[i] if isinstance(metadatas, list) and metadatas and i < len(metadatas) else None),
)
result.append(entry)
+7 -1
View File
@@ -146,7 +146,13 @@ class FAISS(VectorStoreBase):
payload_copy = payload.copy()
score = float(scores[i])
raw_score = float(scores[i])
# Convert to similarity score (higher = more similar)
if self.distance_strategy.lower() == "euclidean":
score = 1.0 / (1.0 + raw_score)
else:
# inner_product and cosine: higher is already better
score = raw_score
entry = OutputData(
id=vector_id,
score=score,
+11 -5
View File
@@ -128,11 +128,17 @@ class MilvusDB(VectorStoreBase):
memory = []
for value in data:
uid, score, metadata = (
value.get("id"),
value.get("distance"),
value.get("entity", {}).get("metadata"),
)
uid = value.get("id")
raw_distance = value.get("distance")
metadata = value.get("entity", {}).get("metadata")
# Convert distance to similarity score (higher = more similar)
# Milvus returns: L2 = distance (lower=better), COSINE/IP = similarity (higher=better)
if raw_distance is not None and self.metric_type in (MetricType.L2, "L2"):
score = 1.0 / (1.0 + raw_distance)
else:
# COSINE and IP: Milvus already returns similarity (higher = better)
score = raw_distance
memory_obj = OutputData(id=uid, score=score, payload=metadata)
memory.append(memory_obj)
+2 -1
View File
@@ -241,7 +241,8 @@ class PGVector(VectorStoreBase):
)
results = cur.fetchall()
return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
# Convert cosine distance to similarity score (higher = more similar)
return [OutputData(id=str(r[0]), score=max(0.0, 1.0 - float(r[1])), payload=r[2]) for r in results]
def delete(self, vector_id: str) -> None:
"""
+1 -1
View File
@@ -158,7 +158,7 @@ class RedisDB(VectorStoreBase):
return [
MemoryResult(
id=result["memory_id"],
score=float(result["vector_distance"]),
score=max(0.0, 1.0 - float(result["vector_distance"])),
payload={
"hash": result["hash"],
"data": result["memory"],
+4 -1
View File
@@ -80,7 +80,10 @@ class S3Vectors(VectorStoreBase):
except json.JSONDecodeError:
logger.warning(f"Failed to parse metadata for key {v.get('key')}")
payload = {}
results.append(OutputData(id=v.get("key"), score=v.get("distance"), payload=payload))
# Convert distance to similarity score (higher = more similar)
raw_distance = v.get("distance")
score = max(0.0, 1.0 - raw_distance) if raw_distance is not None else None
results.append(OutputData(id=v.get("key"), score=score, payload=payload))
return results
def insert(self, vectors, payloads=None, ids=None):
+2 -1
View File
@@ -135,7 +135,8 @@ class Supabase(VectorStoreBase):
data=vectors, limit=limit, filters=filters, include_metadata=True, include_value=True
)
return [OutputData(id=str(result[0]), score=float(result[1]), payload=result[2]) for result in results]
# Convert cosine distance to similarity score (higher = more similar)
return [OutputData(id=str(result[0]), score=max(0.0, 1.0 - float(result[1])), payload=result[2]) for result in results]
def delete(self, vector_id: str):
"""
+3 -1
View File
@@ -379,7 +379,9 @@ class ValkeyDB(VectorStoreBase):
memory_results = []
for doc in results.docs:
# Extract the score
score = float(doc.vector_score) if hasattr(doc, "vector_score") else None
# Convert cosine distance to similarity score (higher = more similar)
raw_distance = float(doc.vector_score) if hasattr(doc, "vector_score") else None
score = max(0.0, 1.0 - raw_distance) if raw_distance is not None else None
# Create the payload
payload = {
@@ -115,7 +115,7 @@ class GoogleMatchingEngine(VectorStoreBase):
output_data.append(
OutputData(
id=result.get("datapoint").get("datapointId"),
score=result.get("distance"),
score=max(0.0, 1.0 - result.get("distance")) if result.get("distance") is not None else None,
payload=result.get("datapoint").get("metadata"),
)
)
@@ -262,7 +262,9 @@ class GoogleMatchingEngine(VectorStoreBase):
logger.debug("Adding %s: %s", restrict.name, restrict.allow_tokens[0])
payload[restrict.name] = restrict.allow_tokens[0]
output_data = OutputData(id=neighbor.id, score=neighbor.distance, payload=payload)
# Convert distance to similarity score (higher = more similar)
score = max(0.0, 1.0 - neighbor.distance) if neighbor.distance is not None else None
output_data = OutputData(id=neighbor.id, score=score, payload=payload)
results.append(output_data)
logger.debug("Returning %d results", len(results))
@@ -406,7 +408,8 @@ class GoogleMatchingEngine(VectorStoreBase):
if restrict.allow_list:
payload[restrict.namespace] = restrict.allow_list[0]
return OutputData(id=neighbor.datapoint.datapoint_id, score=neighbor.distance, payload=payload)
score = max(0.0, 1.0 - neighbor.distance) if neighbor.distance is not None else None
return OutputData(id=neighbor.datapoint.datapoint_id, score=score, payload=payload)
logger.debug("No results found")
return None
+300
View File
@@ -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)
+365
View File
@@ -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)
+2 -2
View File
@@ -45,7 +45,7 @@ def test_search_vectors(chromadb_instance, mock_chromadb_client):
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == 0.1
assert abs(results[0].score - (1.0 / (1.0 + 0.1))) < 1e-9 # L2 distance 0.1 -> similarity
assert results[0].payload == {"name": "vector1"}
@@ -150,7 +150,7 @@ def test_get_vector(chromadb_instance):
chromadb_instance.collection.get.assert_called_once_with(ids=["id1"])
assert result.id == "id1"
assert result.score == 0.1
assert abs(result.score - (1.0 / (1.0 + 0.1))) < 1e-9 # L2 distance 0.1 -> similarity
assert result.payload == {"name": "vector1"}
+3 -3
View File
@@ -129,7 +129,7 @@ class TestMilvusDB:
# Verify results are parsed correctly
assert len(results) == 1
assert results[0].id == "mem1"
assert results[0].score == 0.8
assert results[0].score == 0.8 # COSINE: Milvus returns similarity directly
def test_search_different_user_ids(self, milvus_db, mock_milvus_client):
"""Test that search works with different user_ids (reproduces reported bug)."""
@@ -228,10 +228,10 @@ class TestMilvusDB:
assert len(parsed) == 2
assert parsed[0].id == "mem1"
assert parsed[0].score == 0.9
assert parsed[0].score == 0.9 # COSINE: Milvus returns similarity directly
assert parsed[0].payload == {"user_id": "alice"}
assert parsed[1].id == "mem2"
assert parsed[1].score == 0.85
assert parsed[1].score == 0.85 # COSINE: Milvus returns similarity directly
def test_update_with_none_vector_fetches_existing(self, milvus_db, mock_milvus_client):
"""Test that update with vector=None fetches the existing vector (fixes #3708)."""
+12 -12
View File
@@ -445,9 +445,9 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 2)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[1].id, self.test_ids[1])
self.assertEqual(results[1].score, 0.2)
self.assertEqual(results[1].score, 0.8)
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 2)
@patch('mem0.vector_stores.pgvector.ConnectionPool')
@@ -494,9 +494,9 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 2)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[1].id, self.test_ids[1])
self.assertEqual(results[1].score, 0.2)
self.assertEqual(results[1].score, 0.8)
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 3)
@patch('mem0.vector_stores.pgvector.ConnectionPool')
@@ -1140,7 +1140,7 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[0].payload["user_id"], "alice")
self.assertEqual(results[0].payload["agent_id"], "agent1")
self.assertEqual(results[0].payload["run_id"], "run1")
@@ -1190,7 +1190,7 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[0].payload["user_id"], "alice")
self.assertEqual(results[0].payload["agent_id"], "agent1")
self.assertEqual(results[0].payload["run_id"], "run1")
@@ -1240,7 +1240,7 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[0].payload["user_id"], "alice")
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 2)
@@ -1288,7 +1288,7 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[0].payload["user_id"], "alice")
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 3)
@@ -1336,9 +1336,9 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 2)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[1].id, self.test_ids[1])
self.assertEqual(results[1].score, 0.2)
self.assertEqual(results[1].score, 0.8)
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 2)
@patch('mem0.vector_stores.pgvector.ConnectionPool')
@@ -1385,9 +1385,9 @@ class TestPGVector(unittest.TestCase):
# Verify results
self.assertEqual(len(results), 2)
self.assertEqual(results[0].id, self.test_ids[0])
self.assertEqual(results[0].score, 0.1)
self.assertEqual(results[0].score, 0.9)
self.assertEqual(results[1].id, self.test_ids[1])
self.assertEqual(results[1].score, 0.2)
self.assertEqual(results[1].score, 0.8)
@patch('mem0.vector_stores.pgvector.PSYCOPG_VERSION', 3)
@patch('mem0.vector_stores.pgvector.ConnectionPool')
+1 -1
View File
@@ -157,7 +157,7 @@ def test_search(mock_boto_client):
mock_boto_client.query_vectors.assert_called_once()
assert len(results) == 1
assert results[0].id == "id1"
assert results[0].score == 0.9
assert abs(results[0].score - 0.1) < 1e-9 # distance 0.9 -> similarity 0.1
def test_get(mock_boto_client):
+1 -1
View File
@@ -75,7 +75,7 @@ def test_search_vectors(supabase_instance, mock_collection):
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == 0.9
assert abs(results[0].score - 0.1) < 1e-9 # cosine distance 0.9 -> similarity 0.1
assert results[0].payload == {"name": "vector1"}
@@ -112,7 +112,7 @@ def test_search_vectors(vector_store, mock_vertex_ai):
assert len(results) == 1
assert results[0].id == "test-id"
assert results[0].score == 0.1
assert results[0].score == 0.9 # distance 0.1 -> similarity 0.9
assert results[0].payload == {"user_id": "test_user"}