Underground-Digital/Workflow-Engine

1from abc import ABC, abstractmethod2from typing import Any, Optional3 4from configs import dify_config5from core.model_manager import ModelManager6from core.model_runtime.entities.model_entities import ModelType7from core.rag.datasource.vdb.vector_base import BaseVector8from core.rag.datasource.vdb.vector_type import VectorType9from core.rag.embedding.cached_embedding import CacheEmbedding10from core.rag.embedding.embedding_base import Embeddings11from core.rag.models.document import Document12from extensions.ext_database import db13from extensions.ext_redis import redis_client14from models.dataset import Dataset, Whitelist15 16 17class AbstractVectorFactory(ABC):18    @abstractmethod19    def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector:20        raise NotImplementedError21 22    @staticmethod23    def gen_index_struct_dict(vector_type: VectorType, collection_name: str) -> dict:24        index_struct_dict = {"type": vector_type, "vector_store": {"class_prefix": collection_name}}25        return index_struct_dict26 27 28class Vector:29    def __init__(self, dataset: Dataset, attributes: Optional[list] = None):30        if attributes is None:31            attributes = ["doc_id", "dataset_id", "document_id", "doc_hash"]32        self._dataset = dataset33        self._embeddings = self._get_embeddings()34        self._attributes = attributes35        self._vector_processor = self._init_vector()36 37    def _init_vector(self) -> BaseVector:38        vector_type = dify_config.VECTOR_STORE39 40        if self._dataset.index_struct_dict:41            vector_type = self._dataset.index_struct_dict["type"]42        else:43            if dify_config.VECTOR_STORE_WHITELIST_ENABLE:44                whitelist = (45                    db.session.query(Whitelist)46                    .filter(Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db")47                    .one_or_none()48                )49                if whitelist:50                    vector_type = VectorType.TIDB_ON_QDRANT51 52        if not vector_type:53            raise ValueError("Vector store must be specified.")54 55        vector_factory_cls = self.get_vector_factory(vector_type)56        return vector_factory_cls().init_vector(self._dataset, self._attributes, self._embeddings)57 58    @staticmethod59    def get_vector_factory(vector_type: str) -> type[AbstractVectorFactory]:60        match vector_type:61            case VectorType.CHROMA:62                from core.rag.datasource.vdb.chroma.chroma_vector import ChromaVectorFactory63 64                return ChromaVectorFactory65            case VectorType.MILVUS:66                from core.rag.datasource.vdb.milvus.milvus_vector import MilvusVectorFactory67 68                return MilvusVectorFactory69            case VectorType.MYSCALE:70                from core.rag.datasource.vdb.myscale.myscale_vector import MyScaleVectorFactory71 72                return MyScaleVectorFactory73            case VectorType.PGVECTOR:74                from core.rag.datasource.vdb.pgvector.pgvector import PGVectorFactory75 76                return PGVectorFactory77            case VectorType.PGVECTO_RS:78                from core.rag.datasource.vdb.pgvecto_rs.pgvecto_rs import PGVectoRSFactory79 80                return PGVectoRSFactory81            case VectorType.QDRANT:82                from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantVectorFactory83 84                return QdrantVectorFactory85            case VectorType.RELYT:86                from core.rag.datasource.vdb.relyt.relyt_vector import RelytVectorFactory87 88                return RelytVectorFactory89            case VectorType.ELASTICSEARCH:90                from core.rag.datasource.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory91 92                return ElasticSearchVectorFactory93            case VectorType.TIDB_VECTOR:94                from core.rag.datasource.vdb.tidb_vector.tidb_vector import TiDBVectorFactory95 96                return TiDBVectorFactory97            case VectorType.WEAVIATE:98                from core.rag.datasource.vdb.weaviate.weaviate_vector import WeaviateVectorFactory99 100                return WeaviateVectorFactory101            case VectorType.TENCENT:102                from core.rag.datasource.vdb.tencent.tencent_vector import TencentVectorFactory103 104                return TencentVectorFactory105            case VectorType.ORACLE:106                from core.rag.datasource.vdb.oracle.oraclevector import OracleVectorFactory107 108                return OracleVectorFactory109            case VectorType.OPENSEARCH:110                from core.rag.datasource.vdb.opensearch.opensearch_vector import OpenSearchVectorFactory111 112                return OpenSearchVectorFactory113            case VectorType.ANALYTICDB:114                from core.rag.datasource.vdb.analyticdb.analyticdb_vector import AnalyticdbVectorFactory115 116                return AnalyticdbVectorFactory117            case VectorType.COUCHBASE:118                from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseVectorFactory119 120                return CouchbaseVectorFactory121            case VectorType.BAIDU:122                from core.rag.datasource.vdb.baidu.baidu_vector import BaiduVectorFactory123 124                return BaiduVectorFactory125            case VectorType.VIKINGDB:126                from core.rag.datasource.vdb.vikingdb.vikingdb_vector import VikingDBVectorFactory127 128                return VikingDBVectorFactory129            case VectorType.UPSTASH:130                from core.rag.datasource.vdb.upstash.upstash_vector import UpstashVectorFactory131 132                return UpstashVectorFactory133            case VectorType.TIDB_ON_QDRANT:134                from core.rag.datasource.vdb.tidb_on_qdrant.tidb_on_qdrant_vector import TidbOnQdrantVectorFactory135 136                return TidbOnQdrantVectorFactory137            case VectorType.LINDORM:138                from core.rag.datasource.vdb.lindorm.lindorm_vector import LindormVectorStoreFactory139 140                return LindormVectorStoreFactory141            case VectorType.OCEANBASE:142                from core.rag.datasource.vdb.oceanbase.oceanbase_vector import OceanBaseVectorFactory143 144                return OceanBaseVectorFactory145            case _:146                raise ValueError(f"Vector store {vector_type} is not supported.")147 148    def create(self, texts: Optional[list] = None, **kwargs):149        if texts:150            embeddings = self._embeddings.embed_documents([document.page_content for document in texts])151            self._vector_processor.create(texts=texts, embeddings=embeddings, **kwargs)152 153    def add_texts(self, documents: list[Document], **kwargs):154        if kwargs.get("duplicate_check", False):155            documents = self._filter_duplicate_texts(documents)156 157        embeddings = self._embeddings.embed_documents([document.page_content for document in documents])158        self._vector_processor.create(texts=documents, embeddings=embeddings, **kwargs)159 160    def text_exists(self, id: str) -> bool:161        return self._vector_processor.text_exists(id)162 163    def delete_by_ids(self, ids: list[str]) -> None:164        self._vector_processor.delete_by_ids(ids)165 166    def delete_by_metadata_field(self, key: str, value: str) -> None:167        self._vector_processor.delete_by_metadata_field(key, value)168 169    def search_by_vector(self, query: str, **kwargs: Any) -> list[Document]:170        query_vector = self._embeddings.embed_query(query)171        return self._vector_processor.search_by_vector(query_vector, **kwargs)172 173    def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:174        return self._vector_processor.search_by_full_text(query, **kwargs)175 176    def delete(self) -> None:177        self._vector_processor.delete()178        # delete collection redis cache179        if self._vector_processor.collection_name:180            collection_exist_cache_key = "vector_indexing_{}".format(self._vector_processor.collection_name)181            redis_client.delete(collection_exist_cache_key)182 183    def _get_embeddings(self) -> Embeddings:184        model_manager = ModelManager()185 186        embedding_model = model_manager.get_model_instance(187            tenant_id=self._dataset.tenant_id,188            provider=self._dataset.embedding_model_provider,189            model_type=ModelType.TEXT_EMBEDDING,190            model=self._dataset.embedding_model,191        )192        return CacheEmbedding(embedding_model)193 194    def _filter_duplicate_texts(self, texts: list[Document]) -> list[Document]:195        for text in texts.copy():196            doc_id = text.metadata["doc_id"]197            exists_duplicate_node = self.text_exists(doc_id)198            if exists_duplicate_node:199                texts.remove(text)200 201        return texts202 203    def __getattr__(self, name):204        if self._vector_processor is not None:205            method = getattr(self._vector_processor, name)206            if callable(method):207                return method208 209        raise AttributeError(f"'vector_processor' object has no attribute '{name}'")210