mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
fix(matrixone): apply score threshold to vector search (#39784)
Co-authored-by: caosurong <surong.cao@thinkingdata.cn>
This commit is contained in:
@@ -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"]}}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user