From bf650ecec56eb5b6fcf4bad2204e2201a2087d25 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Wed, 21 Aug 2024 09:34:15 -0300
Subject: [PATCH] refactor: improve attribute compatibility and error messages
 (#3367)

* feat: Refactor how configuration attributes are handled in Component class.

* refactor: Update method name to handle singular and plural inputs.

* refactor: Add value validation for input and attribute assignment.

* test: Add test for setting invalid input in task component.

* fix: add param without underscore to inputs

* test(astra): Refactor AstraDB imports in test_astra_component (#3413)

* test(astra): Refactor AstraDB imports in test_astra_component

* fix: Add AstraVectorizeComponent to AstraDB tests.

* fix: Refactor custom component error message for setting parameter or attribute.

* test(astra): Refactor AstraDB imports in test_astra_component (#3413)

* test(astra): Refactor AstraDB imports in test_astra_component

* fix: Add AstraVectorizeComponent to AstraDB tests.

* chore: Refactor joining methods list in _set_parameter_or_attribute function.

* refactor: Refactor error message string formatting in Component class.
---
 .../custom/custom_component/component.py      | 25 ++++++++++++++++++-
 .../custom/custom_component/test_component.py | 12 +++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/backend/base/langflow/custom/custom_component/component.py b/src/backend/base/langflow/custom/custom_component/component.py
index c78856d7f..880af91df 100644
--- a/src/backend/base/langflow/custom/custom_component/component.py
+++ b/src/backend/base/langflow/custom/custom_component/component.py
@@ -26,6 +26,7 @@ if TYPE_CHECKING:
     from langflow.inputs.inputs import InputTypes
 
 BACKWARDS_COMPATIBLE_ATTRIBUTES = ["user_id", "vertex", "tracing_service"]
+CONFIG_ATTRIBUTES = ["_display_name", "_description", "_icon", "_name"]
 
 
 class Component(CustomComponent):
@@ -43,6 +44,8 @@ class Component(CustomComponent):
         for key, value in kwargs.items():
             if key.startswith("_"):
                 config[key] = value
+            elif key in CONFIG_ATTRIBUTES:
+                config[key[1:]] = value
             else:
                 inputs[key] = value
         self._inputs: dict[str, "InputTypes"] = {}
@@ -142,7 +145,7 @@ class Component(CustomComponent):
             KeyError: If the specified input name does not exist.
         """
         for key, value in kwargs.items():
-            self._process_connection_or_parameter(key, value)
+            self._process_connection_or_parameters(key, value)
         return self
 
     def list_inputs(self):
@@ -320,6 +323,14 @@ class Component(CustomComponent):
         else:
             self._set_parameter_or_attribute(key, value)
 
+    def _process_connection_or_parameters(self, key, value):
+        # if value is a list of components, we need to process each component
+        if isinstance(value, list):
+            for val in value:
+                self._process_connection_or_parameter(key, val)
+        else:
+            self._process_connection_or_parameter(key, value)
+
     def _get_or_create_input(self, key):
         try:
             return self._inputs[key]
@@ -358,6 +369,12 @@ class Component(CustomComponent):
         )
 
     def _set_parameter_or_attribute(self, key, value):
+        if isinstance(value, Component):
+            methods = ", ".join([f"'{output.method}'" for output in value.outputs])
+            raise ValueError(
+                f"You set {value.display_name} as value for `{key}`. "
+                f"You should pass one of the following: {methods}"
+            )
         self._set_input_value(key, value)
         self._parameters[key] = value
         self._attributes[key] = value
@@ -396,6 +413,12 @@ class Component(CustomComponent):
     def _set_input_value(self, name: str, value: Any):
         if name in self._inputs:
             input_value = self._inputs[name].value
+            if isinstance(input_value, Component):
+                methods = ", ".join([f"'{output.method}'" for output in input_value.outputs])
+                raise ValueError(
+                    f"You set {input_value.display_name} as value for `{name}`. "
+                    f"You should pass one of the following: {methods}"
+                )
             if callable(input_value):
                 raise ValueError(
                     f"Input {name} is connected to {input_value.__self__.display_name}.{input_value.__name__}"
diff --git a/src/backend/tests/unit/custom/custom_component/test_component.py b/src/backend/tests/unit/custom/custom_component/test_component.py
index f28ce861b..b52d993e0 100644
--- a/src/backend/tests/unit/custom/custom_component/test_component.py
+++ b/src/backend/tests/unit/custom/custom_component/test_component.py
@@ -1,5 +1,7 @@
 import pytest
 
+from langflow.components.agents.CrewAIAgent import CrewAIAgentComponent
+from langflow.components.helpers.SequentialTask import SequentialTaskComponent
 from langflow.components.inputs.ChatInput import ChatInput
 from langflow.components.outputs import ChatOutput
 
@@ -14,3 +16,13 @@ def test_set_invalid_output():
     chatoutput = ChatOutput()
     with pytest.raises(ValueError):
         chatoutput.set(input_value=chatinput.build_config)
+
+
+def test_set_invalid_input():
+    crewai_agent = CrewAIAgentComponent()
+    task = SequentialTaskComponent()
+    with pytest.raises(
+        ValueError,
+        match="You set CrewAI Agent as value for `agent`. You should pass one of the following: 'build_output'",
+    ):
+        task.set(agent=crewai_agent)