From 224fd6c7d8e2af8bc570d260c6f9f6b89100d533 Mon Sep 17 00:00:00 2001 From: utkarsh240799 Date: Fri, 20 Mar 2026 18:33:33 +0530 Subject: [PATCH] fix: normalize vector store scores to similarity (higher = better) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- mem0/vector_stores/azure_mysql.py | 11 ++++----- mem0/vector_stores/base.py | 10 +++++++- mem0/vector_stores/cassandra.py | 11 ++++----- mem0/vector_stores/chroma.py | 5 +++- mem0/vector_stores/faiss.py | 8 ++++++- mem0/vector_stores/milvus.py | 16 +++++++++---- mem0/vector_stores/pgvector.py | 3 ++- mem0/vector_stores/redis.py | 2 +- mem0/vector_stores/s3_vectors.py | 5 +++- mem0/vector_stores/supabase.py | 3 ++- mem0/vector_stores/valkey.py | 4 +++- mem0/vector_stores/vertex_ai_vector_search.py | 9 ++++--- tests/vector_stores/test_chroma.py | 4 ++-- tests/vector_stores/test_milvus.py | 6 ++--- tests/vector_stores/test_pgvector.py | 24 +++++++++---------- tests/vector_stores/test_s3_vectors.py | 2 +- tests/vector_stores/test_supabase.py | 2 +- .../test_vertex_ai_vector_search.py | 2 +- 18 files changed, 79 insertions(+), 48 deletions(-) diff --git a/mem0/vector_stores/azure_mysql.py b/mem0/vector_stores/azure_mysql.py index 2d9ab373b..dfe5bd53b 100644 --- a/mem0/vector_stores/azure_mysql.py +++ b/mem0/vector_stores/azure_mysql.py @@ -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 [ diff --git a/mem0/vector_stores/base.py b/mem0/vector_stores/base.py index 3e22499d7..9b476d169 100644 --- a/mem0/vector_stores/base.py +++ b/mem0/vector_stores/base.py @@ -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 diff --git a/mem0/vector_stores/cassandra.py b/mem0/vector_stores/cassandra.py index 24e4fea88..da13aea31 100644 --- a/mem0/vector_stores/cassandra.py +++ b/mem0/vector_stores/cassandra.py @@ -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 [ diff --git a/mem0/vector_stores/chroma.py b/mem0/vector_stores/chroma.py index 63818a5ba..9d04b59f7 100644 --- a/mem0/vector_stores/chroma.py +++ b/mem0/vector_stores/chroma.py @@ -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) diff --git a/mem0/vector_stores/faiss.py b/mem0/vector_stores/faiss.py index 03865c0ac..4192d9a95 100644 --- a/mem0/vector_stores/faiss.py +++ b/mem0/vector_stores/faiss.py @@ -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, diff --git a/mem0/vector_stores/milvus.py b/mem0/vector_stores/milvus.py index 4a0cd7961..0f0caf685 100644 --- a/mem0/vector_stores/milvus.py +++ b/mem0/vector_stores/milvus.py @@ -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) diff --git a/mem0/vector_stores/pgvector.py b/mem0/vector_stores/pgvector.py index e2d020a66..0f776d305 100644 --- a/mem0/vector_stores/pgvector.py +++ b/mem0/vector_stores/pgvector.py @@ -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: """ diff --git a/mem0/vector_stores/redis.py b/mem0/vector_stores/redis.py index 6e2544a0a..0aed7a029 100644 --- a/mem0/vector_stores/redis.py +++ b/mem0/vector_stores/redis.py @@ -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"], diff --git a/mem0/vector_stores/s3_vectors.py b/mem0/vector_stores/s3_vectors.py index f6504c379..fbb3fcbbd 100644 --- a/mem0/vector_stores/s3_vectors.py +++ b/mem0/vector_stores/s3_vectors.py @@ -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): diff --git a/mem0/vector_stores/supabase.py b/mem0/vector_stores/supabase.py index e55a979cb..4d066cc87 100644 --- a/mem0/vector_stores/supabase.py +++ b/mem0/vector_stores/supabase.py @@ -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): """ diff --git a/mem0/vector_stores/valkey.py b/mem0/vector_stores/valkey.py index 273aaba11..7ab764ee5 100644 --- a/mem0/vector_stores/valkey.py +++ b/mem0/vector_stores/valkey.py @@ -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 = { diff --git a/mem0/vector_stores/vertex_ai_vector_search.py b/mem0/vector_stores/vertex_ai_vector_search.py index 9e2a9a5c4..1608ada00 100644 --- a/mem0/vector_stores/vertex_ai_vector_search.py +++ b/mem0/vector_stores/vertex_ai_vector_search.py @@ -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 diff --git a/tests/vector_stores/test_chroma.py b/tests/vector_stores/test_chroma.py index 57c16d4d9..d00017e9f 100644 --- a/tests/vector_stores/test_chroma.py +++ b/tests/vector_stores/test_chroma.py @@ -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"} diff --git a/tests/vector_stores/test_milvus.py b/tests/vector_stores/test_milvus.py index 976056a25..c380f1f64 100644 --- a/tests/vector_stores/test_milvus.py +++ b/tests/vector_stores/test_milvus.py @@ -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).""" diff --git a/tests/vector_stores/test_pgvector.py b/tests/vector_stores/test_pgvector.py index 436c9708c..4c1c1654c 100644 --- a/tests/vector_stores/test_pgvector.py +++ b/tests/vector_stores/test_pgvector.py @@ -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') diff --git a/tests/vector_stores/test_s3_vectors.py b/tests/vector_stores/test_s3_vectors.py index e8141e2f5..012213158 100644 --- a/tests/vector_stores/test_s3_vectors.py +++ b/tests/vector_stores/test_s3_vectors.py @@ -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): diff --git a/tests/vector_stores/test_supabase.py b/tests/vector_stores/test_supabase.py index e051ccf1b..b0d984923 100644 --- a/tests/vector_stores/test_supabase.py +++ b/tests/vector_stores/test_supabase.py @@ -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"} diff --git a/tests/vector_stores/test_vertex_ai_vector_search.py b/tests/vector_stores/test_vertex_ai_vector_search.py index d0d1f4c9b..d3eb82aa8 100644 --- a/tests/vector_stores/test_vertex_ai_vector_search.py +++ b/tests/vector_stores/test_vertex_ai_vector_search.py @@ -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"}