import json import logging import uuid from collections.abc import Callable from functools import wraps from typing import Any, Concatenate from mo_vector.client import MoVectorClient # type: ignore from pydantic import BaseModel, model_validator from configs import dify_config from core.rag.datasource.vdb.field import parse_metadata_json from core.rag.datasource.vdb.vector_base import BaseVector from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory from core.rag.datasource.vdb.vector_type import VectorType from core.rag.embedding.embedding_base import Embeddings from core.rag.models.document import Document from extensions.ext_redis import redis_client from models.dataset import Dataset logger = logging.getLogger(__name__) def ensure_client[T: MatrixoneVector, **P, R]( func: Callable[Concatenate[T, P], R], ) -> Callable[Concatenate[T, P], R]: @wraps(func) def wrapper(self: T, *args: P.args, **kwargs: P.kwargs) -> R: if self.client is None: self.client = self._get_client(None, False) return func(self, *args, **kwargs) return wrapper class MatrixoneConfig(BaseModel): host: str = "localhost" port: int = 6001 user: str = "dump" password: str = "111" database: str = "dify" metric: str = "l2" @model_validator(mode="before") @classmethod def validate_config(cls, values: dict[str, Any]): if not values["host"]: raise ValueError("config host is required") if not values["port"]: raise ValueError("config port is required") if not values["user"]: raise ValueError("config user is required") if not values["password"]: raise ValueError("config password is required") if not values["database"]: raise ValueError("config database is required") return values class MatrixoneVector(BaseVector): """ Matrixone vector storage implementation. """ def __init__(self, collection_name: str, config: MatrixoneConfig): super().__init__(collection_name) self.config = config self.collection_name = collection_name.lower() self.client = None @property def collection_name(self): return self._collection_name @collection_name.setter def collection_name(self, value): self._collection_name = value def get_type(self) -> str: return VectorType.MATRIXONE def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): if self.client is None: self.client = self._get_client(len(embeddings[0]), True) return self.add_texts(texts, embeddings) def _get_client(self, dimension: int | None = None, create_table: bool = False) -> MoVectorClient: """ Create a new client for the collection. The collection will be created if it doesn't exist. """ lock_name = f"vector_indexing_lock_{self._collection_name}" with redis_client.lock(lock_name, timeout=20): client = MoVectorClient( connection_string=f"mysql+pymysql://{self.config.user}:{self.config.password}@{self.config.host}:{self.config.port}/{self.config.database}", table_name=self.collection_name, vector_dimension=dimension, create_table=create_table, ) collection_exist_cache_key = f"vector_indexing_{self._collection_name}" if redis_client.get(collection_exist_cache_key): return client try: client.create_full_text_index() redis_client.set(collection_exist_cache_key, 1, ex=3600) except Exception: logger.exception("Failed to create full text index") return client def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): if self.client is None: self.client = self._get_client(len(embeddings[0]), True) assert self.client is not None ids = [] for doc in documents: if doc.metadata is not None: doc_id = doc.metadata.get("doc_id", str(uuid.uuid4())) ids.append(doc_id) self.client.insert( texts=[doc.page_content for doc in documents], embeddings=embeddings, metadatas=[doc.metadata for doc in documents], ids=ids, ) return ids @ensure_client def text_exists(self, id: str) -> bool: assert self.client is not None result = self.client.get(ids=[id]) return len(result) > 0 @ensure_client def delete_by_ids(self, ids: list[str]): assert self.client is not None if not ids: return self.client.delete(ids=ids) @ensure_client def get_ids_by_metadata_field(self, key: str, value: str): assert self.client is not None results = self.client.query_by_metadata(filter={key: value}) return [result.id for result in results] @ensure_client def delete_by_metadata_field(self, key: str, value: str): assert self.client is not None self.client.delete(filter={key: value}) @ensure_client def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: assert self.client is not None top_k = kwargs.get("top_k", 5) document_ids_filter = kwargs.get("document_ids_filter") filter = None if document_ids_filter: filter = {"document_id": {"$in": document_ids_filter}} results = self.client.query( query_vector=query_vector, k=top_k, filter=filter, ) 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, ) ) return docs @ensure_client def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: assert self.client is not None top_k = kwargs.get("top_k", 5) document_ids_filter = kwargs.get("document_ids_filter") filter = None if document_ids_filter: filter = {"document_id": {"$in": document_ids_filter}} score_threshold = float(kwargs.get("score_threshold", 0.0)) results = self.client.full_text_query( keywords=[query], k=top_k, filter=filter, ) docs = [] for result in results: metadata = parse_metadata_json(result.metadata) score = 1 - result.distance if score >= score_threshold: metadata["score"] = score docs.append( Document( page_content=result.document, metadata=metadata, ) ) return docs @ensure_client def delete(self): assert self.client is not None self.client.delete() class MatrixoneVectorFactory(AbstractVectorFactory): def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> MatrixoneVector: if dataset.index_struct_dict: class_prefix: str = dataset.index_struct_dict["vector_store"]["class_prefix"] collection_name = class_prefix else: dataset_id = dataset.id collection_name = Dataset.gen_collection_name_by_id(dataset_id) dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.MATRIXONE, collection_name)) config = MatrixoneConfig( host=dify_config.MATRIXONE_HOST or "localhost", port=dify_config.MATRIXONE_PORT or 6001, user=dify_config.MATRIXONE_USER or "dump", password=dify_config.MATRIXONE_PASSWORD or "111", database=dify_config.MATRIXONE_DATABASE or "dify", metric=dify_config.MATRIXONE_METRIC or "l2", ) return MatrixoneVector(collection_name=collection_name, config=config)