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>
This commit is contained in:
utkarsh240799
2026-03-20 18:33:33 +05:30
parent 431cba20e9
commit 224fd6c7d8
18 changed files with 79 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
+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"}