import json import os import uuid from collections.abc import Generator, Iterable, Sequence from itertools import islice from typing import TYPE_CHECKING, Any, cast import qdrant_client from flask import current_app from pydantic import BaseModel from qdrant_client.http import models as rest from qdrant_client.http.models import ( FilterSelector, HnswConfigDiff, PayloadSchemaType, TextIndexParams, TextIndexType, TokenizerType, ) from qdrant_client.local.qdrant_local import QdrantLocal from sqlalchemy import select from configs import dify_config from core.rag.datasource.vdb.field import Field from core.rag.datasource.vdb.vector_base import BaseVector, VectorIndexStructDict 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_database import db from extensions.ext_redis import redis_client from models.dataset import Dataset, DatasetCollectionBinding if TYPE_CHECKING: from qdrant_client.conversions import common_types from qdrant_client.http import models as rest type DictFilter = dict[str, str | int | bool | dict | list] type MetadataFilter = DictFilter | common_types.Filter class PathQdrantParams(BaseModel): path: str class UrlQdrantParams(BaseModel): url: str api_key: str | None timeout: float verify: bool grpc_port: int prefer_grpc: bool class QdrantConfig(BaseModel): endpoint: str api_key: str | None = None timeout: float = 20 root_path: str | None = None grpc_port: int = 6334 prefer_grpc: bool = False replication_factor: int = 1 write_consistency_factor: int = 1 def to_qdrant_params(self) -> PathQdrantParams | UrlQdrantParams: if self.endpoint and self.endpoint.startswith("path:"): path = self.endpoint.replace("path:", "") if not os.path.isabs(path): if not self.root_path: raise ValueError("Root path is not set") path = os.path.join(self.root_path, path) return PathQdrantParams(path=path) else: return UrlQdrantParams( url=self.endpoint, api_key=self.api_key, timeout=self.timeout, verify=self.endpoint.startswith("https"), grpc_port=self.grpc_port, prefer_grpc=self.prefer_grpc, ) class QdrantVector(BaseVector): def __init__(self, collection_name: str, group_id: str, config: QdrantConfig, distance_func: str = "Cosine"): super().__init__(collection_name) self._client_config = config self._client = qdrant_client.QdrantClient(**self._client_config.to_qdrant_params().model_dump()) self._distance_func = distance_func.upper() self._group_id = group_id def get_type(self) -> str: return VectorType.QDRANT def to_index_struct(self) -> VectorIndexStructDict: result: VectorIndexStructDict = { "type": self.get_type(), "vector_store": {"class_prefix": self._collection_name}, } return result def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): if texts: # get embedding vector size vector_size = len(embeddings[0]) # get collection name collection_name = self._collection_name # create collection self.create_collection(collection_name, vector_size) self.add_texts(texts, embeddings, **kwargs) def create_collection(self, collection_name: str, vector_size: int): lock_name = f"vector_indexing_lock_{collection_name}" with redis_client.lock(lock_name, timeout=20): collection_exist_cache_key = f"vector_indexing_{self._collection_name}" if redis_client.get(collection_exist_cache_key): return collection_name = collection_name or uuid.uuid4().hex all_collection_name = [] collections_response = self._client.get_collections() collection_list = collections_response.collections for collection in collection_list: all_collection_name.append(collection.name) if collection_name not in all_collection_name: from qdrant_client.http import models as rest vectors_config = rest.VectorParams( size=vector_size, distance=rest.Distance[self._distance_func], ) hnsw_config = HnswConfigDiff( m=0, payload_m=16, ef_construct=100, full_scan_threshold=10000, max_indexing_threads=0, on_disk=False, ) self._client.create_collection( collection_name=collection_name, vectors_config=vectors_config, hnsw_config=hnsw_config, timeout=int(self._client_config.timeout), replication_factor=self._client_config.replication_factor, write_consistency_factor=self._client_config.write_consistency_factor, ) # create group_id payload index self._client.create_payload_index( collection_name, Field.GROUP_KEY, field_schema=PayloadSchemaType.KEYWORD ) # create doc_id payload index self._client.create_payload_index(collection_name, Field.DOC_ID, field_schema=PayloadSchemaType.KEYWORD) # create document_id payload index self._client.create_payload_index( collection_name, Field.DOCUMENT_ID, field_schema=PayloadSchemaType.KEYWORD ) # create full text index text_index_params = TextIndexParams( type=TextIndexType.TEXT, tokenizer=TokenizerType.MULTILINGUAL, min_token_len=2, max_token_len=20, lowercase=True, ) self._client.create_payload_index(collection_name, Field.CONTENT_KEY, field_schema=text_index_params) redis_client.set(collection_exist_cache_key, 1, ex=3600) def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs): uuids = self._get_uuids(documents) texts = [d.page_content for d in documents] metadatas = [d.metadata for d in documents] added_ids = [] # Filter out None values from metadatas list to match expected type filtered_metadatas = [m for m in metadatas if m is not None] for batch_ids, points in self._generate_rest_batches( texts, embeddings, filtered_metadatas, uuids, 64, self._group_id ): self._client.upsert(collection_name=self._collection_name, points=cast("common_types.Points", points)) added_ids.extend(batch_ids) return added_ids def _generate_rest_batches( self, texts: Iterable[str], embeddings: list[list[float]], metadatas: list[dict] | None = None, ids: Sequence[str] | None = None, batch_size: int = 64, group_id: str | None = None, ) -> Generator[tuple[list[str], list[rest.PointStruct]], None, None]: from qdrant_client.http import models as rest texts_iterator = iter(texts) embeddings_iterator = iter(embeddings) metadatas_iterator = iter(metadatas or []) ids_iterator = iter(ids or [uuid.uuid4().hex for _ in iter(texts)]) while batch_texts := list(islice(texts_iterator, batch_size)): # Take the corresponding metadata and id for each text in a batch batch_metadatas = list(islice(metadatas_iterator, batch_size)) or None batch_ids = list(islice(ids_iterator, batch_size)) # Generate the embeddings for all the texts in a batch batch_embeddings = list(islice(embeddings_iterator, batch_size)) points = [ rest.PointStruct( id=point_id, vector=vector, payload=payload, ) for point_id, vector, payload in zip( batch_ids, batch_embeddings, self._build_payloads( batch_texts, batch_metadatas, Field.CONTENT_KEY, Field.METADATA_KEY, group_id or "", # Ensure group_id is never None Field.GROUP_KEY, ), ) ] yield batch_ids, points @classmethod def _build_payloads( cls, texts: Iterable[str], metadatas: list[dict] | None, content_payload_key: str, metadata_payload_key: str, group_id: str, group_payload_key: str, ) -> list[dict]: payloads = [] for i, text in enumerate(texts): if text is None: raise ValueError( "At least one of the texts is None. Please remove it before " "calling .from_texts or .add_texts on Qdrant instance." ) metadata = metadatas[i] if metadatas is not None else None payloads.append({content_payload_key: text, metadata_payload_key: metadata, group_payload_key: group_id}) return payloads def delete_by_metadata_field(self, key: str, value: str): from qdrant_client.http import models from qdrant_client.http.exceptions import UnexpectedResponse try: filter = models.Filter( must=[ models.FieldCondition( key=f"metadata.{key}", match=models.MatchValue(value=value), ), ], ) self._reload_if_needed() self._client.delete( collection_name=self._collection_name, points_selector=FilterSelector(filter=filter), ) except UnexpectedResponse as e: # Collection does not exist, so return if e.status_code == 404: return # Some other error occurred, so re-raise the exception else: raise e def delete(self): from qdrant_client.http import models from qdrant_client.http.exceptions import UnexpectedResponse try: filter = models.Filter( must=[ models.FieldCondition( key="group_id", match=models.MatchValue(value=self._group_id), ), ], ) self._client.delete( collection_name=self._collection_name, points_selector=FilterSelector(filter=filter), ) except UnexpectedResponse as e: # Collection does not exist, so return if e.status_code == 404: return # Some other error occurred, so re-raise the exception else: raise e def delete_by_ids(self, ids: list[str]): from qdrant_client.http import models from qdrant_client.http.exceptions import UnexpectedResponse try: filter = models.Filter( must=[ models.FieldCondition( key="metadata.doc_id", match=models.MatchAny(any=ids), ), ], ) self._client.delete( collection_name=self._collection_name, points_selector=FilterSelector(filter=filter), ) except UnexpectedResponse as e: # Collection does not exist, so return if e.status_code == 404: return # Some other error occurred, so re-raise the exception else: raise e def text_exists(self, id: str) -> bool: all_collection_name = [] collections_response = self._client.get_collections() collection_list = collections_response.collections for collection in collection_list: all_collection_name.append(collection.name) if self._collection_name not in all_collection_name: return False response = self._client.retrieve(collection_name=self._collection_name, ids=[id]) return len(response) > 0 def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: from qdrant_client.http import models score_threshold = float(kwargs.get("score_threshold") or 0.0) if score_threshold >= 1: # return empty list because some versions of qdrant may response with 400 bad request， # and at the same time, the score_threshold with value 1 may be valid for other vector stores return [] filter = models.Filter( must=[ models.FieldCondition( key="group_id", match=models.MatchValue(value=self._group_id), ), ], ) document_ids_filter = kwargs.get("document_ids_filter") if document_ids_filter: if filter.must: filter.must.append( models.FieldCondition( key="metadata.document_id", match=models.MatchAny(any=document_ids_filter), ) ) results = self._client.search( collection_name=self._collection_name, query_vector=query_vector, query_filter=filter, limit=kwargs.get("top_k", 4), with_payload=True, with_vectors=True, score_threshold=score_threshold, ) docs = [] for result in results: if result.payload is None: continue metadata = result.payload.get(Field.METADATA_KEY) or {} # duplicate check score threshold if result.score >= score_threshold: metadata["score"] = result.score doc = Document( page_content=result.payload.get(Field.CONTENT_KEY, ""), metadata=metadata, ) docs.append(doc) # Sort the documents by score in descending order docs = sorted(docs, key=lambda x: x.metadata["score"] if x.metadata is not None else 0, reverse=True) return docs def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: """Return docs most similar by full-text search. Searches each keyword separately and merges results to ensure documents matching ANY keyword are returned (OR logic). Results are capped at top_k. Args: query: Search query text. Multi-word queries are split into keywords, with each keyword searched separately. Limited to 10 keywords. **kwargs: Additional search parameters (top_k, document_ids_filter) Returns: List of up to top_k unique documents matching any query keyword. """ from qdrant_client.http import models # Build base must conditions (AND logic) for metadata filters base_must_conditions: list = [ models.FieldCondition( key="group_id", match=models.MatchValue(value=self._group_id), ), ] document_ids_filter = kwargs.get("document_ids_filter") if document_ids_filter: base_must_conditions.append( models.FieldCondition( key="metadata.document_id", match=models.MatchAny(any=document_ids_filter), ) ) # Split query into keywords, deduplicate and limit to prevent DoS keywords = list(dict.fromkeys(kw.strip() for kw in query.strip().split() if kw.strip()))[:10] if not keywords: return [] top_k = kwargs.get("top_k", 2) seen_ids: set[str | int] = set() documents: list[Document] = [] # Search each keyword separately and merge results. # This ensures each keyword gets its own search, preventing one keyword's # results from completely overshadowing another's due to scroll ordering. for keyword in keywords: scroll_filter = models.Filter( must=[ *base_must_conditions, models.FieldCondition( key="page_content", match=models.MatchText(text=keyword), ), ] ) response = self._client.scroll( collection_name=self._collection_name, scroll_filter=scroll_filter, limit=top_k, with_payload=True, with_vectors=True, ) results = response[0] for result in results: if result and result.id not in seen_ids: seen_ids.add(result.id) document = self._document_from_scored_point(result, Field.CONTENT_KEY, Field.METADATA_KEY) documents.append(document) if len(documents) >= top_k: return documents return documents def _reload_if_needed(self): if isinstance(self._client, QdrantLocal): self._client._load() @classmethod def _document_from_scored_point( cls, scored_point: Any, content_payload_key: str, metadata_payload_key: str, ) -> Document: return Document( page_content=scored_point.payload.get(content_payload_key), vector=scored_point.vector, metadata=scored_point.payload.get(metadata_payload_key) or {}, ) class QdrantVectorFactory(AbstractVectorFactory): def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> QdrantVector: if dataset.collection_binding_id: stmt = select(DatasetCollectionBinding).where(DatasetCollectionBinding.id == dataset.collection_binding_id) dataset_collection_binding = db.session.scalars(stmt).one_or_none() if dataset_collection_binding: collection_name = dataset_collection_binding.collection_name else: raise ValueError("Dataset Collection Bindings does not exist!") else: 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) if not dataset.index_struct_dict: dataset.index_struct = json.dumps(self.gen_index_struct_dict(VectorType.QDRANT, collection_name)) return QdrantVector( collection_name=collection_name, group_id=dataset.id, config=QdrantConfig( endpoint=dify_config.QDRANT_URL or "", api_key=dify_config.QDRANT_API_KEY, root_path=str(current_app.config.root_path), timeout=dify_config.QDRANT_CLIENT_TIMEOUT, grpc_port=dify_config.QDRANT_GRPC_PORT, prefer_grpc=dify_config.QDRANT_GRPC_ENABLED, replication_factor=dify_config.QDRANT_REPLICATION_FACTOR, ), )