mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from typing import Optional
|
|
|
|
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'),
|
|
callbacks=[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",
|
|
callbacks=[DifyStdOutCallbackHandler()],
|
|
**kwargs
|
|
)
|
|
|
|
return None
|