Underground-Digital/Workflow-Engine

1from abc import ABC, abstractmethod2from typing import Optional3 4from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk5from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool6from core.model_runtime.model_providers.__base.ai_model import AIModel7 8_TEXT_COLOR_MAPPING = {9    "blue": "36;1",10    "yellow": "33;1",11    "pink": "38;5;200",12    "green": "32;1",13    "red": "31;1",14}15 16 17class Callback(ABC):18    """19    Base class for callbacks.20    Only for LLM.21    """22 23    raise_error: bool = False24 25    @abstractmethod26    def on_before_invoke(27        self,28        llm_instance: AIModel,29        model: str,30        credentials: dict,31        prompt_messages: list[PromptMessage],32        model_parameters: dict,33        tools: Optional[list[PromptMessageTool]] = None,34        stop: Optional[list[str]] = None,35        stream: bool = True,36        user: Optional[str] = None,37    ) -> None:38        """39        Before invoke callback40 41        :param llm_instance: LLM instance42        :param model: model name43        :param credentials: model credentials44        :param prompt_messages: prompt messages45        :param model_parameters: model parameters46        :param tools: tools for tool calling47        :param stop: stop words48        :param stream: is stream response49        :param user: unique user id50        """51        raise NotImplementedError()52 53    @abstractmethod54    def on_new_chunk(55        self,56        llm_instance: AIModel,57        chunk: LLMResultChunk,58        model: str,59        credentials: dict,60        prompt_messages: list[PromptMessage],61        model_parameters: dict,62        tools: Optional[list[PromptMessageTool]] = None,63        stop: Optional[list[str]] = None,64        stream: bool = True,65        user: Optional[str] = None,66    ):67        """68        On new chunk callback69 70        :param llm_instance: LLM instance71        :param chunk: chunk72        :param model: model name73        :param credentials: model credentials74        :param prompt_messages: prompt messages75        :param model_parameters: model parameters76        :param tools: tools for tool calling77        :param stop: stop words78        :param stream: is stream response79        :param user: unique user id80        """81        raise NotImplementedError()82 83    @abstractmethod84    def on_after_invoke(85        self,86        llm_instance: AIModel,87        result: LLMResult,88        model: str,89        credentials: dict,90        prompt_messages: list[PromptMessage],91        model_parameters: dict,92        tools: Optional[list[PromptMessageTool]] = None,93        stop: Optional[list[str]] = None,94        stream: bool = True,95        user: Optional[str] = None,96    ) -> None:97        """98        After invoke callback99 100        :param llm_instance: LLM instance101        :param result: result102        :param model: model name103        :param credentials: model credentials104        :param prompt_messages: prompt messages105        :param model_parameters: model parameters106        :param tools: tools for tool calling107        :param stop: stop words108        :param stream: is stream response109        :param user: unique user id110        """111        raise NotImplementedError()112 113    @abstractmethod114    def on_invoke_error(115        self,116        llm_instance: AIModel,117        ex: Exception,118        model: str,119        credentials: dict,120        prompt_messages: list[PromptMessage],121        model_parameters: dict,122        tools: Optional[list[PromptMessageTool]] = None,123        stop: Optional[list[str]] = None,124        stream: bool = True,125        user: Optional[str] = None,126    ) -> None:127        """128        Invoke error callback129 130        :param llm_instance: LLM instance131        :param ex: exception132        :param model: model name133        :param credentials: model credentials134        :param prompt_messages: prompt messages135        :param model_parameters: model parameters136        :param tools: tools for tool calling137        :param stop: stop words138        :param stream: is stream response139        :param user: unique user id140        """141        raise NotImplementedError()142 143    def print_text(self, text: str, color: Optional[str] = None, end: str = "") -> None:144        """Print text with highlighting and no end characters."""145        text_to_print = self._get_colored_text(text, color) if color else text146        print(text_to_print, end=end)147 148    def _get_colored_text(self, text: str, color: str) -> str:149        """Get colored text."""150        color_str = _TEXT_COLOR_MAPPING[color]151        return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m"152