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:
@@ -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 [
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 [
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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"],
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
|
||||
@@ -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)."""
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user