🔧 refactor(constants.py, conftest.py): improve code structure and readability

🔨 refactor(constants.py): make YourComponent inherit from CustomComponent for better code reuse and structure

🔨 refactor(conftest.py): restructure custom_chain fixture for better readability and maintainability

src/backend/langflow/interface/custom/constants.py

@@ -25,17 +25,19 @@ LANGCHAIN_BASE_TYPES = {
 
 DEFAULT_CUSTOM_COMPONENT_CODE = """
 from langflow import Prompt
+from langflow.interface.custom.custom_component import CustomComponent
+
 from langchain.llms.base import BaseLLM
 from langchain.chains import LLMChain
-from langflow.interface.custom import CustomComponent
 from langchain import PromptTemplate
 from langchain.schema import Document
+
 import requests
 
-class YourComponent:
+class YourComponent(CustomComponent):
-    display_name: str = "Your Component"
+    #display_name: str = "Your Component"
-    description: str = "Your description"
+    #description: str = "Your description"
-    field_config = { "url": { "multiline": True, "required": True } }
+    #field_config = { "url": { "multiline": True, "required": True } }
 
     def build(self, url: str, llm: BaseLLM, template: Prompt) -> Document:
         response = requests.get(url)
@@ -44,7 +46,3 @@ class YourComponent:
         result = chain.run(response.text[:300])
         return Document(page_content=str(result))
 """
-
-
-# Create a new class that can be used as a type
-# that returns type "prompt" if we get a certain param

tests/conftest.py

@@ -120,22 +120,22 @@ def client_fixture(session: Session):  #
 
 @pytest.fixture
 def custom_chain():
-    return '''from __future__ import annotations
+    return '''
+from __future__ import annotations
+from typing import Any, Dict, List, Optional
 
-    from typing import Any, Dict, List, Optional
+from pydantic import Extra
 
-    from pydantic import Extra
+from langchain.schema import BaseLanguageModel, Document
-
+from langchain.callbacks.manager import (
-    from langchain.schema import BaseLanguageModel, Document
-    from langchain.callbacks.manager import (
     AsyncCallbackManagerForChainRun,
     CallbackManagerForChainRun,
-    )
+)
-    from langchain.chains.base import Chain
+from langchain.chains.base import Chain
-    from langchain.prompts import StringPromptTemplate
+from langchain.prompts import StringPromptTemplate
-    from langflow.interface.custom.base import CustomComponent
+from langflow.interface.custom.base import CustomComponent
 
-    class MyCustomChain(Chain):
+class MyCustomChain(Chain):
     """
     An example of a custom chain.
     """
@@ -225,13 +225,14 @@ def custom_chain():
     def _chain_type(self) -> str:
         return "my_custom_chain"
 
-    class CustomChain(CustomComponent):
+class CustomChain(CustomComponent):
     display_name: str = "Custom Chain"
     field_config = {
         "prompt": {"field_type": "prompt"},
         "llm": {"field_type": "BaseLanguageModel"},
     }
 
-        def build(self, prompt, llm, input: str) -> Document:
+    def build(self, prompt: StringPromptTemplate, llm: BaseLanguageModel, input: str) -> Document:
         chain = MyCustomChain(prompt=prompt, llm=llm)
-            return chain(input)'''
+        return chain(input)
+'''