diff --git a/mem0/vector_stores/s3_vectors.py b/mem0/vector_stores/s3_vectors.py index da059a409..55dadb878 100644 --- a/mem0/vector_stores/s3_vectors.py +++ b/mem0/vector_stores/s3_vectors.py @@ -69,6 +69,16 @@ class S3Vectors(VectorStoreBase): else: raise + def _distance_to_score(self, raw_distance: Optional[float]) -> Optional[float]: + if raw_distance is None: + return None + # Euclidean distance is unbounded, so 1 - distance would collapse most + # scores to 0. Use a bounded monotonic map instead. Cosine distance is + # in [0, 2], where 1 - distance is already a valid similarity. + if self.distance_metric == "euclidean": + return 1.0 / (1.0 + raw_distance) + return max(0.0, 1.0 - raw_distance) + def _parse_output(self, vectors: List[Dict]) -> List[OutputData]: results = [] for v in vectors: @@ -81,7 +91,7 @@ class S3Vectors(VectorStoreBase): logger.warning(f"Failed to parse metadata for key {v.get('key')}") payload = {} raw_distance = v.get("distance") - score = max(0.0, 1.0 - raw_distance) if raw_distance is not None else None + score = self._distance_to_score(raw_distance) results.append(OutputData(id=v.get("key"), score=score, payload=payload)) return results diff --git a/tests/vector_stores/test_s3_vectors.py b/tests/vector_stores/test_s3_vectors.py index 68b7fa0f6..b72db32f4 100644 --- a/tests/vector_stores/test_s3_vectors.py +++ b/tests/vector_stores/test_s3_vectors.py @@ -186,6 +186,45 @@ def test_search(mock_boto_client): assert results[0].score == pytest.approx(0.1) +def test_search_score_cosine_uses_one_minus_distance(mock_boto_client): + """Cosine metric keeps the 1 - distance similarity mapping.""" + mock_boto_client.query_vectors.return_value = {"vectors": [{"key": "id1", "distance": 0.25, "metadata": {}}]} + store = S3Vectors( + vector_bucket_name=BUCKET_NAME, + collection_name=INDEX_NAME, + embedding_model_dims=EMBEDDING_DIMS, + distance_metric="cosine", + ) + + results = store.search(query="test", vectors=[0.1, 0.2], top_k=1) + + assert results[0].score == pytest.approx(0.75) + + +def test_search_score_euclidean_uses_bounded_map(mock_boto_client): + """Euclidean distance is unbounded; scores must not collapse to 0.""" + mock_boto_client.query_vectors.return_value = { + "vectors": [ + {"key": "near", "distance": 0.5, "metadata": {}}, + {"key": "far", "distance": 4.0, "metadata": {}}, + ] + } + store = S3Vectors( + vector_bucket_name=BUCKET_NAME, + collection_name=INDEX_NAME, + embedding_model_dims=EMBEDDING_DIMS, + distance_metric="euclidean", + ) + + results = store.search(query="test", vectors=[0.1, 0.2], top_k=2) + + # 1 / (1 + d): a distance > 1 would give a negative (clamped 0) score under + # the old cosine-only formula; the bounded map keeps ranking intact. + assert results[0].score == pytest.approx(1.0 / 1.5) + assert results[1].score == pytest.approx(1.0 / 5.0) + assert results[0].score > results[1].score > 0.0 + + def test_get(mock_boto_client): """Test retrieving a vector by ID.""" mock_boto_client.get_vectors.return_value = {