Initial commit

api/core/chain/chain_builder.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from langchain.callbacks import CallbackManager
+
+from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
+from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceChain
+from core.chain.tool_chain import ToolChain
+
+
+class ChainBuilder:
+    @classmethod
+    def to_tool_chain(cls, tool, **kwargs) -> ToolChain:
+        return ToolChain(
+            tool=tool,
+            input_key=kwargs.get('input_key', 'input'),
+            output_key=kwargs.get('output_key', 'tool_output'),
+            callback_manager=CallbackManager([DifyStdOutCallbackHandler()])
+        )
+
+    @classmethod
+    def to_sensitive_word_avoidance_chain(cls, tool_config: dict, **kwargs) -> Optional[
+        SensitiveWordAvoidanceChain]:
+        sensitive_words = tool_config.get("words", "")
+        if tool_config.get("enabled", False) \
+                and sensitive_words:
+            return SensitiveWordAvoidanceChain(
+                sensitive_words=sensitive_words.split(","),
+                canned_response=tool_config.get("canned_response", ''),
+                output_key="sensitive_word_avoidance_output",
+                callback_manager=CallbackManager([DifyStdOutCallbackHandler()]),
+                **kwargs
+            )
+
+        return None