Underground-Digital/Workflow-Engine

1import json2 3from core.agent.cot_agent_runner import CotAgentRunner4from core.file import file_manager5from core.model_runtime.entities import (6    AssistantPromptMessage,7    PromptMessage,8    PromptMessageContent,9    SystemPromptMessage,10    TextPromptMessageContent,11    UserPromptMessage,12)13from core.model_runtime.utils.encoders import jsonable_encoder14 15 16class CotChatAgentRunner(CotAgentRunner):17    def _organize_system_prompt(self) -> SystemPromptMessage:18        """19        Organize system prompt20        """21        prompt_entity = self.app_config.agent.prompt22        first_prompt = prompt_entity.first_prompt23 24        system_prompt = (25            first_prompt.replace("{{instruction}}", self._instruction)26            .replace("{{tools}}", json.dumps(jsonable_encoder(self._prompt_messages_tools)))27            .replace("{{tool_names}}", ", ".join([tool.name for tool in self._prompt_messages_tools]))28        )29 30        return SystemPromptMessage(content=system_prompt)31 32    def _organize_user_query(self, query, prompt_messages: list[PromptMessage]) -> list[PromptMessage]:33        """34        Organize user query35        """36        if self.files:37            prompt_message_contents: list[PromptMessageContent] = []38            prompt_message_contents.append(TextPromptMessageContent(data=query))39            for file_obj in self.files:40                prompt_message_contents.append(file_manager.to_prompt_message_content(file_obj))41 42            prompt_messages.append(UserPromptMessage(content=prompt_message_contents))43        else:44            prompt_messages.append(UserPromptMessage(content=query))45 46        return prompt_messages47 48    def _organize_prompt_messages(self) -> list[PromptMessage]:49        """50        Organize51        """52        # organize system prompt53        system_message = self._organize_system_prompt()54 55        # organize current assistant messages56        agent_scratchpad = self._agent_scratchpad57        if not agent_scratchpad:58            assistant_messages = []59        else:60            assistant_message = AssistantPromptMessage(content="")61            for unit in agent_scratchpad:62                if unit.is_final():63                    assistant_message.content += f"Final Answer: {unit.agent_response}"64                else:65                    assistant_message.content += f"Thought: {unit.thought}\n\n"66                    if unit.action_str:67                        assistant_message.content += f"Action: {unit.action_str}\n\n"68                    if unit.observation:69                        assistant_message.content += f"Observation: {unit.observation}\n\n"70 71            assistant_messages = [assistant_message]72 73        # query messages74        query_messages = self._organize_user_query(self._query, [])75 76        if assistant_messages:77            # organize historic prompt messages78            historic_messages = self._organize_historic_prompt_messages(79                [system_message, *query_messages, *assistant_messages, UserPromptMessage(content="continue")]80            )81            messages = [82                system_message,83                *historic_messages,84                *query_messages,85                *assistant_messages,86                UserPromptMessage(content="continue"),87            ]88        else:89            # organize historic prompt messages90            historic_messages = self._organize_historic_prompt_messages([system_message, *query_messages])91            messages = [system_message, *historic_messages, *query_messages]92 93        # join all messages94        return messages95