From 8a2cf65ec55c9b4e9e370a4776c2a6203247f027 Mon Sep 17 00:00:00 2001
From: csurong <csurong1@gmail.com>
Date: Sun, 2 Aug 2026 10:43:56 +0800
Subject: [PATCH] fix(matrixone): apply score threshold to vector search
 (#39784)

Co-authored-by: caosurong <surong.cao@thinkingdata.cn>
---
 .../src/dify_vdb_matrixone/matrixone_vector.py | 17 ++++++++++-------
 .../tests/unit_tests/test_matrixone_vector.py  | 18 ++++++++++++------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/api/providers/vdb/vdb-matrixone/src/dify_vdb_matrixone/matrixone_vector.py b/api/providers/vdb/vdb-matrixone/src/dify_vdb_matrixone/matrixone_vector.py
index 6039baca9aa..7a51776ad26 100644
--- a/api/providers/vdb/vdb-matrixone/src/dify_vdb_matrixone/matrixone_vector.py
+++ b/api/providers/vdb/vdb-matrixone/src/dify_vdb_matrixone/matrixone_vector.py
@@ -166,6 +166,7 @@ class MatrixoneVector(BaseVector):
         filter = None
         if document_ids_filter:
             filter = {"document_id": {"$in": document_ids_filter}}
+        score_threshold = float(kwargs.get("score_threshold") or 0.0)
 
         results = self.client.query(
             query_vector=query_vector,
@@ -174,15 +175,17 @@ class MatrixoneVector(BaseVector):
         )
 
         docs = []
-        # TODO: add the score threshold to the query
         for result in results:
-            metadata = result.metadata
-            docs.append(
-                Document(
-                    page_content=result.document,
-                    metadata=metadata,
+            metadata = parse_metadata_json(result.metadata)
+            score = 1.0 / (1.0 + float(result.distance))
+            if score >= score_threshold:
+                metadata["score"] = score
+                docs.append(
+                    Document(
+                        page_content=result.document,
+                        metadata=metadata,
+                    )
                 )
-            )
         return docs
 
     @ensure_client
diff --git a/api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py b/api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py
index f8db50455cd..c14c9529fdf 100644
--- a/api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py
+++ b/api/providers/vdb/vdb-matrixone/tests/unit_tests/test_matrixone_vector.py
@@ -206,19 +206,25 @@ def test_delete_and_metadata_methods(matrixone_module):
     assert vector.client.delete.call_count == 3
 
 
-def test_search_by_vector_builds_documents(matrixone_module):
+def test_search_by_vector_applies_score_threshold(matrixone_module):
     vector = matrixone_module.MatrixoneVector("collection_1", _valid_config(matrixone_module))
     vector.client = MagicMock()
     vector.client.query.return_value = [
-        SimpleNamespace(document="doc-a", metadata={"doc_id": "1"}),
-        SimpleNamespace(document="doc-b", metadata={"doc_id": "2"}),
+        SimpleNamespace(document="doc-a", metadata={"doc_id": "1"}, distance=0.25),
+        SimpleNamespace(document="doc-b", metadata={"doc_id": "2"}, distance=2.0),
     ]
 
-    docs = vector.search_by_vector([0.1, 0.2], top_k=2, document_ids_filter=["d-1"])
+    docs = vector.search_by_vector(
+        [0.1, 0.2],
+        top_k=2,
+        score_threshold=0.5,
+        document_ids_filter=["d-1"],
+    )
 
-    assert len(docs) == 2
+    assert len(docs) == 1
     assert docs[0].page_content == "doc-a"
-    assert docs[1].metadata["doc_id"] == "2"
+    assert docs[0].metadata["doc_id"] == "1"
+    assert docs[0].metadata["score"] == pytest.approx(0.8)
     assert vector.client.query.call_args.kwargs["filter"] == {"document_id": {"$in": ["d-1"]}}