Underground-Digital/Workflow-Engine

1import json2import logging3import sys4from typing import Optional5 6from core.model_runtime.callbacks.base_callback import Callback7from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk8from core.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool9from core.model_runtime.model_providers.__base.ai_model import AIModel10 11logger = logging.getLogger(__name__)12 13 14class LoggingCallback(Callback):15    def on_before_invoke(16        self,17        llm_instance: AIModel,18        model: str,19        credentials: dict,20        prompt_messages: list[PromptMessage],21        model_parameters: dict,22        tools: Optional[list[PromptMessageTool]] = None,23        stop: Optional[list[str]] = None,24        stream: bool = True,25        user: Optional[str] = None,26    ) -> None:27        """28        Before invoke callback29 30        :param llm_instance: LLM instance31        :param model: model name32        :param credentials: model credentials33        :param prompt_messages: prompt messages34        :param model_parameters: model parameters35        :param tools: tools for tool calling36        :param stop: stop words37        :param stream: is stream response38        :param user: unique user id39        """40        self.print_text("\n[on_llm_before_invoke]\n", color="blue")41        self.print_text(f"Model: {model}\n", color="blue")42        self.print_text("Parameters:\n", color="blue")43        for key, value in model_parameters.items():44            self.print_text(f"\t{key}: {value}\n", color="blue")45 46        if stop:47            self.print_text(f"\tstop: {stop}\n", color="blue")48 49        if tools:50            self.print_text("\tTools:\n", color="blue")51            for tool in tools:52                self.print_text(f"\t\t{tool.name}\n", color="blue")53 54        self.print_text(f"Stream: {stream}\n", color="blue")55 56        if user:57            self.print_text(f"User: {user}\n", color="blue")58 59        self.print_text("Prompt messages:\n", color="blue")60        for prompt_message in prompt_messages:61            if prompt_message.name:62                self.print_text(f"\tname: {prompt_message.name}\n", color="blue")63 64            self.print_text(f"\trole: {prompt_message.role.value}\n", color="blue")65            self.print_text(f"\tcontent: {prompt_message.content}\n", color="blue")66 67        if stream:68            self.print_text("\n[on_llm_new_chunk]")69 70    def on_new_chunk(71        self,72        llm_instance: AIModel,73        chunk: LLMResultChunk,74        model: str,75        credentials: dict,76        prompt_messages: list[PromptMessage],77        model_parameters: dict,78        tools: Optional[list[PromptMessageTool]] = None,79        stop: Optional[list[str]] = None,80        stream: bool = True,81        user: Optional[str] = None,82    ):83        """84        On new chunk callback85 86        :param llm_instance: LLM instance87        :param chunk: chunk88        :param model: model name89        :param credentials: model credentials90        :param prompt_messages: prompt messages91        :param model_parameters: model parameters92        :param tools: tools for tool calling93        :param stop: stop words94        :param stream: is stream response95        :param user: unique user id96        """97        sys.stdout.write(chunk.delta.message.content)98        sys.stdout.flush()99 100    def on_after_invoke(101        self,102        llm_instance: AIModel,103        result: LLMResult,104        model: str,105        credentials: dict,106        prompt_messages: list[PromptMessage],107        model_parameters: dict,108        tools: Optional[list[PromptMessageTool]] = None,109        stop: Optional[list[str]] = None,110        stream: bool = True,111        user: Optional[str] = None,112    ) -> None:113        """114        After invoke callback115 116        :param llm_instance: LLM instance117        :param result: result118        :param model: model name119        :param credentials: model credentials120        :param prompt_messages: prompt messages121        :param model_parameters: model parameters122        :param tools: tools for tool calling123        :param stop: stop words124        :param stream: is stream response125        :param user: unique user id126        """127        self.print_text("\n[on_llm_after_invoke]\n", color="yellow")128        self.print_text(f"Content: {result.message.content}\n", color="yellow")129 130        if result.message.tool_calls:131            self.print_text("Tool calls:\n", color="yellow")132            for tool_call in result.message.tool_calls:133                self.print_text(f"\t{tool_call.id}\n", color="yellow")134                self.print_text(f"\t{tool_call.function.name}\n", color="yellow")135                self.print_text(f"\t{json.dumps(tool_call.function.arguments)}\n", color="yellow")136 137        self.print_text(f"Model: {result.model}\n", color="yellow")138        self.print_text(f"Usage: {result.usage}\n", color="yellow")139        self.print_text(f"System Fingerprint: {result.system_fingerprint}\n", color="yellow")140 141    def on_invoke_error(142        self,143        llm_instance: AIModel,144        ex: Exception,145        model: str,146        credentials: dict,147        prompt_messages: list[PromptMessage],148        model_parameters: dict,149        tools: Optional[list[PromptMessageTool]] = None,150        stop: Optional[list[str]] = None,151        stream: bool = True,152        user: Optional[str] = None,153    ) -> None:154        """155        Invoke error callback156 157        :param llm_instance: LLM instance158        :param ex: exception159        :param model: model name160        :param credentials: model credentials161        :param prompt_messages: prompt messages162        :param model_parameters: model parameters163        :param tools: tools for tool calling164        :param stop: stop words165        :param stream: is stream response166        :param user: unique user id167        """168        self.print_text("\n[on_llm_invoke_error]\n", color="red")169        logger.exception(ex)170