Underground-Digital/Workflow-Engine

1from collections.abc import Mapping, Sequence2from typing import Any, Optional, TextIO, Union3 4from pydantic import BaseModel5 6from configs import dify_config7from core.ops.entities.trace_entity import TraceTaskName8from core.ops.ops_trace_manager import TraceQueueManager, TraceTask9from core.tools.entities.tool_entities import ToolInvokeMessage10 11_TEXT_COLOR_MAPPING = {12    "blue": "36;1",13    "yellow": "33;1",14    "pink": "38;5;200",15    "green": "32;1",16    "red": "31;1",17}18 19 20def get_colored_text(text: str, color: str) -> str:21    """Get colored text."""22    color_str = _TEXT_COLOR_MAPPING[color]23    return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m"24 25 26def print_text(text: str, color: Optional[str] = None, end: str = "", file: Optional[TextIO] = None) -> None:27    """Print text with highlighting and no end characters."""28    text_to_print = get_colored_text(text, color) if color else text29    print(text_to_print, end=end, file=file)30    if file:31        file.flush()  # ensure all printed content are written to file32 33 34class DifyAgentCallbackHandler(BaseModel):35    """Callback Handler that prints to std out."""36 37    color: Optional[str] = ""38    current_loop: int = 139 40    def __init__(self, color: Optional[str] = None) -> None:41        super().__init__()42        """Initialize callback handler."""43        # use a specific color is not specified44        self.color = color or "green"45        self.current_loop = 146 47    def on_tool_start(48        self,49        tool_name: str,50        tool_inputs: Mapping[str, Any],51    ) -> None:52        """Do nothing."""53        if dify_config.DEBUG:54            print_text("\n[on_tool_start] ToolCall:" + tool_name + "\n" + str(tool_inputs) + "\n", color=self.color)55 56    def on_tool_end(57        self,58        tool_name: str,59        tool_inputs: Mapping[str, Any],60        tool_outputs: Sequence[ToolInvokeMessage],61        message_id: Optional[str] = None,62        timer: Optional[Any] = None,63        trace_manager: Optional[TraceQueueManager] = None,64    ) -> None:65        """If not the final action, print out observation."""66        if dify_config.DEBUG:67            print_text("\n[on_tool_end]\n", color=self.color)68            print_text("Tool: " + tool_name + "\n", color=self.color)69            print_text("Inputs: " + str(tool_inputs) + "\n", color=self.color)70            print_text("Outputs: " + str(tool_outputs)[:1000] + "\n", color=self.color)71            print_text("\n")72 73        if trace_manager:74            trace_manager.add_trace_task(75                TraceTask(76                    TraceTaskName.TOOL_TRACE,77                    message_id=message_id,78                    tool_name=tool_name,79                    tool_inputs=tool_inputs,80                    tool_outputs=tool_outputs,81                    timer=timer,82                )83            )84 85    def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:86        """Do nothing."""87        if dify_config.DEBUG:88            print_text("\n[on_tool_error] Error: " + str(error) + "\n", color="red")89 90    def on_agent_start(self, thought: str) -> None:91        """Run on agent start."""92        if dify_config.DEBUG:93            if thought:94                print_text(95                    "\n[on_agent_start] \nCurrent Loop: " + str(self.current_loop) + "\nThought: " + thought + "\n",96                    color=self.color,97                )98            else:99                print_text("\n[on_agent_start] \nCurrent Loop: " + str(self.current_loop) + "\n", color=self.color)100 101    def on_agent_finish(self, color: Optional[str] = None, **kwargs: Any) -> None:102        """Run on agent end."""103        if dify_config.DEBUG:104            print_text("\n[on_agent_finish]\n Loop: " + str(self.current_loop) + "\n", color=self.color)105 106        self.current_loop += 1107 108    @property109    def ignore_agent(self) -> bool:110        """Whether to ignore agent callbacks."""111        return not dify_config.DEBUG112 113    @property114    def ignore_chat_model(self) -> bool:115        """Whether to ignore chat model callbacks."""116        return not dify_config.DEBUG117