From bfb366240cc6eb3296e6dd7027b973ceb062121c Mon Sep 17 00:00:00 2001 From: LZC6244 Date: Fri, 18 Jul 2025 11:23:44 +0800 Subject: [PATCH] =?UTF-8?q?=E7=9F=A5=E8=AF=86=E5=BA=93=20=5Finit=5Fvector?= =?UTF-8?q?=20import=20=E6=8F=90=E9=80=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/core/rag/datasource/vdb/vector_factory.py | 119 +++++++++++++++++- 1 file changed, 117 insertions(+), 2 deletions(-) diff --git a/api/core/rag/datasource/vdb/vector_factory.py b/api/core/rag/datasource/vdb/vector_factory.py index 00080b0fae..ed5a7b7632 100644 --- a/api/core/rag/datasource/vdb/vector_factory.py +++ b/api/core/rag/datasource/vdb/vector_factory.py @@ -28,6 +28,122 @@ class AbstractVectorFactory(ABC): index_struct_dict = {"type": vector_type, "vector_store": {"class_prefix": collection_name}} return index_struct_dict +VECTOR_FACTORY_CLS: type[AbstractVectorFactory] + +match dify_config.VECTOR_STORE: + case VectorType.CHROMA: + from core.rag.datasource.vdb.chroma.chroma_vector import ChromaVectorFactory + + VECTOR_FACTORY_CLS = ChromaVectorFactory + case VectorType.MILVUS: + from core.rag.datasource.vdb.milvus.milvus_vector import MilvusVectorFactory + + VECTOR_FACTORY_CLS = MilvusVectorFactory + case VectorType.MYSCALE: + from core.rag.datasource.vdb.myscale.myscale_vector import MyScaleVectorFactory + + VECTOR_FACTORY_CLS = MyScaleVectorFactory + case VectorType.PGVECTOR: + from core.rag.datasource.vdb.pgvector.pgvector import PGVectorFactory + + VECTOR_FACTORY_CLS = PGVectorFactory + case VectorType.VASTBASE: + from core.rag.datasource.vdb.pyvastbase.vastbase_vector import VastbaseVectorFactory + + VECTOR_FACTORY_CLS = VastbaseVectorFactory + case VectorType.PGVECTO_RS: + from core.rag.datasource.vdb.pgvecto_rs.pgvecto_rs import PGVectoRSFactory + + VECTOR_FACTORY_CLS = PGVectoRSFactory + case VectorType.QDRANT: + from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantVectorFactory + + VECTOR_FACTORY_CLS = QdrantVectorFactory + case VectorType.RELYT: + from core.rag.datasource.vdb.relyt.relyt_vector import RelytVectorFactory + + VECTOR_FACTORY_CLS = RelytVectorFactory + case VectorType.ELASTICSEARCH: + from core.rag.datasource.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory + + VECTOR_FACTORY_CLS = ElasticSearchVectorFactory + case VectorType.ELASTICSEARCH_JA: + from core.rag.datasource.vdb.elasticsearch.elasticsearch_ja_vector import ( + ElasticSearchJaVectorFactory, + ) + + VECTOR_FACTORY_CLS = ElasticSearchJaVectorFactory + case VectorType.TIDB_VECTOR: + from core.rag.datasource.vdb.tidb_vector.tidb_vector import TiDBVectorFactory + + VECTOR_FACTORY_CLS = TiDBVectorFactory + case VectorType.WEAVIATE: + from core.rag.datasource.vdb.weaviate.weaviate_vector import WeaviateVectorFactory + + VECTOR_FACTORY_CLS = WeaviateVectorFactory + case VectorType.TENCENT: + from core.rag.datasource.vdb.tencent.tencent_vector import TencentVectorFactory + + VECTOR_FACTORY_CLS = TencentVectorFactory + case VectorType.ORACLE: + from core.rag.datasource.vdb.oracle.oraclevector import OracleVectorFactory + + VECTOR_FACTORY_CLS = OracleVectorFactory + case VectorType.OPENSEARCH: + from core.rag.datasource.vdb.opensearch.opensearch_vector import OpenSearchVectorFactory + + VECTOR_FACTORY_CLS = OpenSearchVectorFactory + case VectorType.ANALYTICDB: + from core.rag.datasource.vdb.analyticdb.analyticdb_vector import AnalyticdbVectorFactory + + VECTOR_FACTORY_CLS = AnalyticdbVectorFactory + case VectorType.COUCHBASE: + from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseVectorFactory + + VECTOR_FACTORY_CLS = CouchbaseVectorFactory + case VectorType.BAIDU: + from core.rag.datasource.vdb.baidu.baidu_vector import BaiduVectorFactory + + VECTOR_FACTORY_CLS = BaiduVectorFactory + case VectorType.VIKINGDB: + from core.rag.datasource.vdb.vikingdb.vikingdb_vector import VikingDBVectorFactory + + VECTOR_FACTORY_CLS = VikingDBVectorFactory + case VectorType.UPSTASH: + from core.rag.datasource.vdb.upstash.upstash_vector import UpstashVectorFactory + + VECTOR_FACTORY_CLS = UpstashVectorFactory + case VectorType.TIDB_ON_QDRANT: + from core.rag.datasource.vdb.tidb_on_qdrant.tidb_on_qdrant_vector import TidbOnQdrantVectorFactory + + VECTOR_FACTORY_CLS = TidbOnQdrantVectorFactory + case VectorType.LINDORM: + from core.rag.datasource.vdb.lindorm.lindorm_vector import LindormVectorStoreFactory + + VECTOR_FACTORY_CLS = LindormVectorStoreFactory + case VectorType.OCEANBASE: + from core.rag.datasource.vdb.oceanbase.oceanbase_vector import OceanBaseVectorFactory + + VECTOR_FACTORY_CLS = OceanBaseVectorFactory + case VectorType.OPENGAUSS: + from core.rag.datasource.vdb.opengauss.opengauss import OpenGaussFactory + + VECTOR_FACTORY_CLS = OpenGaussFactory + case VectorType.TABLESTORE: + from core.rag.datasource.vdb.tablestore.tablestore_vector import TableStoreVectorFactory + + VECTOR_FACTORY_CLS = TableStoreVectorFactory + case VectorType.HUAWEI_CLOUD: + from core.rag.datasource.vdb.huawei.huawei_cloud_vector import HuaweiCloudVectorFactory + + VECTOR_FACTORY_CLS = HuaweiCloudVectorFactory + case VectorType.MATRIXONE: + from core.rag.datasource.vdb.matrixone.matrixone_vector import MatrixoneVectorFactory + + VECTOR_FACTORY_CLS = MatrixoneVectorFactory + case _: + raise ValueError(f"Vector store {dify_config.VECTOR_STORE} is not supported.") + class Vector: def __init__(self, dataset: Dataset, attributes: Optional[list] = None): @@ -56,8 +172,7 @@ class Vector: if not vector_type: raise ValueError("Vector store must be specified.") - vector_factory_cls = self.get_vector_factory(vector_type) - return vector_factory_cls().init_vector(self._dataset, self._attributes, self._embeddings) + return VECTOR_FACTORY_CLS().init_vector(self._dataset, self._attributes, self._embeddings) @staticmethod def get_vector_factory(vector_type: str) -> type[AbstractVectorFactory]: