diff --git a/mem0/vector_stores/turbopuffer.py b/mem0/vector_stores/turbopuffer.py index 76281fc96..97109704f 100644 --- a/mem0/vector_stores/turbopuffer.py +++ b/mem0/vector_stores/turbopuffer.py @@ -122,7 +122,14 @@ class TurbopufferDB(VectorStoreBase): dist = row_dict.pop("$dist", None) row_dict.pop("vector", None) - score = 1 - dist if dist is not None else None + if dist is None: + score = None + elif self.distance_metric == "euclidean_squared": + # $dist is unbounded squared-L2 (lower = closer); map to a bounded + # higher-is-better score, mirroring milvus/baidu. Cosine returns 1 - dist. + score = 1.0 / (1.0 + dist) + else: + score = 1 - dist results.append(OutputData( id=row_id, diff --git a/tests/vector_stores/test_turbopuffer.py b/tests/vector_stores/test_turbopuffer.py index e1b9233bb..df0670570 100644 --- a/tests/vector_stores/test_turbopuffer.py +++ b/tests/vector_stores/test_turbopuffer.py @@ -232,6 +232,35 @@ class TestParseOutput: def test_parse_empty_rows(self, db): assert db._parse_output([]) == [] + def test_parse_cosine_score_is_one_minus_dist(self, db): + # Default cosine metric: score = 1 - dist, unchanged by the metric fix. + results = db._parse_output([_make_row("id1", dist=0.25)]) + assert results[0].score == pytest.approx(0.75) + + def test_parse_euclidean_squared_score_is_bounded(self, mock_client): + # euclidean_squared $dist is unbounded (e.g. 4.0). 1 - dist would give -3.0, + # violating the higher-is-better contract; map it to 1/(1+dist) instead. + db = TurbopufferDB( + collection_name="test_ns", + embedding_model_dims=4, + api_key="tpuf_test_key", + region="gcp-us-central1", + distance_metric="euclidean_squared", + ) + results = db._parse_output([_make_row("id1", dist=4.0)]) + assert results[0].score == pytest.approx(0.2) + assert 0.0 <= results[0].score <= 1.0 + + def test_parse_euclidean_squared_preserves_none(self, mock_client): + db = TurbopufferDB( + collection_name="test_ns", + embedding_model_dims=4, + api_key="tpuf_test_key", + region="gcp-us-central1", + distance_metric="euclidean_squared", + ) + assert db._parse_output([_make_row("id1")])[0].score is None + # ── _convert_filters ─────────────────────────────────────────────────