diff --git a/api/providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py b/api/providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py index 30974156ab9352..d1f2491ec8a0a1 100644 --- a/api/providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py +++ b/api/providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py @@ -1,9 +1,9 @@ import json import logging import math -from typing import Any, TypedDict, override +from typing import Any, TypedDict, cast, override -from pydantic import BaseModel +from pydantic import BaseModel, model_validator from sqlalchemy.orm import Session from tcvdb_text.encoder import BM25Encoder # type: ignore from tcvectordb import RPCVectorDBClient, VectorDBException # type: ignore @@ -36,7 +36,7 @@ class TencentConfig(BaseModel): api_key: str | None = None timeout: float = 30 username: str | None = None - database: str | None = None + database: str index_type: str = "HNSW" metric_type: str = "IP" shard: int = 1 @@ -44,6 +44,13 @@ class TencentConfig(BaseModel): max_upsert_batch_size: int = 128 enable_hybrid_search: bool = False # Flag to enable hybrid search + @model_validator(mode="before") + @classmethod + def validate_config(cls, values: dict[str, Any]) -> dict[str, Any]: + if not values.get("database"): + raise ValueError("config TENCENT_VECTOR_DB_DATABASE is required") + return values + def to_tencent_params(self) -> TencentParamsDict: result: TencentParamsDict = { "url": self.url, @@ -87,7 +94,10 @@ def _load_collection(self): if idx.name == "sparse_vector": has_hybrid_search = True elif idx.name == "vector": - self._dimension = idx.dimension + # The "vector" index is the VectorIndex _create_collection made. + dimension = cast(vdb_index.VectorIndex, idx).dimension + if dimension is not None: + self._dimension = dimension if not has_hybrid_search: self._enable_hybrid_search = False @@ -151,7 +161,7 @@ def _create_collection(self, dimension: int): index_type=enum.IndexType.SPARSE_INVERTED, metric_type=enum.MetricType.IP, ) - indexes = [index_id, index_vector, index_metadate] + indexes: list[vdb_index.IndexField] = [index_id, index_vector, index_metadate] if self._enable_hybrid_search: indexes.append(index_sparse_vector) try: @@ -170,7 +180,7 @@ def _create_collection(self, dimension: int): index_metadate = vdb_index.FilterIndex( self.field_metadata, enum.FieldType.String, enum.IndexType.FILTER ) - indexes = [index_id, index_vector, index_metadate] + indexes: list[vdb_index.IndexField] = [index_id, index_vector, index_metadate] if self._enable_hybrid_search: indexes.append(index_sparse_vector) self._client.create_collection( @@ -196,7 +206,7 @@ def add_texts(self, documents: list[Document], embeddings: list[list[float]], ** batch_size = self._client_config.max_upsert_batch_size batch = math.ceil(total_count / batch_size) for j in range(batch): - docs = [] + docs: list[document.Document | dict[str, Any]] = [] start_idx = j * batch_size end_idx = min(total_count, (j + 1) * batch_size) for i in range(start_idx, end_idx): @@ -267,7 +277,8 @@ def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Doc database_name=self._client_config.database, collection_name=self.collection_name, vectors=[query_vector], - filter=filter, + # tcvectordb declares `filter: Union[Filter, str] = None`; None means no filter. + filter=filter, # pyrefly: ignore[bad-argument-type] params=document.HNSWSearchParams(ef=kwargs.get("ef", 10)), retrieve_vector=False, limit=kwargs.get("top_k", 4), @@ -296,7 +307,8 @@ def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: match=[ KeywordSearch( field_name="sparse_vector", - data=bm25.encode_queries(query), + # A single query string encodes to a single sparse vector. + data=cast(list[list[float | int]], bm25.encode_queries(query)), ), ], rerank=WeightedRerank( @@ -305,7 +317,8 @@ def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: ), retrieve_vector=False, limit=kwargs.get("top_k", 4), - filter=filter, + # tcvectordb declares `filter: Union[Filter, str] = None`; None means no filter. + filter=filter, # pyrefly: ignore[bad-argument-type] ) score_threshold = float(kwargs.get("score_threshold") or 0.0) return self._get_search_res(res, score_threshold) @@ -354,7 +367,7 @@ def init_vector( api_key=dify_config.TENCENT_VECTOR_DB_API_KEY, timeout=dify_config.TENCENT_VECTOR_DB_TIMEOUT, username=dify_config.TENCENT_VECTOR_DB_USERNAME, - database=dify_config.TENCENT_VECTOR_DB_DATABASE, + database=dify_config.TENCENT_VECTOR_DB_DATABASE or "", shard=dify_config.TENCENT_VECTOR_DB_SHARD, replicas=dify_config.TENCENT_VECTOR_DB_REPLICAS, enable_hybrid_search=dify_config.TENCENT_VECTOR_DB_ENABLE_HYBRID_SEARCH or False, diff --git a/api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py b/api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py index 5d4d05aec28bcd..259d058297731d 100644 --- a/api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py +++ b/api/providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py @@ -96,35 +96,36 @@ class _SparseIndex: def __init__(self, **kwargs): self.kwargs = kwargs - tcvectordb_enum.IndexType = SimpleNamespace( + # ModuleType does not declare these attributes, so set them through the module namespace. + vars(tcvectordb_enum)["IndexType"] = SimpleNamespace( __members__={"HNSW": "HNSW", "PRIMARY_KEY": "PRIMARY_KEY", "FILTER": "FILTER", "SPARSE_INVERTED": "SPARSE"}, PRIMARY_KEY="PRIMARY_KEY", FILTER="FILTER", SPARSE_INVERTED="SPARSE", ) - tcvectordb_enum.MetricType = SimpleNamespace(__members__={"IP": "IP"}, IP="IP") - tcvectordb_enum.FieldType = SimpleNamespace(String="String", Json="Json", SparseVector="SparseVector") + vars(tcvectordb_enum)["MetricType"] = SimpleNamespace(__members__={"IP": "IP"}, IP="IP") + vars(tcvectordb_enum)["FieldType"] = SimpleNamespace(String="String", Json="Json", SparseVector="SparseVector") - tcvectordb_document.Document = _Document - tcvectordb_document.HNSWSearchParams = _HNSWSearchParams - tcvectordb_document.AnnSearch = _AnnSearch - tcvectordb_document.Filter = _Filter - tcvectordb_document.KeywordSearch = _KeywordSearch - tcvectordb_document.WeightedRerank = _WeightedRerank + vars(tcvectordb_document)["Document"] = _Document + vars(tcvectordb_document)["HNSWSearchParams"] = _HNSWSearchParams + vars(tcvectordb_document)["AnnSearch"] = _AnnSearch + vars(tcvectordb_document)["Filter"] = _Filter + vars(tcvectordb_document)["KeywordSearch"] = _KeywordSearch + vars(tcvectordb_document)["WeightedRerank"] = _WeightedRerank - tcvectordb_index.HNSWParams = _HNSWParams - tcvectordb_index.FilterIndex = _FilterIndex - tcvectordb_index.VectorIndex = _VectorIndex - tcvectordb_index.SparseIndex = _SparseIndex + vars(tcvectordb_index)["HNSWParams"] = _HNSWParams + vars(tcvectordb_index)["FilterIndex"] = _FilterIndex + vars(tcvectordb_index)["VectorIndex"] = _VectorIndex + vars(tcvectordb_index)["SparseIndex"] = _SparseIndex - tcvdb_text_encoder.BM25Encoder = _BM25Encoder + vars(tcvdb_text_encoder)["BM25Encoder"] = _BM25Encoder - tcvectordb_model.document = tcvectordb_document - tcvectordb_model.enum = tcvectordb_enum - tcvectordb_model.index = tcvectordb_index + vars(tcvectordb_model)["document"] = tcvectordb_document + vars(tcvectordb_model)["enum"] = tcvectordb_enum + vars(tcvectordb_model)["index"] = tcvectordb_index - tcvectordb.RPCVectorDBClient = RPCVectorDBClient - tcvectordb.VectorDBException = VectorDBError + vars(tcvectordb)["RPCVectorDBClient"] = RPCVectorDBClient + vars(tcvectordb)["VectorDBException"] = VectorDBError return { "tcvdb_text": tcvdb_text, @@ -186,7 +187,14 @@ def test_config_and_init_paths(tencent_module): indexes=[SimpleNamespace(name="vector", dimension=512)] ) vector._load_collection() - assert vector._enable_hybrid_search is False + # Without a sparse_vector index, hybrid search is turned off. + assert not vector._enable_hybrid_search + + +@pytest.mark.parametrize("database", [None, ""]) +def test_config_requires_a_database(tencent_module, database): + with pytest.raises(ValueError, match="config TENCENT_VECTOR_DB_DATABASE is required"): + _config(tencent_module, database=database) def test_create_collection_branches(tencent_module, monkeypatch: pytest.MonkeyPatch): diff --git a/api/pyrefly-local-excludes.txt b/api/pyrefly-local-excludes.txt index 8009889baf7277..e127244b8c3955 100644 --- a/api/pyrefly-local-excludes.txt +++ b/api/pyrefly-local-excludes.txt @@ -31,9 +31,7 @@ providers/vdb/vdb-qdrant/tests/unit_tests/test_qdrant_vector.py providers/vdb/vdb-relyt/tests/unit_tests/test_relyt_vector.py providers/vdb/vdb-tablestore/tests/integration_tests/test_tablestore.py providers/vdb/vdb-tablestore/tests/unit_tests/test_tablestore_vector.py -providers/vdb/vdb-tencent/src/dify_vdb_tencent/tencent_vector.py providers/vdb/vdb-tencent/tests/integration_tests/conftest.py -providers/vdb/vdb-tencent/tests/unit_tests/test_tencent_vector.py providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_service.py providers/vdb/vdb-tidb-on-qdrant/tests/unit_tests/test_tidb_service.py