# 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 hashlib
import os
import re
import uuid
from typing import Any, Callable, Literal, Optional, Union

from ...code_utils import extract_code
from ...formatting_utils import colored
from ...import_utils import optional_import_block, require_optional_import
from ...retrieve_utils import (
    TEXT_FORMATS,
    create_vector_db_from_dir,
    get_files_from_dir,
    query_vector_db,
    split_files_to_chunks,
)
from ...token_count_utils import count_token
from .. import UserProxyAgent
from ..agent import Agent
from ..contrib.vectordb.base import Document, QueryResults, VectorDB, VectorDBFactory
from ..contrib.vectordb.utils import (
    chroma_results_to_query_results,
    filter_results_by_distance,
    get_logger,
)

with optional_import_block():
    import chromadb
    from IPython import get_ipython

logger = get_logger(__name__)

PROMPT_DEFAULT = """You're a retrieve augmented chatbot. You answer user's questions based on your own knowledge and the
context provided by the user. You should follow the following steps to answer a question:
Step 1, you estimate the user's intent based on the question and context. The intent can be a code generation task or
a question answering task.
Step 2, you reply based on the intent.
If you can't answer the question with or without the current context, you should reply exactly `UPDATE CONTEXT`.
If user's intent is code generation, you must obey the following rules:
Rule 1. You MUST NOT install any packages because all the packages needed are already installed.
Rule 2. You must follow the formats below to write your code:
```language
# your code
```

If user's intent is question answering, you must give as short an answer as possible.

User's question is: {input_question}

Context is: {input_context}

The source of the context is: {input_sources}

If you can answer the question, in the end of your answer, add the source of the context in the format of `Sources: source1, source2, ...`.
"""

PROMPT_CODE = """You're a retrieve augmented coding assistant. You answer user's questions based on your own knowledge and the
context provided by the user.
If you can't answer the question with or without the current context, you should reply exactly `UPDATE CONTEXT`.
For code generation, you must obey the following rules:
Rule 1. You MUST NOT install any packages because all the packages needed are already installed.
Rule 2. You must follow the formats below to write your code:
```language
# your code
```

User's question is: {input_question}

Context is: {input_context}
"""

PROMPT_QA = """You're a retrieve augmented chatbot. You answer user's questions based on your own knowledge and the
context provided by the user.
If you can't answer the question with or without the current context, you should reply exactly `UPDATE CONTEXT`.
You must give as short an answer as possible.

User's question is: {input_question}

Context is: {input_context}
"""

HASH_LENGTH = int(os.environ.get("HASH_LENGTH", 8))
UPDATE_CONTEXT_IN_PROMPT = "you should reply exactly `UPDATE CONTEXT`"


@require_optional_import(["chromadb", "IPython"], "retrievechat")
class RetrieveUserProxyAgent(UserProxyAgent):
    """(In preview) The Retrieval-Augmented User Proxy retrieves document chunks based on the embedding
    similarity, and sends them along with the question to the Retrieval-Augmented Assistant
    """

    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,
    ):
        r"""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, and provides
                     source information in the end of the response.
                - `vector_db` (Optional, Union[str, VectorDB]) - the vector db for the retrieve chat.
                    If it's a string, it should be the type of the vector db, such as "chroma"; otherwise,
                    it should be an instance of the VectorDB protocol. Default is "chroma".
                    Set `None` to use the deprecated `client`.
                - `db_config` (Optional, Dict) - the config for the vector db. Default is `{}`. Please make
                    sure you understand the config for the vector db you are using, otherwise, leave it as `{}`.
                    Only valid when `vector_db` is a string.
                - `client` (Optional, chromadb.Client) - the chromadb client. If key not provided, a
                     default client `chromadb.Client()` will be used. If you want to use other
                     vector db, extend this class and override the `retrieve_docs` function.
                     *[Deprecated]* use `vector_db` instead.
                - `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.
                    *[Deprecated]* use `new_docs` when use `vector_db` instead of `client`.
                - `new_docs` (Optional, bool) - when True, only adds new documents to the collection;
                    when False, updates existing documents and adds new ones. Default is True.
                    Document id is used to determine if a document is new or existing. By default, the
                    id is the hash value of the content.
                - `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 `all-MiniLM-L6-v2` will be used. All available
                    models can be found at `https://www.sbert.net/docs/sentence_transformer/pretrained_models.html`.
                    The default model is a fast model. If you want to use a high performance model,
                    `all-mpnet-base-v2` is recommended.
                    *[Deprecated]* no need when use `vector_db` instead of `client`.
                - `embedding_function` (Optional, Callable) - the embedding function for creating the
                    vector db. Default is None, SentenceTransformer with the given `embedding_model`
                    will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
                    functions, you can pass it here,
                    follow the examples in `https://docs.trychroma.com/embeddings`.
                - `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.
                - `collection_name` (Optional, str) - the name of the collection.
                    If key not provided, a default name `ag2-docs` will be used.
                - `get_or_create` (Optional, bool) - Whether to get the collection if it exists. Default is False.
                - `overwrite` (Optional, bool) - Whether to overwrite the collection if it exists. Default is False.
                    Case 1. if the collection does not exist, create the collection.
                    Case 2. the collection exists, if overwrite is True, it will overwrite the collection.
                    Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection,
                        otherwise it raise a ValueError.
                - `custom_token_count_function` (Optional, Callable) - a custom function to count the
                    number of tokens in a string.
                    The function should take (text:str, model:str) as input and return the
                    token_count(int). the retrieve_config["model"] will be passed in the function.
                    Default is autogen.token_count_utils.count_token that uses tiktoken, which 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.
                - `distance_threshold` (Optional, float) - the threshold for the distance score, only
                    distance smaller than it will be returned. Will be ignored if < 0. Default is -1.

            `**kwargs` (dict): other kwargs in [UserProxyAgent](/docs/api-reference/autogen/UserProxyAgent#userproxyagent).

        Example:
        Example of overriding retrieve_docs - If you have set up a customized vector db, and it's
        not compatible with chromadb, you can easily plug in it with below code.
        *[Deprecated]* use `vector_db` instead. You can extend VectorDB and pass it to the agent.
        ```python
        class MyRetrieveUserProxyAgent(RetrieveUserProxyAgent):
            def query_vector_db(
                self,
                query_texts: List[str],
                n_results: int = 10,
                search_string: str = "",
                **kwargs: Any,
            ) -> Dict[str, Union[List[str], List[List[str]]]]:
                # define your own query function here
                pass

            def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = "", **kwargs):
                results = self.query_vector_db(
                    query_texts=[problem],
                    n_results=n_results,
                    search_string=search_string,
                    **kwargs: Any,
                )

                self._results = results
                print("doc_ids: ", results["ids"])
        ```
        """
        super().__init__(
            name=name,
            human_input_mode=human_input_mode,
            **kwargs,
        )

        self._retrieve_config = {} if retrieve_config is None else retrieve_config
        self._task = self._retrieve_config.get("task", "default")
        self._vector_db = self._retrieve_config.get("vector_db", "chroma")
        self._db_config = self._retrieve_config.get("db_config", {})
        self._docs_path = self._retrieve_config.get("docs_path", None)
        self._extra_docs = self._retrieve_config.get("extra_docs", False)
        self._new_docs = self._retrieve_config.get("new_docs", True)
        self._collection_name = self._retrieve_config.get("collection_name", "ag2-docs")
        if "docs_path" not in self._retrieve_config:
            logger.warning(
                "docs_path is not provided in retrieve_config. "
                f"Will raise ValueError if the collection `{self._collection_name}` doesn't exist. "
                "Set docs_path to None to suppress this warning."
            )
        self._model = self._retrieve_config.get("model", "gpt-4")
        self._max_tokens = self.get_max_tokens(self._model)
        self._chunk_token_size = int(self._retrieve_config.get("chunk_token_size", self._max_tokens * 0.4))
        self._chunk_mode = self._retrieve_config.get("chunk_mode", "multi_lines")
        self._must_break_at_empty_line = self._retrieve_config.get("must_break_at_empty_line", True)
        self._embedding_model = self._retrieve_config.get("embedding_model", "all-MiniLM-L6-v2")
        self._embedding_function = self._retrieve_config.get("embedding_function", None)
        self.customized_prompt = self._retrieve_config.get("customized_prompt", None)
        self.customized_answer_prefix = self._retrieve_config.get("customized_answer_prefix", "").upper()
        self.update_context = self._retrieve_config.get("update_context", True)
        self._get_or_create = self._retrieve_config.get("get_or_create", False) if self._docs_path is not None else True
        self._overwrite = self._retrieve_config.get("overwrite", False)
        self.custom_token_count_function = self._retrieve_config.get("custom_token_count_function", count_token)
        self.custom_text_split_function = self._retrieve_config.get("custom_text_split_function", None)
        self._custom_text_types = self._retrieve_config.get("custom_text_types", TEXT_FORMATS)
        self._recursive = self._retrieve_config.get("recursive", True)
        self._context_max_tokens = self._retrieve_config.get("context_max_tokens", self._max_tokens * 0.8)
        self._collection = self._docs_path is None  # whether the collection is created
        self._ipython = get_ipython()
        self._doc_idx = -1  # the index of the current used doc
        self._results = []  # the results of the current query
        self._intermediate_answers = set()  # the intermediate answers
        self._doc_contents = []  # the contents of the current used doc
        self._doc_ids = []  # the ids of the current used doc
        self._current_docs_in_context = []  # the ids of the current context sources
        self._search_string = ""  # the search string used in the current query
        self._distance_threshold = self._retrieve_config.get("distance_threshold", -1)
        # update the termination message function
        self._is_termination_msg = (
            self._is_termination_msg_retrievechat if is_termination_msg is None else is_termination_msg
        )
        if isinstance(self._vector_db, str):
            if not isinstance(self._db_config, dict):
                raise ValueError("`db_config` should be a dictionary.")
            if "embedding_function" in self._retrieve_config:
                self._db_config["embedding_function"] = self._embedding_function
            self._vector_db = VectorDBFactory.create_vector_db(db_type=self._vector_db, **self._db_config)
        self._client = self._retrieve_config.get("client", None)
        if self._client is None and hasattr(self._vector_db, "client"):
            # Since the client arg is deprecated, let's check
            # if the `vector_db` instance has a 'client' attribute.
            self._client = getattr(self._vector_db, "client", None)
        if self._client is None:
            self._client = chromadb.Client()
        self.register_reply(Agent, RetrieveUserProxyAgent._generate_retrieve_user_reply, position=2)
        self.register_hook(
            hookable_method="process_message_before_send",
            hook=self._check_update_context_before_send,
        )

    def _init_db(self):
        if not self._vector_db:
            return

        is_to_chunk = False  # whether to chunk the raw files
        if self._new_docs:
            is_to_chunk = True
        if not self._docs_path:
            try:
                self._vector_db.get_collection(self._collection_name)
                logger.warning(f"`docs_path` is not provided. Use the existing collection `{self._collection_name}`.")
                self._overwrite = False
                self._get_or_create = True
                is_to_chunk = False
            except ValueError:
                raise ValueError(
                    "`docs_path` is not provided. "
                    f"The collection `{self._collection_name}` doesn't exist either. "
                    "Please provide `docs_path` or create the collection first."
                )
        elif self._get_or_create and not self._overwrite:
            try:
                self._vector_db.get_collection(self._collection_name)
                logger.info(f"Use the existing collection `{self._collection_name}`.", color="green")
            except ValueError:
                is_to_chunk = True
        else:
            is_to_chunk = True

        self._vector_db.active_collection = self._vector_db.create_collection(
            self._collection_name, overwrite=self._overwrite, get_or_create=self._get_or_create
        )

        docs = None
        if is_to_chunk:
            if self.custom_text_split_function is not None:
                chunks, sources = split_files_to_chunks(
                    get_files_from_dir(self._docs_path, self._custom_text_types, self._recursive),
                    custom_text_split_function=self.custom_text_split_function,
                )
            else:
                chunks, sources = split_files_to_chunks(
                    get_files_from_dir(self._docs_path, self._custom_text_types, self._recursive),
                    self._chunk_token_size,
                    self._chunk_mode,
                    self._must_break_at_empty_line,
                )
            logger.info(f"Found {len(chunks)} chunks.")

            if self._new_docs:
                all_docs_ids = {
                    doc["id"]
                    for doc in self._vector_db.get_docs_by_ids(ids=None, collection_name=self._collection_name)
                }
            else:
                all_docs_ids = set()

            chunk_ids = (
                [hashlib.blake2b(chunk.encode("utf-8")).hexdigest()[:HASH_LENGTH] for chunk in chunks]
                if self._vector_db.type != "qdrant"
                else [str(uuid.UUID(hex=hashlib.md5(chunk.encode("utf-8")).hexdigest())) for chunk in chunks]
            )
            chunk_ids_set = set(chunk_ids)
            chunk_ids_set_idx = [chunk_ids.index(hash_value) for hash_value in chunk_ids_set]
            docs = [
                Document(id=chunk_ids[idx], content=chunks[idx], metadata=sources[idx])
                for idx in chunk_ids_set_idx
                if chunk_ids[idx] not in all_docs_ids
            ]

        self._vector_db.insert_docs(docs=docs, collection_name=self._collection_name, upsert=True)

    def _is_termination_msg_retrievechat(self, message):
        """Check if a message is a termination message.
        For code generation, terminate when no code block is detected. Currently only detect python code blocks.
        For question answering, terminate when don't update context, i.e., answer is given.
        """
        if isinstance(message, dict):
            message = message.get("content")
            if message is None:
                return False
        cb = extract_code(message)
        contain_code = False
        for c in cb:
            # todo: support more languages
            if c[0] == "python":
                contain_code = True
                break
        update_context_case1, update_context_case2 = self._check_update_context(message)
        return not (contain_code or update_context_case1 or update_context_case2)

    def _check_update_context_before_send(self, sender, message, recipient, silent):
        if not isinstance(message, (str, dict)):
            return message
        elif isinstance(message, dict):
            msg_text = message.get("content", message)
        else:
            msg_text = message

        if msg_text.strip().upper() == "UPDATE CONTEXT":
            doc_contents = self._get_context(self._results)

            # Always use self.problem as the query text to retrieve docs, but each time we replace the context with the
            # next similar docs in the retrieved doc results.
            if not doc_contents:
                for _tmp_retrieve_count in range(1, 5):
                    self._reset(intermediate=True)
                    self.retrieve_docs(
                        self.problem, self.n_results * (2 * _tmp_retrieve_count + 1), self._search_string
                    )
                    doc_contents = self._get_context(self._results)
                    if doc_contents or self.n_results * (2 * _tmp_retrieve_count + 1) >= len(self._results[0]):
                        break
            msg_text = self._generate_message(doc_contents, task=self._task)

        if isinstance(message, dict):
            message["content"] = msg_text
        return message

    @staticmethod
    def get_max_tokens(model="gpt-3.5-turbo"):
        if "32k" in model:
            return 32000
        elif "16k" in model:
            return 16000
        elif "gpt-4" in model:
            return 8000
        else:
            return 4000

    def _reset(self, intermediate=False):
        self._doc_idx = -1  # the index of the current used doc
        self._results = []  # the results of the current query
        if not intermediate:
            self._intermediate_answers = set()  # the intermediate answers
            self._doc_contents = []  # the contents of the current used doc
            self._doc_ids = []  # the ids of the current used doc

    def _get_context(self, results: QueryResults):
        doc_contents = ""
        self._current_docs_in_context = []
        current_tokens = 0
        _doc_idx = self._doc_idx
        _tmp_retrieve_count = 0
        for idx, doc in enumerate(results[0]):
            doc = doc[0]
            if idx <= _doc_idx:
                continue
            if doc["id"] in self._doc_ids:
                continue
            _doc_tokens = self.custom_token_count_function(doc["content"], self._model)
            if _doc_tokens > self._context_max_tokens:
                func_print = f"Skip doc_id {doc['id']} as it is too long to fit in the context."
                print(colored(func_print, "green"), flush=True)
                self._doc_idx = idx
                continue
            if current_tokens + _doc_tokens > self._context_max_tokens:
                break
            func_print = f"Adding content of doc {doc['id']} to context."
            print(colored(func_print, "green"), flush=True)
            current_tokens += _doc_tokens
            doc_contents += doc["content"] + "\n"
            _metadata = doc.get("metadata")
            if isinstance(_metadata, dict):
                self._current_docs_in_context.append(_metadata.get("source", ""))
            self._doc_idx = idx
            self._doc_ids.append(doc["id"])
            self._doc_contents.append(doc["content"])
            _tmp_retrieve_count += 1
            if _tmp_retrieve_count >= self.n_results:
                break
        return doc_contents

    def _generate_message(self, doc_contents, task="default"):
        if not doc_contents:
            print(colored("No more context, will terminate.", "green"), flush=True)
            return "TERMINATE"
        if self.customized_prompt:
            message = self.customized_prompt.format(input_question=self.problem, input_context=doc_contents)
        elif task.upper() == "CODE":
            message = PROMPT_CODE.format(input_question=self.problem, input_context=doc_contents)
        elif task.upper() == "QA":
            message = PROMPT_QA.format(input_question=self.problem, input_context=doc_contents)
        elif task.upper() == "DEFAULT":
            message = PROMPT_DEFAULT.format(
                input_question=self.problem, input_context=doc_contents, input_sources=self._current_docs_in_context
            )
        else:
            raise NotImplementedError(f"task {task} is not implemented.")
        return message

    def _check_update_context(self, message):
        if isinstance(message, dict):
            message = message.get("content", "")
        elif not isinstance(message, str):
            message = ""
        update_context_case1 = "UPDATE CONTEXT" in message.upper() and UPDATE_CONTEXT_IN_PROMPT not in message
        update_context_case2 = self.customized_answer_prefix and self.customized_answer_prefix not in message.upper()
        return update_context_case1, update_context_case2

    def _generate_retrieve_user_reply(
        self,
        messages: Optional[list[dict[str, Any]]] = None,
        sender: Optional[Agent] = None,
        config: Optional[Any] = None,
    ) -> tuple[bool, Optional[Union[str, dict[str, Any]]]]:
        """In this function, we will update the context and reset the conversation based on different conditions.
        We'll update the context and reset the conversation if update_context is True and either of the following:
        (1) the last message contains "UPDATE CONTEXT",
        (2) the last message doesn't contain "UPDATE CONTEXT" and the customized_answer_prefix is not in the message.
        """
        if config is None:
            config = self
        if messages is None:
            messages = self._oai_messages[sender]
        message = messages[-1]
        update_context_case1, update_context_case2 = self._check_update_context(message)
        if (update_context_case1 or update_context_case2) and self.update_context:
            print(colored("Updating context and resetting conversation.", "green"), flush=True)
            # extract the first sentence in the response as the intermediate answer
            _message = message.get("content", "").split("\n")[0].strip()
            _intermediate_info = re.split(r"(?<=[.!?])\s+", _message)
            self._intermediate_answers.add(_intermediate_info[0])

            if update_context_case1:
                # try to get more context from the current retrieved doc results because the results may be too long to fit
                # in the LLM context.
                doc_contents = self._get_context(self._results)

                # Always use self.problem as the query text to retrieve docs, but each time we replace the context with the
                # next similar docs in the retrieved doc results.
                if not doc_contents:
                    for _tmp_retrieve_count in range(1, 5):
                        self._reset(intermediate=True)
                        self.retrieve_docs(
                            self.problem, self.n_results * (2 * _tmp_retrieve_count + 1), self._search_string
                        )
                        doc_contents = self._get_context(self._results)
                        if doc_contents or self.n_results * (2 * _tmp_retrieve_count + 1) >= len(self._results[0]):
                            break
            elif update_context_case2:
                # Use the current intermediate info as the query text to retrieve docs, and each time we append the top similar
                # docs in the retrieved doc results to the context.
                for _tmp_retrieve_count in range(5):
                    self._reset(intermediate=True)
                    self.retrieve_docs(
                        _intermediate_info[0], self.n_results * (2 * _tmp_retrieve_count + 1), self._search_string
                    )
                    self._get_context(self._results)
                    doc_contents = "\n".join(self._doc_contents)  # + "\n" + "\n".join(self._intermediate_answers)
                    if doc_contents or self.n_results * (2 * _tmp_retrieve_count + 1) >= len(self._results[0]):
                        break

            self.clear_history()
            sender.clear_history()
            return True, self._generate_message(doc_contents, task=self._task)
        else:
            return False, None

    def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
        """Retrieve docs based on the given problem and assign the results to the class property `_results`.
        The retrieved docs should be type of `QueryResults` which is a list of tuples containing the document and
        the distance.

        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 "".
                Not used if the vector_db doesn't support it.

        Returns:
            None.
        """
        if isinstance(self._vector_db, VectorDB):
            if not self._collection or not self._get_or_create:
                print("Trying to create collection.")
                self._init_db()
                self._collection = True
                self._get_or_create = True

            kwargs = {}
            if hasattr(self._vector_db, "type") and self._vector_db.type == "chroma":
                kwargs["where_document"] = {"$contains": search_string} if search_string else None
            results = self._vector_db.retrieve_docs(
                queries=[problem],
                n_results=n_results,
                collection_name=self._collection_name,
                distance_threshold=self._distance_threshold,
                **kwargs,
            )
            self._search_string = search_string
            self._results = results
            print("VectorDB returns doc_ids: ", [[r[0]["id"] for r in rr] for rr in results])
            return

        if not self._collection or not self._get_or_create:
            print("Trying to create collection.")
            self._client = create_vector_db_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,
                get_or_create=self._get_or_create,
                embedding_function=self._embedding_function,
                custom_text_split_function=self.custom_text_split_function,
                custom_text_types=self._custom_text_types,
                recursive=self._recursive,
                extra_docs=self._extra_docs,
            )
            self._collection = True
            self._get_or_create = True

        results = query_vector_db(
            query_texts=[problem],
            n_results=n_results,
            search_string=search_string,
            client=self._client,
            collection_name=self._collection_name,
            embedding_model=self._embedding_model,
            embedding_function=self._embedding_function,
        )
        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
        print("doc_ids: ", [[r[0]["id"] for r in rr] for rr in results])

    @staticmethod
    def message_generator(sender, recipient, context):
        """Generate an initial message with the given context for the RetrieveUserProxyAgent.

        Args:
            sender (Agent): the sender agent. It should be the instance of RetrieveUserProxyAgent.
            recipient (Agent): the recipient agent. Usually it's the assistant agent.
            context (dict): the context for the message generation. It should contain the following keys:
                - `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 "".

        Returns:
            str: the generated message ready to be sent to the recipient agent.
        """
        sender._reset()

        problem = context.get("problem", "")
        n_results = context.get("n_results", 20)
        search_string = context.get("search_string", "")

        sender.retrieve_docs(problem, n_results, search_string)
        sender.problem = problem
        sender.n_results = n_results
        doc_contents = sender._get_context(sender._results)
        message = sender._generate_message(doc_contents, sender._task)
        return message

    def run_code(self, code, **kwargs):
        lang = kwargs.get("lang")
        if code.startswith("!") or code.startswith("pip") or lang in ["bash", "shell", "sh"]:
            return (
                0,
                "You MUST NOT install any packages because all the packages needed are already installed.",
                None,
            )
        if self._ipython is None or lang != "python":
            return super().run_code(code, **kwargs)
        else:
            result = self._ipython.run_cell(code)
            log = str(result.result)
            exitcode = 0 if result.success else 1
            if result.error_before_exec is not None:
                log += f"\n{result.error_before_exec}"
                exitcode = 1
            if result.error_in_exec is not None:
                log += f"\n{result.error_in_exec}"
                exitcode = 1
            return exitcode, log, None
