Skip to main content

agentchat.contrib.retrieve_user_proxy_agent

RetrieveUserProxyAgent​

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

__init__​

def __init__(name="RetrieveChatAgent",
             human_input_mode: Literal["ALWAYS", "NEVER",
                                       "TERMINATE"] = "ALWAYS",
             is_termination_msg: Optional[Callable[[Dict], bool]] = None,
             retrieve_config: Optional[Dict] = None,
             **kwargs)

Arguments:

  • 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/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/guides/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 autogen-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.

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.

class MyRetrieveUserProxyAgent(RetrieveUserProxyAgent):
    def query_vector_db(
        self,
        query_texts: List[str],
        n_results: int = 10,
        search_string: str = "",
        **kwargs,
    ) -> 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,
        )

        self._results = results
        print("doc_ids: ", results["ids"])

retrieve_docs​

def retrieve_docs(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.

Arguments:

  • 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.

message_generator​

@staticmethod
def message_generator(sender, recipient, context)

Generate an initial message with the given context for the RetrieveUserProxyAgent.

Arguments:

  • 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.