fix(matrixone): apply score threshold to vector search (#39784)

Co-authored-by: caosurong <surong.cao@thinkingdata.cn>
This commit is contained in:
csurong
2026-08-02 02:43:56 +00:00
committed by GitHub
co-authored by caosurong
parent a42d486bc0
commit 8a2cf65ec5
2 changed files with 22 additions and 13 deletions
@@ -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
@@ -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"]}}