Model Runtime (#1858)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com> Co-authored-by: Garfield Dai <dai.hai@foxmail.com> Co-authored-by: chenhe <guchenhe@gmail.com> Co-authored-by: jyong <jyong@dify.ai> Co-authored-by: Joel <iamjoel007@gmail.com> Co-authored-by: Yeuoly <admin@srmxy.cn>
2025-12-15 22:06:52 +08:00 · 2024-01-02 23:42:00 +08:00
parent e91dd28a76
commit d069c668f8
807 changed files with 171310 additions and 23806 deletions
--- a/api/core/model_runtime/callbacks/init.py
+++ b/api/core/model_runtime/callbacks/init.py
--- a/api/core/model_runtime/callbacks/base_callback.py
+++ b/api/core/model_runtime/callbacks/base_callback.py
@@ -0,0 +1,113 @@
+from abc import ABC
+from typing import Optional, List
+
+from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk
+from core.model_runtime.entities.message_entities import PromptMessageTool, PromptMessage
+from core.model_runtime.model_providers.__base.ai_model import AIModel
+
+_TEXT_COLOR_MAPPING = {
+    "blue": "36;1",
+    "yellow": "33;1",
+    "pink": "38;5;200",
+    "green": "32;1",
+    "red": "31;1",
+}
+
+
+class Callback(ABC):
+    """
+    Base class for callbacks.
+    Only for LLM.
+    """
+    raise_error: bool = False
+
+    def on_before_invoke(self, llm_instance: AIModel, model: str, credentials: dict,
+                         prompt_messages: list[PromptMessage], model_parameters: dict,
+                         tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                         stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        Before invoke callback
+
+        :param llm_instance: LLM instance
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        raise NotImplementedError()
+
+    def on_new_chunk(self, llm_instance: AIModel, chunk: LLMResultChunk, model: str, credentials: dict,
+                     prompt_messages: list[PromptMessage], model_parameters: dict,
+                     tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                     stream: bool = True, user: Optional[str] = None):
+        """
+        On new chunk callback
+
+        :param llm_instance: LLM instance
+        :param chunk: chunk
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        raise NotImplementedError()
+
+    def on_after_invoke(self, llm_instance: AIModel, result: LLMResult, model: str, credentials: dict,
+                        prompt_messages: list[PromptMessage], model_parameters: dict,
+                        tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                        stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        After invoke callback
+
+        :param llm_instance: LLM instance
+        :param result: result
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        raise NotImplementedError()
+
+    def on_invoke_error(self, llm_instance: AIModel, ex: Exception, model: str, credentials: dict,
+                        prompt_messages: list[PromptMessage], model_parameters: dict,
+                        tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                        stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        Invoke error callback
+
+        :param llm_instance: LLM instance
+        :param ex: exception
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        raise NotImplementedError()
+
+    def print_text(
+            self, text: str, color: Optional[str] = None, end: str = ""
+    ) -> None:
+        """Print text with highlighting and no end characters."""
+        text_to_print = self._get_colored_text(text, color) if color else text
+        print(text_to_print, end=end)
+
+    def _get_colored_text(self, text: str, color: str) -> str:
+        """Get colored text."""
+        color_str = _TEXT_COLOR_MAPPING[color]
+        return f"\u001b[{color_str}m\033[1;3m{text}\u001b[0m"
--- a/api/core/model_runtime/callbacks/logging_callback.py
+++ b/api/core/model_runtime/callbacks/logging_callback.py
@@ -0,0 +1,133 @@
+import json
+import logging
+import sys
+from typing import Optional, List
+
+from core.model_runtime.callbacks.base_callback import Callback
+from core.model_runtime.entities.llm_entities import LLMResultChunk, LLMResult
+from core.model_runtime.entities.message_entities import PromptMessageTool, PromptMessage
+from core.model_runtime.model_providers.__base.ai_model import AIModel
+
+logger = logging.getLogger(__name__)
+
+class LoggingCallback(Callback):
+    def on_before_invoke(self, llm_instance: AIModel, model: str, credentials: dict,
+                         prompt_messages: list[PromptMessage], model_parameters: dict,
+                         tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                         stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        Before invoke callback
+
+        :param llm_instance: LLM instance
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        self.print_text("\n[on_llm_before_invoke]\n", color='blue')
+        self.print_text(f"Model: {model}\n", color='blue')
+        self.print_text(f"Parameters:\n", color='blue')
+        for key, value in model_parameters.items():
+            self.print_text(f"\t{key}: {value}\n", color='blue')
+
+        if stop:
+            self.print_text(f"\tstop: {stop}\n", color='blue')
+
+        if tools:
+            self.print_text(f"\tTools:\n", color='blue')
+            for tool in tools:
+                self.print_text(f"\t\t{tool.name}\n", color='blue')
+
+        self.print_text(f"Stream: {stream}\n", color='blue')
+
+        if user:
+            self.print_text(f"User: {user}\n", color='blue')
+
+        self.print_text(f"Prompt messages:\n", color='blue')
+        for prompt_message in prompt_messages:
+            if prompt_message.name:
+                self.print_text(f"\tname: {prompt_message.name}\n", color='blue')
+
+            self.print_text(f"\trole: {prompt_message.role.value}\n", color='blue')
+            self.print_text(f"\tcontent: {prompt_message.content}\n", color='blue')
+
+        if stream:
+            self.print_text("\n[on_llm_new_chunk]")
+
+    def on_new_chunk(self, llm_instance: AIModel, chunk: LLMResultChunk, model: str, credentials: dict,
+                     prompt_messages: list[PromptMessage], model_parameters: dict,
+                     tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                     stream: bool = True, user: Optional[str] = None):
+        """
+        On new chunk callback
+
+        :param llm_instance: LLM instance
+        :param chunk: chunk
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        sys.stdout.write(chunk.delta.message.content)
+        sys.stdout.flush()
+
+    def on_after_invoke(self, llm_instance: AIModel, result: LLMResult, model: str, credentials: dict,
+                        prompt_messages: list[PromptMessage], model_parameters: dict,
+                        tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                        stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        After invoke callback
+
+        :param llm_instance: LLM instance
+        :param result: result
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        self.print_text("\n[on_llm_after_invoke]\n", color='yellow')
+        self.print_text(f"Content: {result.message.content}\n", color='yellow')
+
+        if result.message.tool_calls:
+            self.print_text(f"Tool calls:\n", color='yellow')
+            for tool_call in result.message.tool_calls:
+                self.print_text(f"\t{tool_call.id}\n", color='yellow')
+                self.print_text(f"\t{tool_call.function.name}\n", color='yellow')
+                self.print_text(f"\t{json.dumps(tool_call.function.arguments)}\n", color='yellow')
+
+        self.print_text(f"Model: {result.model}\n", color='yellow')
+        self.print_text(f"Usage: {result.usage}\n", color='yellow')
+        self.print_text(f"System Fingerprint: {result.system_fingerprint}\n", color='yellow')
+
+    def on_invoke_error(self, llm_instance: AIModel, ex: Exception, model: str, credentials: dict,
+                        prompt_messages: list[PromptMessage], model_parameters: dict,
+                        tools: Optional[list[PromptMessageTool]] = None, stop: Optional[List[str]] = None,
+                        stream: bool = True, user: Optional[str] = None) -> None:
+        """
+        Invoke error callback
+
+        :param llm_instance: LLM instance
+        :param ex: exception
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param tools: tools for tool calling
+        :param stop: stop words
+        :param stream: is stream response
+        :param user: unique user id
+        """
+        self.print_text("\n[on_llm_invoke_error]\n", color='red')
+        logger.exception(ex)