feat: migrate chains and memories to Component syntax (langflow-ai#2528)

* feat: migrate chains and memories to Component syntax

* use base class

* add classes

* [autofix.ci] apply automated fixes

* fix tests

* fix tests

* ✅ (filterSidebar.spec.ts): increase waitForTimeout from 1000ms to 2000ms to ensure elements are fully loaded before interaction

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <[email protected]>
Co-authored-by: cristhianzl <[email protected]>
Co-authored-by: Cristhian Zanforlin Lousa <[email protected]>

Author: 4 people

src/backend/base/langflow/base/chains/__init__.py

@@ -0,0 +1,17 @@
+from langflow.custom import Component
+from langflow.template import Output
+
+
+class LCChainComponent(Component):
+    trace_type = "chain"
+
+    outputs = [Output(display_name="Text", name="text", method="invoke_chain")]
+
+    def _validate_outputs(self):
+        required_output_methods = ["invoke_chain"]
+        output_names = [output.name for output in self.outputs]
+        for method_name in required_output_methods:
+            if method_name not in output_names:
+                raise ValueError(f"Output with name '{method_name}' must be defined.")
+            elif not hasattr(self, method_name):
+                raise ValueError(f"Method '{method_name}' must be defined.")

src/backend/base/langflow/base/chains/model.py

@@ -0,0 +1,35 @@
+from abc import abstractmethod
+
+from langflow.custom import Component
+from langflow.field_typing import BaseChatMessageHistory, BaseChatMemory
+from langflow.template import Output
+from langchain.memory import ConversationBufferMemory
+
+
+class LCChatMemoryComponent(Component):
+    trace_type = "chat_memory"
+    outputs = [
+        Output(
+            display_name="Memory",
+            name="base_memory",
+            method="build_base_memory",
+        )
+    ]
+
+    def _validate_outputs(self):
+        required_output_methods = ["build_base_memory"]
+        output_names = [output.name for output in self.outputs]
+        for method_name in required_output_methods:
+            if method_name not in output_names:
+                raise ValueError(f"Output with name '{method_name}' must be defined.")
+            elif not hasattr(self, method_name):
+                raise ValueError(f"Method '{method_name}' must be defined.")
+
+    def build_base_memory(self) -> BaseChatMemory:
+        return ConversationBufferMemory(chat_memory=self.build_message_history())
+
+    @abstractmethod
+    def build_message_history(self) -> BaseChatMessageHistory:
+        """
+        Builds the chat message history memory.
+        """

src/backend/base/langflow/base/memory/model.py

@@ -1,47 +1,41 @@
-from typing import Optional
-
 from langchain.chains import ConversationChain
 
-from langflow.custom import CustomComponent
-from langflow.field_typing import BaseMemory, LanguageModel, Text
+from langflow.base.chains.model import LCChainComponent
+from langflow.field_typing import Message
+from langflow.inputs import MultilineInput, HandleInput
 
 
-class ConversationChainComponent(CustomComponent):
+class ConversationChainComponent(LCChainComponent):
     display_name = "ConversationChain"
     description = "Chain to have a conversation and load context from memory."
     name = "ConversationChain"
 
-    def build_config(self):
-        return {
-            "prompt": {"display_name": "Prompt"},
-            "llm": {"display_name": "LLM"},
-            "memory": {
-                "display_name": "Memory",
-                "info": "Memory to load context from. If none is provided, a ConversationBufferMemory will be used.",
-            },
-            "input_value": {
-                "display_name": "Input Value",
-                "info": "The input value to pass to the chain.",
-            },
-        }
+    inputs = [
+        MultilineInput(
+            name="input_value", display_name="Input", info="The input value to pass to the chain.", required=True
+        ),
+        HandleInput(name="llm", display_name="Language Model", input_types=["LanguageModel"], required=True),
+        HandleInput(
+            name="memory",
+            display_name="Memory",
+            input_types=["BaseChatMemory"],
+        ),
+    ]
 
-    def build(
-        self,
-        input_value: Text,
-        llm: LanguageModel,
-        memory: Optional[BaseMemory] = None,
-    ) -> Text:
-        if memory is None:
-            chain = ConversationChain(llm=llm)
+    def invoke_chain(self) -> Message:
+        if not self.memory:
+            chain = ConversationChain(llm=self.llm)
         else:
-            chain = ConversationChain(llm=llm, memory=memory)
-        result = chain.invoke({"input": input_value})
+            chain = ConversationChain(llm=self.llm, memory=self.memory)
+
+        result = chain.invoke({"input": self.input_value})
         if isinstance(result, dict):
             result = result.get(chain.output_key, "")  # type: ignore
 
         elif isinstance(result, str):
             result = result
         else:
             result = result.get("response")
+        result = str(result)
         self.status = result
-        return str(result)
+        return Message(text=result)

src/backend/base/langflow/components/chains/ConversationChain.py

@@ -1,32 +1,27 @@
 from langchain.chains import LLMCheckerChain
 
-from langflow.custom import CustomComponent
-from langflow.field_typing import LanguageModel, Text
+from langflow.base.chains.model import LCChainComponent
+from langflow.field_typing import Message
+from langflow.inputs import MultilineInput, HandleInput
 
 
-class LLMCheckerChainComponent(CustomComponent):
+class LLMCheckerChainComponent(LCChainComponent):
     display_name = "LLMCheckerChain"
-    description = ""
+    description = "Chain for question-answering with self-verification."
     documentation = "https://python.langchain.com/docs/modules/chains/additional/llm_checker"
     name = "LLMCheckerChain"
 
-    def build_config(self):
-        return {
-            "llm": {"display_name": "LLM"},
-            "input_value": {
-                "display_name": "Input Value",
-                "info": "The input value to pass to the chain.",
-            },
-        }
+    inputs = [
+        MultilineInput(
+            name="input_value", display_name="Input", info="The input value to pass to the chain.", required=True
+        ),
+        HandleInput(name="llm", display_name="Language Model", input_types=["LanguageModel"], required=True),
+    ]
 
-    def build(
-        self,
-        input_value: Text,
-        llm: LanguageModel,
-    ) -> Text:
-        chain = LLMCheckerChain.from_llm(llm=llm)
-        response = chain.invoke({chain.input_key: input_value})
+    def invoke_chain(self) -> Message:
+        chain = LLMCheckerChain.from_llm(llm=self.llm)
+        response = chain.invoke({chain.input_key: self.input_value})
         result = response.get(chain.output_key, "")
-        result_str = str(result)
-        self.status = result_str
-        return result_str
+        result = str(result)
+        self.status = result
+        return Message(text=result)

src/backend/base/langflow/components/chains/LLMChain.py

@@ -1,48 +1,30 @@
-from typing import Optional
+from langchain.chains import LLMMathChain
 
-from langchain.chains import LLMChain, LLMMathChain
+from langflow.base.chains.model import LCChainComponent
+from langflow.field_typing import Message
+from langflow.inputs import MultilineInput, HandleInput
+from langflow.template import Output
 
-from langflow.custom import CustomComponent
-from langflow.field_typing import BaseMemory, LanguageModel, Text
 
-
-class LLMMathChainComponent(CustomComponent):
+class LLMMathChainComponent(LCChainComponent):
     display_name = "LLMMathChain"
     description = "Chain that interprets a prompt and executes python code to do math."
     documentation = "https://python.langchain.com/docs/modules/chains/additional/llm_math"
     name = "LLMMathChain"
 
-    def build_config(self):
-        return {
-            "llm": {"display_name": "LLM"},
-            "llm_chain": {"display_name": "LLM Chain"},
-            "memory": {"display_name": "Memory"},
-            "input_key": {"display_name": "Input Key"},
-            "output_key": {"display_name": "Output Key"},
-            "input_value": {
-                "display_name": "Input Value",
-                "info": "The input value to pass to the chain.",
-            },
-        }
+    inputs = [
+        MultilineInput(
+            name="input_value", display_name="Input", info="The input value to pass to the chain.", required=True
+        ),
+        HandleInput(name="llm", display_name="Language Model", input_types=["LanguageModel"], required=True),
+    ]
+
+    outputs = [Output(display_name="Text", name="text", method="invoke_chain")]
 
-    def build(
-        self,
-        input_value: Text,
-        llm: LanguageModel,
-        llm_chain: LLMChain,
-        input_key: str = "question",
-        output_key: str = "answer",
-        memory: Optional[BaseMemory] = None,
-    ) -> Text:
-        chain = LLMMathChain(
-            llm=llm,
-            llm_chain=llm_chain,
-            input_key=input_key,
-            output_key=output_key,
-            memory=memory,
-        )
-        response = chain.invoke({input_key: input_value})
-        result = response.get(output_key)
-        result_str = str(result)
-        self.status = result_str
-        return result_str
+    def invoke_chain(self) -> Message:
+        chain = LLMMathChain.from_llm(llm=self.llm)
+        response = chain.invoke({chain.input_key: self.input_value})
+        result = response.get(chain.output_key, "")
+        result = str(result)
+        self.status = result
+        return Message(text=result)