feat: universal chat in explore (#649)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>

api/core/agent/agent/calc_token_mixin.py

@@ -0,0 +1,35 @@
+from typing import cast, List
+
+from langchain import OpenAI
+from langchain.base_language import BaseLanguageModel
+from langchain.chat_models.openai import ChatOpenAI
+from langchain.schema import BaseMessage
+
+from core.constant import llm_constant
+
+
+class CalcTokenMixin:
+
+    def get_num_tokens_from_messages(self, llm: BaseLanguageModel, messages: List[BaseMessage], **kwargs) -> int:
+        llm = cast(ChatOpenAI, llm)
+        return llm.get_num_tokens_from_messages(messages)
+
+    def get_message_rest_tokens(self, llm: BaseLanguageModel, messages: List[BaseMessage], **kwargs) -> int:
+        """
+        Got the rest tokens available for the model after excluding messages tokens and completion max tokens
+
+        :param llm:
+        :param messages:
+        :return:
+        """
+        llm = cast(ChatOpenAI, llm)
+        llm_max_tokens = llm_constant.max_context_token_length[llm.model_name]
+        completion_max_tokens = llm.max_tokens
+        used_tokens = self.get_num_tokens_from_messages(llm, messages, **kwargs)
+        rest_tokens = llm_max_tokens - completion_max_tokens - used_tokens
+
+        return rest_tokens
+
+
+class ExceededLLMTokensLimitError(Exception):
+    pass

api/core/agent/agent/multi_dataset_router_agent.py

@@ -0,0 +1,84 @@
+from typing import Tuple, List, Any, Union, Sequence, Optional, cast
+
+from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, BaseLanguageModel, SystemMessage
+from langchain.tools import BaseTool
+
+from core.tool.dataset_retriever_tool import DatasetRetrieverTool
+
+
+class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
+    """
+    An Multi Dataset Retrieve Agent driven by Router.
+    """
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        return True
+
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        if len(self.tools) == 0:
+            return AgentFinish(return_values={"output": ''}, log='')
+        elif len(self.tools) == 1:
+            tool = next(iter(self.tools))
+            tool = cast(DatasetRetrieverTool, tool)
+            rst = tool.run(tool_input={'dataset_id': tool.dataset_id, 'query': kwargs['input']})
+            return AgentFinish(return_values={"output": rst}, log=rst)
+
+        if intermediate_steps:
+            _, observation = intermediate_steps[-1]
+            return AgentFinish(return_values={"output": observation}, log=observation)
+
+        return super().plan(intermediate_steps, callbacks, **kwargs)
+
+    async def aplan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        raise NotImplementedError()
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseSingleActionAgent:
+        llm.model_name = 'gpt-3.5-turbo'
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=system_message,
+            **kwargs,
+        )

api/core/agent/agent/openai_function_call.py

@@ -0,0 +1,120 @@
+from datetime import datetime
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+import pytz
+from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
+from langchain.agents.openai_functions_agent.base import _parse_ai_message, \
+    _format_intermediate_steps
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, SystemMessage, BaseLanguageModel
+from langchain.tools import BaseTool
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError
+from core.agent.agent.openai_function_call_summarize_mixin import OpenAIFunctionCallSummarizeMixin
+
+
+class AutoSummarizingOpenAIFunctionCallAgent(OpenAIFunctionsAgent, OpenAIFunctionCallSummarizeMixin):
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseSingleActionAgent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=cls.get_system_message(),
+            **kwargs,
+        )
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        original_max_tokens = self.llm.max_tokens
+        self.llm.max_tokens = 15
+
+        prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
+        messages = prompt.to_messages()
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=None
+        )
+
+        function_call = predicted_message.additional_kwargs.get("function_call", {})
+
+        self.llm.max_tokens = original_max_tokens
+
+        return True if function_call else False
+
+    def plan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
+        selected_inputs = {
+            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
+        }
+        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
+        prompt = self.prompt.format_prompt(**full_inputs)
+        messages = prompt.to_messages()
+
+        # summarize messages if rest_tokens < 0
+        try:
+            messages = self.summarize_messages_if_needed(self.llm, messages, functions=self.functions)
+        except ExceededLLMTokensLimitError as e:
+            return AgentFinish(return_values={"output": str(e)}, log=str(e))
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=callbacks
+        )
+        agent_decision = _parse_ai_message(predicted_message)
+        return agent_decision
+
+    @classmethod
+    def get_system_message(cls):
+        # get current time
+        current_time = datetime.now()
+        current_timezone = pytz.timezone('UTC')
+        current_time = current_timezone.localize(current_time)
+
+        return SystemMessage(content="You are a helpful AI assistant.\n"
+                                     "Current time: {}\n"
+                                     "Respond directly if appropriate.".format(
+            current_time.strftime("%Y-%m-%d %H:%M:%S %Z%z")))
+
+    def return_stopped_response(
+            self,
+            early_stopping_method: str,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            **kwargs: Any,
+    ) -> AgentFinish:
+        try:
+            return super().return_stopped_response(early_stopping_method, intermediate_steps, **kwargs)
+        except ValueError:
+            return AgentFinish({"output": "I'm sorry, I don't know how to respond to that."}, "")

api/core/agent/agent/openai_function_call_summarize_mixin.py

@@ -0,0 +1,132 @@
+from typing import cast, List
+
+from langchain.chat_models import ChatOpenAI
+from langchain.chat_models.openai import _convert_message_to_dict
+from langchain.memory.summary import SummarizerMixin
+from langchain.schema import SystemMessage, HumanMessage, BaseMessage, AIMessage, BaseLanguageModel
+from pydantic import BaseModel
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError, CalcTokenMixin
+
+
+class OpenAIFunctionCallSummarizeMixin(BaseModel, CalcTokenMixin):
+    moving_summary_buffer: str = ""
+    moving_summary_index: int = 0
+    summary_llm: BaseLanguageModel
+
+    def summarize_messages_if_needed(self, llm: BaseLanguageModel, messages: List[BaseMessage], **kwargs) -> List[BaseMessage]:
+        # calculate rest tokens and summarize previous function observation messages if rest_tokens < 0
+        rest_tokens = self.get_message_rest_tokens(llm, messages, **kwargs)
+        rest_tokens = rest_tokens - 20  # to deal with the inaccuracy of rest_tokens
+        if rest_tokens >= 0:
+            return messages
+
+        system_message = None
+        human_message = None
+        should_summary_messages = []
+        for message in messages:
+            if isinstance(message, SystemMessage):
+                system_message = message
+            elif isinstance(message, HumanMessage):
+                human_message = message
+            else:
+                should_summary_messages.append(message)
+
+        if len(should_summary_messages) > 2:
+            ai_message = should_summary_messages[-2]
+            function_message = should_summary_messages[-1]
+            should_summary_messages = should_summary_messages[self.moving_summary_index:-2]
+            self.moving_summary_index = len(should_summary_messages)
+        else:
+            error_msg = "Exceeded LLM tokens limit, stopped."
+            raise ExceededLLMTokensLimitError(error_msg)
+
+        new_messages = [system_message, human_message]
+
+        if self.moving_summary_index == 0:
+            should_summary_messages.insert(0, human_message)
+
+        summary_handler = SummarizerMixin(llm=self.summary_llm)
+        self.moving_summary_buffer = summary_handler.predict_new_summary(
+            messages=should_summary_messages,
+            existing_summary=self.moving_summary_buffer
+        )
+
+        new_messages.append(AIMessage(content=self.moving_summary_buffer))
+        new_messages.append(ai_message)
+        new_messages.append(function_message)
+
+        return new_messages
+
+    def get_num_tokens_from_messages(self, llm: BaseLanguageModel, messages: List[BaseMessage], **kwargs) -> int:
+        """Calculate num tokens for gpt-3.5-turbo and gpt-4 with tiktoken package.
+
+        Official documentation: https://github.com/openai/openai-cookbook/blob/
+        main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb"""
+        llm = cast(ChatOpenAI, llm)
+        model, encoding = llm._get_encoding_model()
+        if model.startswith("gpt-3.5-turbo"):
+            # every message follows <im_start>{role/name}\n{content}<im_end>\n
+            tokens_per_message = 4
+            # if there's a name, the role is omitted
+            tokens_per_name = -1
+        elif model.startswith("gpt-4"):
+            tokens_per_message = 3
+            tokens_per_name = 1
+        else:
+            raise NotImplementedError(
+                f"get_num_tokens_from_messages() is not presently implemented "
+                f"for model {model}."
+                "See https://github.com/openai/openai-python/blob/main/chatml.md for "
+                "information on how messages are converted to tokens."
+            )
+        num_tokens = 0
+        for m in messages:
+            message = _convert_message_to_dict(m)
+            num_tokens += tokens_per_message
+            for key, value in message.items():
+                if key == "function_call":
+                    for f_key, f_value in value.items():
+                        num_tokens += len(encoding.encode(f_key))
+                        num_tokens += len(encoding.encode(f_value))
+                else:
+                    num_tokens += len(encoding.encode(value))
+
+                if key == "name":
+                    num_tokens += tokens_per_name
+        # every reply is primed with <im_start>assistant
+        num_tokens += 3
+
+        if kwargs.get('functions'):
+            for function in kwargs.get('functions'):
+                num_tokens += len(encoding.encode('name'))
+                num_tokens += len(encoding.encode(function.get("name")))
+                num_tokens += len(encoding.encode('description'))
+                num_tokens += len(encoding.encode(function.get("description")))
+                parameters = function.get("parameters")
+                num_tokens += len(encoding.encode('parameters'))
+                if 'title' in parameters:
+                    num_tokens += len(encoding.encode('title'))
+                    num_tokens += len(encoding.encode(parameters.get("title")))
+                num_tokens += len(encoding.encode('type'))
+                num_tokens += len(encoding.encode(parameters.get("type")))
+                if 'properties' in parameters:
+                    num_tokens += len(encoding.encode('properties'))
+                    for key, value in parameters.get('properties').items():
+                        num_tokens += len(encoding.encode(key))
+                        for field_key, field_value in value.items():
+                            num_tokens += len(encoding.encode(field_key))
+                            if field_key == 'enum':
+                                for enum_field in field_value:
+                                    num_tokens += 3
+                                    num_tokens += len(encoding.encode(enum_field))
+                            else:
+                                num_tokens += len(encoding.encode(field_key))
+                                num_tokens += len(encoding.encode(str(field_value)))
+                if 'required' in parameters:
+                    num_tokens += len(encoding.encode('required'))
+                    for required_field in parameters['required']:
+                        num_tokens += 3
+                        num_tokens += len(encoding.encode(required_field))
+
+        return num_tokens

api/core/agent/agent/openai_multi_function_call.py

@@ -0,0 +1,109 @@
+from datetime import datetime
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+import pytz
+from langchain.agents import BaseMultiActionAgent
+from langchain.agents.openai_functions_multi_agent.base import OpenAIMultiFunctionsAgent, _format_intermediate_steps, \
+    _parse_ai_message
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, SystemMessage, BaseLanguageModel
+from langchain.tools import BaseTool
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError
+from core.agent.agent.openai_function_call_summarize_mixin import OpenAIFunctionCallSummarizeMixin
+
+
+class AutoSummarizingOpenMultiAIFunctionCallAgent(OpenAIMultiFunctionsAgent, OpenAIFunctionCallSummarizeMixin):
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseMultiActionAgent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=cls.get_system_message(),
+            **kwargs,
+        )
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        original_max_tokens = self.llm.max_tokens
+        self.llm.max_tokens = 15
+
+        prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
+        messages = prompt.to_messages()
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=None
+        )
+
+        function_call = predicted_message.additional_kwargs.get("function_call", {})
+
+        self.llm.max_tokens = original_max_tokens
+
+        return True if function_call else False
+
+    def plan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
+        selected_inputs = {
+            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
+        }
+        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
+        prompt = self.prompt.format_prompt(**full_inputs)
+        messages = prompt.to_messages()
+
+        # summarize messages if rest_tokens < 0
+        try:
+            messages = self.summarize_messages_if_needed(self.llm, messages, functions=self.functions)
+        except ExceededLLMTokensLimitError as e:
+            return AgentFinish(return_values={"output": str(e)}, log=str(e))
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=callbacks
+        )
+        agent_decision = _parse_ai_message(predicted_message)
+        return agent_decision
+
+    @classmethod
+    def get_system_message(cls):
+        # get current time
+        current_time = datetime.now()
+        current_timezone = pytz.timezone('UTC')
+        current_time = current_timezone.localize(current_time)
+
+        return SystemMessage(content="You are a helpful AI assistant.\n"
+                                     "Current time: {}\n"
+                                     "Respond directly if appropriate.".format(
+            current_time.strftime("%Y-%m-%d %H:%M:%S %Z%z")))

api/core/agent/agent/output_parser/structured_chat.py

@@ -0,0 +1,29 @@
+import json
+import re
+from typing import Union
+
+from langchain.agents.structured_chat.output_parser import StructuredChatOutputParser as LCStructuredChatOutputParser, \
+    logger
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+
+
+class StructuredChatOutputParser(LCStructuredChatOutputParser):
+    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
+        try:
+            action_match = re.search(r"```(.*?)\n(.*?)```?", text, re.DOTALL)
+            if action_match is not None:
+                response = json.loads(action_match.group(2).strip(), strict=False)
+                if isinstance(response, list):
+                    # gpt turbo frequently ignores the directive to emit a single action
+                    logger.warning("Got multiple action responses: %s", response)
+                    response = response[0]
+                if response["action"] == "Final Answer":
+                    return AgentFinish({"output": response["action_input"]}, text)
+                else:
+                    return AgentAction(
+                        response["action"], response.get("action_input", {}), text
+                    )
+            else:
+                return AgentFinish({"output": text}, text)
+        except Exception as e:
+            raise OutputParserException(f"Could not parse LLM output: {text}") from e

api/core/agent/agent/structured_chat.py

@@ -0,0 +1,182 @@
+import re
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+from langchain import BasePromptTemplate
+from langchain.agents import StructuredChatAgent, AgentOutputParser, Agent
+from langchain.agents.structured_chat.base import HUMAN_MESSAGE_TEMPLATE
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.memory.summary import SummarizerMixin
+from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain.schema import AgentAction, AgentFinish, AIMessage, HumanMessage
+from langchain.tools import BaseTool
+from langchain.agents.structured_chat.prompt import PREFIX, SUFFIX
+
+from core.agent.agent.calc_token_mixin import CalcTokenMixin, ExceededLLMTokensLimitError
+
+
+FORMAT_INSTRUCTIONS = """Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
+The nouns in the format of "Thought", "Action", "Action Input", "Final Answer" must be expressed in English.
+Valid "action" values: "Final Answer" or {tool_names}
+
+Provide only ONE action per $JSON_BLOB, as shown:
+
+```
+{{{{
+  "action": $TOOL_NAME,
+  "action_input": $INPUT
+}}}}
+```
+
+Follow this format:
+
+Question: input question to answer
+Thought: consider previous and subsequent steps
+Action:
+```
+$JSON_BLOB
+```
+Observation: action result
+... (repeat Thought/Action/Observation N times)
+Thought: I know what to respond
+Action:
+```
+{{{{
+  "action": "Final Answer",
+  "action_input": "Final response to human"
+}}}}
+```"""
+
+
+class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
+    moving_summary_buffer: str = ""
+    moving_summary_index: int = 0
+    summary_llm: BaseLanguageModel
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+        Using the ReACT mode to determine whether an agent is needed is costly,
+        so it's better to just use an Agent for reasoning, which is cheaper.
+
+        :param query:
+        :return:
+        """
+        return True
+
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date,
+                along with observations
+            callbacks: Callbacks to run.
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
+
+        prompts, _ = self.llm_chain.prep_prompts(input_list=[self.llm_chain.prep_inputs(full_inputs)])
+        messages = []
+        if prompts:
+            messages = prompts[0].to_messages()
+
+        rest_tokens = self.get_message_rest_tokens(self.llm_chain.llm, messages)
+        if rest_tokens < 0:
+            full_inputs = self.summarize_messages(intermediate_steps, **kwargs)
+
+        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
+        return self.output_parser.parse(full_output)
+
+    def summarize_messages(self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs):
+        if len(intermediate_steps) >= 2:
+            should_summary_intermediate_steps = intermediate_steps[self.moving_summary_index:-1]
+            should_summary_messages = [AIMessage(content=observation)
+                                       for _, observation in should_summary_intermediate_steps]
+            if self.moving_summary_index == 0:
+                should_summary_messages.insert(0, HumanMessage(content=kwargs.get("input")))
+
+            self.moving_summary_index = len(intermediate_steps)
+        else:
+            error_msg = "Exceeded LLM tokens limit, stopped."
+            raise ExceededLLMTokensLimitError(error_msg)
+
+        summary_handler = SummarizerMixin(llm=self.summary_llm)
+        if self.moving_summary_buffer and 'chat_history' in kwargs:
+            kwargs["chat_history"].pop()
+
+        self.moving_summary_buffer = summary_handler.predict_new_summary(
+            messages=should_summary_messages,
+            existing_summary=self.moving_summary_buffer
+        )
+
+        if 'chat_history' in kwargs:
+            kwargs["chat_history"].append(AIMessage(content=self.moving_summary_buffer))
+
+        return self.get_full_inputs([intermediate_steps[-1]], **kwargs)
+
+    @classmethod
+    def create_prompt(
+            cls,
+            tools: Sequence[BaseTool],
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+    ) -> BasePromptTemplate:
+        tool_strings = []
+        for tool in tools:
+            args_schema = re.sub("}", "}}}}", re.sub("{", "{{{{", str(tool.args)))
+            tool_strings.append(f"{tool.name}: {tool.description}, args: {args_schema}")
+        formatted_tools = "\n".join(tool_strings)
+        tool_names = ", ".join([('"' + tool.name + '"') for tool in tools])
+        format_instructions = format_instructions.format(tool_names=tool_names)
+        template = "\n\n".join([prefix, formatted_tools, format_instructions, suffix])
+        if input_variables is None:
+            input_variables = ["input", "agent_scratchpad"]
+        _memory_prompts = memory_prompts or []
+        messages = [
+            SystemMessagePromptTemplate.from_template(template),
+            *_memory_prompts,
+            HumanMessagePromptTemplate.from_template(human_message_template),
+        ]
+        return ChatPromptTemplate(input_variables=input_variables, messages=messages)
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            output_parser: Optional[AgentOutputParser] = None,
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+            **kwargs: Any,
+    ) -> Agent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            output_parser=output_parser,
+            prefix=prefix,
+            suffix=suffix,
+            human_message_template=human_message_template,
+            format_instructions=format_instructions,
+            input_variables=input_variables,
+            memory_prompts=memory_prompts,
+            **kwargs,
+        )

api/core/agent/agent_builder.py

@@ -1,86 +0,0 @@
-from typing import Optional
-
-from langchain import LLMChain
-from langchain.agents import ZeroShotAgent, AgentExecutor, ConversationalAgent
-from langchain.callbacks.manager import CallbackManager
-from langchain.memory.chat_memory import BaseChatMemory
-
-from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
-from core.callback_handler.dataset_tool_callback_handler import DatasetToolCallbackHandler
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.llm.llm_builder import LLMBuilder
-
-
-class AgentBuilder:
-    @classmethod
-    def to_agent_chain(cls, tenant_id: str, tools, memory: Optional[BaseChatMemory],
-                       dataset_tool_callback_handler: DatasetToolCallbackHandler,
-                       agent_loop_gather_callback_handler: AgentLoopGatherCallbackHandler):
-        llm = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name=agent_loop_gather_callback_handler.model_name,
-            temperature=0,
-            max_tokens=1024,
-            callbacks=[agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()]
-        )
-
-        for tool in tools:
-            tool.callbacks = [
-                agent_loop_gather_callback_handler,
-                dataset_tool_callback_handler,
-                DifyStdOutCallbackHandler()
-            ]
-
-        prompt = cls.build_agent_prompt_template(
-            tools=tools,
-            memory=memory,
-        )
-
-        agent_llm_chain = LLMChain(
-            llm=llm,
-            prompt=prompt,
-        )
-
-        agent = cls.build_agent(agent_llm_chain=agent_llm_chain, memory=memory)
-
-        agent_callback_manager = CallbackManager(
-            [agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()]
-        )
-
-        agent_chain = AgentExecutor.from_agent_and_tools(
-            tools=tools,
-            agent=agent,
-            memory=memory,
-            callbacks=agent_callback_manager,
-            max_iterations=6,
-            early_stopping_method="generate",
-            # `generate` will continue to complete the last inference after reaching the iteration limit or request time limit
-        )
-
-        return agent_chain
-
-    @classmethod
-    def build_agent_prompt_template(cls, tools, memory: Optional[BaseChatMemory]):
-        if memory:
-            prompt = ConversationalAgent.create_prompt(
-                tools=tools,
-            )
-        else:
-            prompt = ZeroShotAgent.create_prompt(
-                tools=tools,
-            )
-
-        return prompt
-
-    @classmethod
-    def build_agent(cls, agent_llm_chain: LLMChain, memory: Optional[BaseChatMemory]):
-        if memory:
-            agent = ConversationalAgent(
-                llm_chain=agent_llm_chain
-            )
-        else:
-            agent = ZeroShotAgent(
-                llm_chain=agent_llm_chain
-            )
-
-        return agent

api/core/agent/agent_executor.py

@@ -0,0 +1,121 @@
+import enum
+import logging
+from typing import Union, Optional
+
+from langchain.agents import BaseSingleActionAgent, BaseMultiActionAgent
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import Callbacks
+from langchain.memory.chat_memory import BaseChatMemory
+from langchain.tools import BaseTool
+from pydantic import BaseModel, Extra
+
+from core.agent.agent.multi_dataset_router_agent import MultiDatasetRouterAgent
+from core.agent.agent.openai_function_call import AutoSummarizingOpenAIFunctionCallAgent
+from core.agent.agent.openai_multi_function_call import AutoSummarizingOpenMultiAIFunctionCallAgent
+from core.agent.agent.output_parser.structured_chat import StructuredChatOutputParser
+from core.agent.agent.structured_chat import AutoSummarizingStructuredChatAgent
+from langchain.agents import AgentExecutor as LCAgentExecutor
+
+from core.tool.dataset_retriever_tool import DatasetRetrieverTool
+
+
+class PlanningStrategy(str, enum.Enum):
+    ROUTER = 'router'
+    REACT = 'react'
+    FUNCTION_CALL = 'function_call'
+    MULTI_FUNCTION_CALL = 'multi_function_call'
+
+
+class AgentConfiguration(BaseModel):
+    strategy: PlanningStrategy
+    llm: BaseLanguageModel
+    tools: list[BaseTool]
+    summary_llm: BaseLanguageModel
+    memory: Optional[BaseChatMemory] = None
+    callbacks: Callbacks = None
+    max_iterations: int = 6
+    max_execution_time: Optional[float] = None
+    early_stopping_method: str = "generate"
+    # `generate` will continue to complete the last inference after reaching the iteration limit or request time limit
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+
+class AgentExecuteResult(BaseModel):
+    strategy: PlanningStrategy
+    output: Optional[str]
+    configuration: AgentConfiguration
+
+
+class AgentExecutor:
+    def __init__(self, configuration: AgentConfiguration):
+        self.configuration = configuration
+        self.agent = self._init_agent()
+
+    def _init_agent(self) -> Union[BaseSingleActionAgent | BaseMultiActionAgent]:
+        if self.configuration.strategy == PlanningStrategy.REACT:
+            agent = AutoSummarizingStructuredChatAgent.from_llm_and_tools(
+                llm=self.configuration.llm,
+                tools=self.configuration.tools,
+                output_parser=StructuredChatOutputParser(),
+                summary_llm=self.configuration.summary_llm,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.FUNCTION_CALL:
+            agent = AutoSummarizingOpenAIFunctionCallAgent.from_llm_and_tools(
+                llm=self.configuration.llm,
+                tools=self.configuration.tools,
+                extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,  # used for read chat histories memory
+                summary_llm=self.configuration.summary_llm,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.MULTI_FUNCTION_CALL:
+            agent = AutoSummarizingOpenMultiAIFunctionCallAgent.from_llm_and_tools(
+                llm=self.configuration.llm,
+                tools=self.configuration.tools,
+                extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,  # used for read chat histories memory
+                summary_llm=self.configuration.summary_llm,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.ROUTER:
+            self.configuration.tools = [t for t in self.configuration.tools if isinstance(t, DatasetRetrieverTool)]
+            agent = MultiDatasetRouterAgent.from_llm_and_tools(
+                llm=self.configuration.llm,
+                tools=self.configuration.tools,
+                extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,
+                verbose=True
+            )
+        else:
+            raise NotImplementedError(f"Unknown Agent Strategy: {self.configuration.strategy}")
+
+        return agent
+
+    def should_use_agent(self, query: str) -> bool:
+        return self.agent.should_use_agent(query)
+
+    def run(self, query: str) -> AgentExecuteResult:
+        agent_executor = LCAgentExecutor.from_agent_and_tools(
+            agent=self.agent,
+            tools=self.configuration.tools,
+            memory=self.configuration.memory,
+            max_iterations=self.configuration.max_iterations,
+            max_execution_time=self.configuration.max_execution_time,
+            early_stopping_method=self.configuration.early_stopping_method,
+            callbacks=self.configuration.callbacks
+        )
+
+        try:
+            output = agent_executor.run(query)
+        except Exception:
+            logging.exception("agent_executor run failed")
+            output = None
+
+        return AgentExecuteResult(
+            output=output,
+            strategy=self.configuration.strategy,
+            configuration=self.configuration
+        )