diff --git a/api/providers/vdb/vdb-analyticdb/src/dify_vdb_analyticdb/analyticdb_vector_sql.py b/api/providers/vdb/vdb-analyticdb/src/dify_vdb_analyticdb/analyticdb_vector_sql.py index 2fac15600e2..729a77f9d43 100644 --- a/api/providers/vdb/vdb-analyticdb/src/dify_vdb_analyticdb/analyticdb_vector_sql.py +++ b/api/providers/vdb/vdb-analyticdb/src/dify_vdb_analyticdb/analyticdb_vector_sql.py @@ -212,7 +212,7 @@ class AnalyticdbVectorBySql: if not isinstance(top_k, int) or top_k <= 0: raise ValueError("top_k must be a positive integer") document_ids_filter = kwargs.get("document_ids_filter") - where_clause = "WHERE 1=1" + where_clause = "WHERE 1=1 " if document_ids_filter: document_ids = ", ".join(f"'{id}'" for id in document_ids_filter) where_clause += f"AND metadata_->>'document_id' IN ({document_ids})" diff --git a/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py b/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py index cd255b37cfc..9ea751ce644 100644 --- a/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py +++ b/api/providers/vdb/vdb-analyticdb/tests/unit_tests/src/dify_vdb_analyticdb/test_analyticdb_vector_sql.py @@ -203,10 +203,34 @@ def test_search_by_vector_validates_top_k(invalid_top_k): vector.search_by_vector([0.1, 0.2], top_k=invalid_top_k) -def test_search_by_vector_returns_documents_above_threshold(): - vector = AnalyticdbVectorBySql.__new__(AnalyticdbVectorBySql) - vector.table_name = "dify.collection" - cursor = MagicMock() +@pytest.mark.parametrize( + ("document_ids_filter", "expected_where_clause"), + [ + pytest.param(None, "WHERE 1=1", id="no-filter"), + pytest.param([], "WHERE 1=1", id="empty-filter"), + pytest.param( + ["doc-1"], + "WHERE 1=1 AND metadata_->>'document_id' IN ('doc-1')", + id="single-document-id", + ), + pytest.param( + ["doc-1", "doc-2"], + "WHERE 1=1 AND metadata_->>'document_id' IN ('doc-1', 'doc-2')", + id="multiple-document-ids", + ), + ], +) +def test_search_by_vector_returns_documents_above_threshold( + monkeypatch: pytest.MonkeyPatch, + document_ids_filter: list[str] | None, + expected_where_clause: str, +) -> None: + """Given an optional ID filter, emit valid SQL and return only documents above the score threshold.""" + pool = MagicMock() + cursor = pool.getconn.return_value.cursor.return_value + monkeypatch.setattr(AnalyticdbVectorBySql, "_initialize", lambda self: None) + monkeypatch.setattr(AnalyticdbVectorBySql, "_create_connection_pool", lambda self: pool) + vector = AnalyticdbVectorBySql("collection", AnalyticdbVectorBySqlConfig(**_config_values())) cursor.__iter__.return_value = iter( [ ("id1", [1.0], 0.8, "content 1", {"doc_id": "id1", "document_id": "doc-1"}), @@ -214,18 +238,18 @@ def test_search_by_vector_returns_documents_above_threshold(): ] ) - @contextmanager - def _cursor_context(): - yield cursor - - vector._get_cursor = _cursor_context - - docs = vector.search_by_vector([0.1, 0.2], top_k=2, score_threshold=0.5, document_ids_filter=["doc-1"]) + docs = vector.search_by_vector([0.1, 0.2], top_k=2, score_threshold=0.5, document_ids_filter=document_ids_filter) assert len(docs) == 1 assert docs[0].page_content == "content 1" assert docs[0].metadata["score"] == 0.8 + executed_sql = cursor.execute.call_args.args[0] + normalized_sql = " ".join(executed_sql.split()) + assert f"FROM dify.collection {expected_where_clause} ORDER BY score" in normalized_sql + assert "1=1AND" not in executed_sql + assert ")ORDER BY" not in executed_sql + @pytest.mark.parametrize("invalid_top_k", [0, "x", -1]) def test_search_by_full_text_validates_top_k(invalid_top_k): diff --git a/api/providers/vdb/vdb-relyt/src/dify_vdb_relyt/relyt_vector.py b/api/providers/vdb/vdb-relyt/src/dify_vdb_relyt/relyt_vector.py index 1033db650f1..ea9ab4cceed 100644 --- a/api/providers/vdb/vdb-relyt/src/dify_vdb_relyt/relyt_vector.py +++ b/api/providers/vdb/vdb-relyt/src/dify_vdb_relyt/relyt_vector.py @@ -255,9 +255,9 @@ class RelytVector(BaseVector): filter_condition = "" if filter is not None: conditions = [ - f"metadata->>'{key!r}' in ({', '.join(map(repr, value))})" + f"metadata->>'{key}' in ({', '.join(map(repr, value))})" if len(value) > 1 - else f"metadata->>'{key!r}' = {value[0]!r}" + else f"metadata->>'{key}' = {value[0]!r}" for key, value in filter.items() ] filter_condition = f"WHERE {' AND '.join(conditions)}" diff --git a/api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py b/api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py index 57fbb4ec1a9..13c65ffb96f 100644 --- a/api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py +++ b/api/providers/vdb/vdb-relyt/tests/unit_tests/src/dify_vdb_relyt/test_relyt_vector.py @@ -249,11 +249,36 @@ def test_text_exists_false(relyt_module, monkeypatch: pytest.MonkeyPatch): # 6. similarity_search_with_score_by_vector returns Documents and scores -def test_similarity_search_with_score_by_vector(relyt_module): - vector = relyt_module.RelytVector.__new__(relyt_module.RelytVector) - vector._collection_name = "collection_1" - vector.client = MagicMock() - vector.embedding_dimension = 3 +@pytest.mark.parametrize( + ("metadata_filter", "expected_condition"), + [ + pytest.param( + {"document_id": ["doc-1"], "group_id": ["g-1"]}, + "metadata->>'document_id' = 'doc-1' AND metadata->>'group_id' = 'g-1'", + id="single-value", + ), + pytest.param( + {"document_id": ["doc-1", "doc-2"], "group_id": ["g-1", "g-2"]}, + "metadata->>'document_id' in ('doc-1', 'doc-2') AND metadata->>'group_id' in ('g-1', 'g-2')", + id="multiple-values", + ), + pytest.param( + {"document_id": ["doc-1", "doc-2"], "group_id": ["g-1"]}, + "metadata->>'document_id' in ('doc-1', 'doc-2') AND metadata->>'group_id' = 'g-1'", + id="mixed-value-counts", + ), + ], +) +def test_similarity_search_with_score_by_vector( + relyt_module: types.ModuleType, + monkeypatch: pytest.MonkeyPatch, + metadata_filter: dict[str, list[str]], + expected_condition: str, +) -> None: + """Given nonempty ID filters, emit correctly quoted predicates and return documents with scores.""" + engine = MagicMock() + monkeypatch.setattr(relyt_module, "create_engine", lambda *_args, **_kwargs: engine) + vector = relyt_module.RelytVector("collection_1", _config(relyt_module), group_id="group-1") result_rows = [ SimpleNamespace(document="doc-a", metadata={"doc_id": "1"}, distance=0.1), SimpleNamespace(document="doc-b", metadata={"doc_id": "2"}, distance=0.8), @@ -263,9 +288,13 @@ def test_similarity_search_with_score_by_vector(relyt_module): conn.__exit__.return_value = None conn.execute.return_value.fetchall.return_value = result_rows vector.client.connect.return_value = conn - similarities = vector.similarity_search_with_score_by_vector([0.1, 0.2], k=2, filter={"document_id": ["d-1"]}) + similarities = vector.similarity_search_with_score_by_vector([0.1, 0.2], k=2, filter=metadata_filter) assert len(similarities) == 2 assert similarities[0][0].page_content == "doc-a" + executed_sql = str(conn.execute.call_args.args[0]) + assert f"WHERE {expected_condition}" in executed_sql + assert "''document_id''" not in executed_sql + assert "''group_id''" not in executed_sql # 7. search_by_vector filters by score and ids