Underground-Digital/Workflow-Engine

1import json2from typing import Optional3 4from core.agent.cot_agent_runner import CotAgentRunner5from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, UserPromptMessage6from core.model_runtime.utils.encoders import jsonable_encoder7 8 9class CotCompletionAgentRunner(CotAgentRunner):10    def _organize_instruction_prompt(self) -> str:11        """12        Organize instruction prompt13        """14        prompt_entity = self.app_config.agent.prompt15        first_prompt = prompt_entity.first_prompt16 17        system_prompt = (18            first_prompt.replace("{{instruction}}", self._instruction)19            .replace("{{tools}}", json.dumps(jsonable_encoder(self._prompt_messages_tools)))20            .replace("{{tool_names}}", ", ".join([tool.name for tool in self._prompt_messages_tools]))21        )22 23        return system_prompt24 25    def _organize_historic_prompt(self, current_session_messages: Optional[list[PromptMessage]] = None) -> str:26        """27        Organize historic prompt28        """29        historic_prompt_messages = self._organize_historic_prompt_messages(current_session_messages)30        historic_prompt = ""31 32        for message in historic_prompt_messages:33            if isinstance(message, UserPromptMessage):34                historic_prompt += f"Question: {message.content}\n\n"35            elif isinstance(message, AssistantPromptMessage):36                historic_prompt += message.content + "\n\n"37 38        return historic_prompt39 40    def _organize_prompt_messages(self) -> list[PromptMessage]:41        """42        Organize prompt messages43        """44        # organize system prompt45        system_prompt = self._organize_instruction_prompt()46 47        # organize historic prompt messages48        historic_prompt = self._organize_historic_prompt()49 50        # organize current assistant messages51        agent_scratchpad = self._agent_scratchpad52        assistant_prompt = ""53        for unit in agent_scratchpad:54            if unit.is_final():55                assistant_prompt += f"Final Answer: {unit.agent_response}"56            else:57                assistant_prompt += f"Thought: {unit.thought}\n\n"58                if unit.action_str:59                    assistant_prompt += f"Action: {unit.action_str}\n\n"60                if unit.observation:61                    assistant_prompt += f"Observation: {unit.observation}\n\n"62 63        # query messages64        query_prompt = f"Question: {self._query}"65 66        # join all messages67        prompt = (68            system_prompt.replace("{{historic_messages}}", historic_prompt)69            .replace("{{agent_scratchpad}}", assistant_prompt)70            .replace("{{query}}", query_prompt)71        )72 73        return [UserPromptMessage(content=prompt)]74