aphilippov/python-server-api

1from autogen.agentchat.agent import Agent2from autogen.agentchat.assistant_agent import AssistantAgent3from typing import Dict, Optional, Union, List, Tuple, Any4 5 6class RetrieveAssistantAgent(AssistantAgent):7    """(Experimental) Retrieve Assistant agent, designed to solve a task with LLM.8 9    RetrieveAssistantAgent is a subclass of AssistantAgent configured with a default system message.10    The default system message is designed to solve a task with LLM,11    including suggesting python code blocks and debugging.12    `human_input_mode` is default to "NEVER"13    and `code_execution_config` is default to False.14    This agent doesn't execute code by default, and expects the user to execute the code.15    """16 17    def __init__(self, *args, **kwargs):18        super().__init__(*args, **kwargs)19        self.register_reply(Agent, RetrieveAssistantAgent._generate_retrieve_assistant_reply)20 21    def _generate_retrieve_assistant_reply(22        self,23        messages: Optional[List[Dict]] = None,24        sender: Optional[Agent] = None,25        config: Optional[Any] = None,26    ) -> Tuple[bool, Union[str, Dict, None]]:27        if config is None:28            config = self29        if messages is None:30            messages = self._oai_messages[sender]31        message = messages[-1]32        if "exitcode: 0 (execution succeeded)" in message.get("content", ""):33            # Terminate the conversation when the code execution succeeds. Although sometimes even when the34            # code execution succeeds, the task is not solved, but it's hard to tell. If the human_input_mode35            # of RetrieveUserProxyAgent is "TERMINATE" or "ALWAYS", user can still continue the conversation.36            return True, "TERMINATE"37        elif (38            "UPDATE CONTEXT" in message.get("content", "")[-20:].upper()39            or "UPDATE CONTEXT" in message.get("content", "")[:20].upper()40        ):41            return True, "UPDATE CONTEXT"42        else:43            return False, None44