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

api/core/chain/main_chain_builder.py

@@ -0,0 +1,116 @@
+from typing import Optional, List
+
+from langchain.callbacks import SharedCallbackManager
+from langchain.chains import SequentialChain
+from langchain.chains.base import Chain
+from langchain.memory.chat_memory import BaseChatMemory
+
+from core.agent.agent_builder import AgentBuilder
+from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
+from core.callback_handler.dataset_tool_callback_handler import DatasetToolCallbackHandler
+from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
+from core.chain.chain_builder import ChainBuilder
+from core.constant import llm_constant
+from core.conversation_message_task import ConversationMessageTask
+from core.tool.dataset_tool_builder import DatasetToolBuilder
+
+
+class MainChainBuilder:
+    @classmethod
+    def to_langchain_components(cls, tenant_id: str, agent_mode: dict, memory: Optional[BaseChatMemory],
+                                conversation_message_task: ConversationMessageTask):
+        first_input_key = "input"
+        final_output_key = "output"
+
+        chains = []
+
+        chain_callback_handler = MainChainGatherCallbackHandler(conversation_message_task)
+
+        # agent mode
+        tool_chains, chains_output_key = cls.get_agent_chains(
+            tenant_id=tenant_id,
+            agent_mode=agent_mode,
+            memory=memory,
+            dataset_tool_callback_handler=DatasetToolCallbackHandler(conversation_message_task),
+            agent_loop_gather_callback_handler=chain_callback_handler.agent_loop_gather_callback_handler
+        )
+        chains += tool_chains
+
+        if chains_output_key:
+            final_output_key = chains_output_key
+
+        if len(chains) == 0:
+            return None
+
+        for chain in chains:
+            # do not add handler into singleton callback manager
+            if not isinstance(chain.callback_manager, SharedCallbackManager):
+                chain.callback_manager.add_handler(chain_callback_handler)
+
+        # build main chain
+        overall_chain = SequentialChain(
+            chains=chains,
+            input_variables=[first_input_key],
+            output_variables=[final_output_key],
+            memory=memory,  # only for use the memory prompt input key
+        )
+
+        return overall_chain
+
+    @classmethod
+    def get_agent_chains(cls, tenant_id: str, agent_mode: dict, memory: Optional[BaseChatMemory],
+                         dataset_tool_callback_handler: DatasetToolCallbackHandler,
+                         agent_loop_gather_callback_handler: AgentLoopGatherCallbackHandler):
+        # agent mode
+        chains = []
+        if agent_mode and agent_mode.get('enabled'):
+            tools = agent_mode.get('tools', [])
+
+            pre_fixed_chains = []
+            agent_tools = []
+            for tool in tools:
+                tool_type = list(tool.keys())[0]
+                tool_config = list(tool.values())[0]
+                if tool_type == 'sensitive-word-avoidance':
+                    chain = ChainBuilder.to_sensitive_word_avoidance_chain(tool_config)
+                    if chain:
+                        pre_fixed_chains.append(chain)
+                elif tool_type == "dataset":
+                    dataset_tool = DatasetToolBuilder.build_dataset_tool(
+                        tenant_id=tenant_id,
+                        dataset_id=tool_config.get("id"),
+                        response_mode='no_synthesizer',  # "compact"
+                        callback_handler=dataset_tool_callback_handler
+                    )
+
+                    if dataset_tool:
+                        agent_tools.append(dataset_tool)
+
+            # add pre-fixed chains
+            chains += pre_fixed_chains
+
+            if len(agent_tools) == 1:
+                # tool to chain
+                tool_chain = ChainBuilder.to_tool_chain(tool=agent_tools[0], output_key='tool_output')
+                chains.append(tool_chain)
+            elif len(agent_tools) > 1:
+                # build agent config
+                agent_chain = AgentBuilder.to_agent_chain(
+                    tenant_id=tenant_id,
+                    tools=agent_tools,
+                    memory=memory,
+                    dataset_tool_callback_handler=dataset_tool_callback_handler,
+                    agent_loop_gather_callback_handler=agent_loop_gather_callback_handler
+                )
+
+                chains.append(agent_chain)
+
+        final_output_key = cls.get_chains_output_key(chains)
+
+        return chains, final_output_key
+
+    @classmethod
+    def get_chains_output_key(cls, chains: List[Chain]):
+        if len(chains) > 0:
+            return chains[-1].output_keys[0]
+        return None

api/core/chain/sensitive_word_avoidance_chain.py

@@ -0,0 +1,42 @@
+from typing import List, Dict
+
+from langchain.chains.base import Chain
+
+
+class SensitiveWordAvoidanceChain(Chain):
+    input_key: str = "input"  #: :meta private:
+    output_key: str = "output"  #: :meta private:
+
+    sensitive_words: List[str] = []
+    canned_response: str = None
+
+    @property
+    def _chain_type(self) -> str:
+        return "sensitive_word_avoidance_chain"
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Expect input key.
+
+        :meta private:
+        """
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Return output key.
+
+        :meta private:
+        """
+        return [self.output_key]
+
+    def _check_sensitive_word(self, text: str) -> str:
+        for word in self.sensitive_words:
+            if word in text:
+                return self.canned_response
+        return text
+
+    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        text = inputs[self.input_key]
+        output = self._check_sensitive_word(text)
+        return {self.output_key: output}

api/core/chain/tool_chain.py

@@ -0,0 +1,42 @@
+from typing import List, Dict
+
+from langchain.chains.base import Chain
+from langchain.tools import BaseTool
+
+
+class ToolChain(Chain):
+    input_key: str = "input"  #: :meta private:
+    output_key: str = "output"  #: :meta private:
+
+    tool: BaseTool
+
+    @property
+    def _chain_type(self) -> str:
+        return "tool_chain"
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Expect input key.
+
+        :meta private:
+        """
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Return output key.
+
+        :meta private:
+        """
+        return [self.output_key]
+
+    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        input = inputs[self.input_key]
+        output = self.tool.run(input, self.verbose)
+        return {self.output_key: output}
+
+    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        """Run the logic of this chain and return the output."""
+        input = inputs[self.input_key]
+        output = await self.tool.arun(input, self.verbose)
+        return {self.output_key: output}