# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors # # SPDX-License-Identifier: Apache-2.0 # # Portions derived from https://github.com/microsoft/autogen are under the MIT License. # SPDX-License-Identifier: MIT import warnings from typing import Any, Callable, Literal, Optional from ...import_utils import optional_import_block, require_optional_import from ...retrieve_utils import TEXT_FORMATS, get_files_from_dir, split_files_to_chunks from ..contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent from ..contrib.vectordb.utils import ( chroma_results_to_query_results, filter_results_by_distance, get_logger, ) logger = get_logger(__name__) with optional_import_block(): import fastembed # noqa: F401 from qdrant_client import QdrantClient, models from qdrant_client.fastembed_common import QueryResponse @require_optional_import(["fastembed", "qdrant_client"], "retrievechat-qdrant") class QdrantRetrieveUserProxyAgent(RetrieveUserProxyAgent): def __init__( self, name="RetrieveChatAgent", # default set to RetrieveChatAgent human_input_mode: Literal["ALWAYS", "NEVER", "TERMINATE"] = "ALWAYS", is_termination_msg: Optional[Callable[[dict[str, Any]], bool]] = None, retrieve_config: Optional[dict[str, Any]] = None, # config for the retrieve agent **kwargs: Any, ): """Args: name (str): name of the agent. human_input_mode (str): whether to ask for human inputs every time a message is received. Possible values are "ALWAYS", "TERMINATE", "NEVER". 1. When "ALWAYS", the agent prompts for human input every time a message is received. Under this mode, the conversation stops when the human input is "exit", or when is_termination_msg is True and there is no human input. 2. When "TERMINATE", the agent only prompts for human input only when a termination message is received or the number of auto reply reaches the max_consecutive_auto_reply. 3. When "NEVER", the agent will never prompt for human input. Under this mode, the conversation stops when the number of auto reply reaches the max_consecutive_auto_reply or when is_termination_msg is True. is_termination_msg (function): a function that takes a message in the form of a dictionary and returns a boolean value indicating if this received message is a termination message. The dict can contain the following keys: "content", "role", "name", "function_call". retrieve_config (dict or None): config for the retrieve agent. To use default config, set to None. Otherwise, set to a dictionary with the following keys: - task (Optional, str): the task of the retrieve chat. Possible values are "code", "qa" and "default". System prompt will be different for different tasks. The default value is `default`, which supports both code and qa. - client (Optional, qdrant_client.QdrantClient(":memory:")): A QdrantClient instance. If not provided, an in-memory instance will be assigned. Not recommended for production. will be used. If you want to use other vector db, extend this class and override the `retrieve_docs` function. - docs_path (Optional, Union[str, List[str]]): the path to the docs directory. It can also be the path to a single file, the url to a single file or a list of directories, files and urls. Default is None, which works only if the collection is already created. - extra_docs (Optional, bool): when true, allows adding documents with unique IDs without overwriting existing ones; when false, it replaces existing documents using default IDs, risking collection overwrite., when set to true it enables the system to assign unique IDs starting from "length+i" for new document chunks, preventing the replacement of existing documents and facilitating the addition of more content to the collection.. By default, "extra_docs" is set to false, starting document IDs from zero. This poses a risk as new documents might overwrite existing ones, potentially causing unintended loss or alteration of data in the collection. - collection_name (Optional, str): the name of the collection. If key not provided, a default name `autogen-docs` will be used. - model (Optional, str): the model to use for the retrieve chat. If key not provided, a default model `gpt-4` will be used. - chunk_token_size (Optional, int): the chunk token size for the retrieve chat. If key not provided, a default size `max_tokens * 0.4` will be used. - context_max_tokens (Optional, int): the context max token size for the retrieve chat. If key not provided, a default size `max_tokens * 0.8` will be used. - chunk_mode (Optional, str): the chunk mode for the retrieve chat. Possible values are "multi_lines" and "one_line". If key not provided, a default mode `multi_lines` will be used. - must_break_at_empty_line (Optional, bool): chunk will only break at empty line if True. Default is True. If chunk_mode is "one_line", this parameter will be ignored. - embedding_model (Optional, str): the embedding model to use for the retrieve chat. If key not provided, a default model `BAAI/bge-small-en-v1.5` will be used. All available models can be found at `https://qdrant.github.io/fastembed/examples/Supported_Models/`. - customized_prompt (Optional, str): the customized prompt for the retrieve chat. Default is None. - customized_answer_prefix (Optional, str): the customized answer prefix for the retrieve chat. Default is "". If not "" and the customized_answer_prefix is not in the answer, `Update Context` will be triggered. - update_context (Optional, bool): if False, will not apply `Update Context` for interactive retrieval. Default is True. - custom_token_count_function (Optional, Callable): a custom function to count the number of tokens in a string. The function should take a string as input and return three integers (token_count, tokens_per_message, tokens_per_name). Default is None, tiktoken will be used and may not be accurate for non-OpenAI models. - custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings. Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`. - custom_text_types (Optional, List[str]): a list of file types to be processed. Default is `autogen.retrieve_utils.TEXT_FORMATS`. This only applies to files under the directories in `docs_path`. Explicitly included files and urls will be chunked regardless of their types. - recursive (Optional, bool): whether to search documents recursively in the docs_path. Default is True. - parallel (Optional, int): How many parallel workers to use for embedding. Defaults to the number of CPU cores. - on_disk (Optional, bool): Whether to store the collection on disk. Default is False. - quantization_config: Quantization configuration. If None, quantization will be disabled. - hnsw_config: HNSW configuration. If None, default configuration will be used. You can find more info about the hnsw configuration options at https://qdrant.tech/documentation/concepts/indexing/#vector-index. API Reference: https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_collection - payload_indexing: Whether to create a payload index for the document field. Default is False. You can find more info about the payload indexing options at https://qdrant.tech/documentation/concepts/indexing/#payload-index API Reference: https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_field_index **kwargs (dict): other kwargs in [UserProxyAgent](/docs/api-reference/autogen/UserProxyAgent#userproxyagent). """ warnings.warn( "The QdrantRetrieveUserProxyAgent is deprecated. Please use the RetrieveUserProxyAgent instead, set `vector_db` to `qdrant`.", DeprecationWarning, stacklevel=2, ) super().__init__(name, human_input_mode, is_termination_msg, retrieve_config, **kwargs) self._client = self._retrieve_config.get("client", QdrantClient(":memory:")) self._embedding_model = self._retrieve_config.get("embedding_model", "BAAI/bge-small-en-v1.5") # Uses all available CPU cores to encode data when set to 0 self._parallel = self._retrieve_config.get("parallel", 0) self._on_disk = self._retrieve_config.get("on_disk", False) self._quantization_config = self._retrieve_config.get("quantization_config", None) self._hnsw_config = self._retrieve_config.get("hnsw_config", None) self._payload_indexing = self._retrieve_config.get("payload_indexing", False) def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""): """Args: problem (str): the problem to be solved. n_results (int): the number of results to be retrieved. Default is 20. search_string (str): only docs that contain an exact match of this string will be retrieved. Default is "". """ if not self._collection: print("Trying to create collection.") create_qdrant_from_dir( dir_path=self._docs_path, max_tokens=self._chunk_token_size, client=self._client, collection_name=self._collection_name, chunk_mode=self._chunk_mode, must_break_at_empty_line=self._must_break_at_empty_line, embedding_model=self._embedding_model, custom_text_split_function=self.custom_text_split_function, custom_text_types=self._custom_text_types, recursive=self._recursive, extra_docs=self._extra_docs, parallel=self._parallel, on_disk=self._on_disk, quantization_config=self._quantization_config, hnsw_config=self._hnsw_config, payload_indexing=self._payload_indexing, ) self._collection = True results = query_qdrant( query_texts=problem, n_results=n_results, search_string=search_string, client=self._client, collection_name=self._collection_name, embedding_model=self._embedding_model, ) results["contents"] = results.pop("documents") results = chroma_results_to_query_results(results, "distances") results = filter_results_by_distance(results, self._distance_threshold) self._search_string = search_string self._results = results @require_optional_import(["fastembed", "qdrant_client"], "retrievechat-qdrant") def create_qdrant_from_dir( dir_path: str, max_tokens: int = 4000, client: "QdrantClient" = None, collection_name: str = "all-my-documents", chunk_mode: str = "multi_lines", must_break_at_empty_line: bool = True, embedding_model: str = "BAAI/bge-small-en-v1.5", custom_text_split_function: Callable = None, custom_text_types: list[str] = TEXT_FORMATS, recursive: bool = True, extra_docs: bool = False, parallel: int = 0, on_disk: bool = False, quantization_config: Optional["models.QuantizationConfig"] = None, hnsw_config: Optional["models.HnswConfigDiff"] = None, payload_indexing: bool = False, qdrant_client_options: Optional[dict[str, Any]] = {}, ): """Create a Qdrant collection from all the files in a given directory, the directory can also be a single file or a url to a single file. Args: dir_path (str): the path to the directory, file or url. max_tokens (Optional, int): the maximum number of tokens per chunk. Default is 4000. client (Optional, QdrantClient): the QdrantClient instance. Default is None. collection_name (Optional, str): the name of the collection. Default is "all-my-documents". chunk_mode (Optional, str): the chunk mode. Default is "multi_lines". must_break_at_empty_line (Optional, bool): Whether to break at empty line. Default is True. embedding_model (Optional, str): the embedding model to use. Default is "BAAI/bge-small-en-v1.5". The list of all the available models can be at https://qdrant.github.io/fastembed/examples/Supported_Models/. custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings. Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`. custom_text_types (Optional, List[str]): a list of file types to be processed. Default is TEXT_FORMATS. recursive (Optional, bool): whether to search documents recursively in the dir_path. Default is True. extra_docs (Optional, bool): whether to add more documents in the collection. Default is False parallel (Optional, int): How many parallel workers to use for embedding. Defaults to the number of CPU cores on_disk (Optional, bool): Whether to store the collection on disk. Default is False. quantization_config: Quantization configuration. If None, quantization will be disabled. Ref: https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_collection hnsw_config: HNSW configuration. If None, default configuration will be used. Ref: https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_collection payload_indexing: Whether to create a payload index for the document field. Default is False. qdrant_client_options: (Optional, dict): the options for instantiating the qdrant client. Ref: https://github.com/qdrant/qdrant-client/blob/master/qdrant_client/qdrant_client.py#L36-L58. """ if client is None: client = QdrantClient(**qdrant_client_options) client.set_model(embedding_model) if custom_text_split_function is not None: chunks, _ = split_files_to_chunks( get_files_from_dir(dir_path, custom_text_types, recursive), custom_text_split_function=custom_text_split_function, ) else: chunks, _ = split_files_to_chunks( get_files_from_dir(dir_path, custom_text_types, recursive), max_tokens, chunk_mode, must_break_at_empty_line ) logger.info(f"Found {len(chunks)} chunks.") collection = None # Check if collection by same name exists, if not, create it with custom options try: collection = client.get_collection(collection_name=collection_name) except Exception: client.create_collection( collection_name=collection_name, vectors_config=client.get_fastembed_vector_params( on_disk=on_disk, quantization_config=quantization_config, hnsw_config=hnsw_config ), ) collection = client.get_collection(collection_name=collection_name) length = 0 if extra_docs: length = len(collection.get()["ids"]) # Upsert in batch of 100 or less if the total number of chunks is less than 100 for i in range(0, len(chunks), min(100, len(chunks))): end_idx = i + min(100, len(chunks) - i) client.add( collection_name, documents=chunks[i:end_idx], ids=[(j + length) for j in range(i, end_idx)], parallel=parallel, ) # Create a payload index for the document field # Enables highly efficient payload filtering. Reference: https://qdrant.tech/documentation/concepts/indexing/#indexing # Creating an index requires additional computational resources and memory. # If filtering performance is critical, we can consider creating an index. if payload_indexing: client.create_payload_index( collection_name=collection_name, field_name="document", field_schema=models.TextIndexParams( type="text", tokenizer=models.TokenizerType.WORD, min_token_len=2, max_token_len=15, ), ) @require_optional_import("qdrant_client", "retrievechat-qdrant") def query_qdrant( query_texts: list[str], n_results: int = 10, client: "QdrantClient" = None, collection_name: str = "all-my-documents", search_string: str = "", embedding_model: str = "BAAI/bge-small-en-v1.5", qdrant_client_options: Optional[dict[str, Any]] = {}, ) -> list[list["QueryResponse"]]: """Perform a similarity search with filters on a Qdrant collection Args: query_texts (List[str]): the query texts. n_results (Optional, int): the number of results to return. Default is 10. client (Optional, API): the QdrantClient instance. A default in-memory client will be instantiated if None. collection_name (Optional, str): the name of the collection. Default is "all-my-documents". search_string (Optional, str): the search string. Default is "". embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if embedding_function is not None. qdrant_client_options: (Optional, dict): the options for instantiating the qdrant client. Reference: https://github.com/qdrant/qdrant-client/blob/master/qdrant_client/qdrant_client.py#L36-L58. Returns: list[list[QueryResponse]]: the query result. The format is: class QueryResponse(BaseModel, extra="forbid"): # type: ignore id: Union[str, int] embedding: Optional[List[float]] metadata: Dict[str, Any] document: str score: float """ if client is None: client = QdrantClient(**qdrant_client_options) client.set_model(embedding_model) results = client.query_batch( collection_name, query_texts, limit=n_results, query_filter=( models.Filter( must=[ models.FieldCondition( key="document", match=models.MatchText(text=search_string), ) ] ) if search_string else None ), ) data = { "ids": [[result.id for result in sublist] for sublist in results], "documents": [[result.document for result in sublist] for sublist in results], "distances": [[result.score for result in sublist] for sublist in results], "metadatas": [[result.metadata for result in sublist] for sublist in results], } return data