refactor: enhance tool creation logic and add FeatureFlags (#3662)

* Add `required_inputs` field to `Output` model in `base.py`

* Refactor ComponentTool to ComponentToolkit and enhance tool creation logic

- Replaced `ComponentTool` with `ComponentToolkit` to better encapsulate component-related tools.
- Introduced `build_description` and `_build_output_function` helper functions for dynamic tool creation.
- Updated tool initialization to handle multiple outputs and required inputs using `StructuredTool`.
- Improved schema creation for tool arguments based on component inputs.

* Refactor `to_tool` method to `to_toolkit` to use `ComponentToolkit` instead of `ComponentTool`

* Refactor `ComponentTool` to `ComponentToolkit` in unit tests

- Updated import statements to reflect the new `ComponentToolkit` class.
- Modified test logic to use `ComponentToolkit` for retrieving tools.
- Adjusted assertions to match the new structure and output format.
- Ensured compatibility with `Message` schema for output validation.

* Refactor `test_component_to_tool` to validate `ComponentToolkit` and tool properties

* Refactor `build_description` to include input types in the output format

* Add method to set required inputs for outputs based on method analysis

- Introduced `_set_output_required_inputs` method to determine and set required inputs for each output by analyzing the method's source code.
- Added necessary imports (`ast` and `dedent`) to support the new functionality.

* Update test to assert full tool description in test_component_to_tool.py

* Add unit tests for verifying required inputs of various components

- Added tests to ensure that required inputs for outputs are present in the inputs of `ChatInput`, `ChatOutput`, `SequentialTaskComponent`, `ToolCallingAgentComponent`, and `OpenAIModelComponent`.
- Included helper functions to check if required inputs are in inputs and to assert that all outputs have different required inputs.

* Add RequiredInputsVisitor to identify required inputs in AST nodes

- Introduced RequiredInputsVisitor class to traverse AST nodes and collect required inputs.
- The visitor checks for 'self' attributes matching the provided inputs and adds them to the required_inputs set.

* Refactor required inputs extraction using `RequiredInputsVisitor`

* Add feature flags configuration for toolkit output in settings

* Add toolkit output handling based on feature flag in custom component utils

* Add method to append 'component_as_tool' output in custom component

* Add unit test for toolkit output feature flag in custom component

* Add utility functions for lazy loading and instantiating input types in langflow

- Introduced `get_InputTypesMap` for lazy loading of `InputTypesMap`.
- Added `instantiate_input` function to create instances of input types dynamically.
- Included type checking and error handling for invalid input types.

* Refactor input instantiation logic and update imports

- Removed `instantiate_input` function from `inputs.py` and moved it to `utils.py`.
- Updated imports in `base.py` to reflect the new location of `instantiate_input`.
- Added missing import for `Callable` in `base.py`.

* Refactor import statement to use `instantiate_input` from `langflow.inputs.utils` in test_inputs.py

* Add TOOL_OUTPUT_NAME constant to tools module

* Add type checking and TOOL_OUTPUT_NAME filter in ComponentToolkit

- Introduced `TYPE_CHECKING` for type hints to avoid circular imports.
- Added `TOOL_OUTPUT_NAME` constant to filter specific outputs in `ComponentToolkit`.
- Updated type annotations to use forward references.

* Refactor component toolkit import to avoid circular dependency and use constant for tool output name

* Refactor `ComponentToolkit` class to remove inheritance from `BaseToolkit` and add an initializer for `component`

* Add unit test for ComponentToolkit in test_component_to_tool

- Added `test_component_to_tool_has_no_component_as_tool` to verify that `ComponentToolkit` correctly initializes with a `ChatInput` component and returns the expected tools.

* Refactor toolkit output handling to `custom_component` module

* fix: mypy errors union-attr and arg-type

* Add 'OTHER' field type to schema in langflow/io/schema.py

* Add tool name formatting to ComponentToolkit to ensure valid characters

* Refactor toolkit output handling and add type hint for `to_toolkit` method

* Add `is_interface_component` attribute to vertex types and update import order

* Add tests for ToolCallingAgentComponent and ChatOutput with API key handling

- Updated `test_component_tool` to reflect new description format.
- Added `test_component_tool_with_api_key` to test `ToolCallingAgentComponent` with `ChatOutput` and OpenAI API key.
- Enabled `add_toolkit_output` feature flag for testing.

* Refactor `_find_matching_output_method` to accept `input_name` parameter for more precise input-output matching

* Replace ValueError with warning in build_description function

* use chat_output component directly in set

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 3/3)

* Refactor: Reorder method calls in `__init__` for logical consistency

Moved `set_class_code` method call to ensure output types and required inputs are set before class code initialization.

* Update _format_tool_name to allow '.' in tool names

* Refactor `_format_tool_name` to remove non-alphanumeric characters

* Update test assertions for component tool name and output mapping

* Handle case where 'required_inputs' is empty in 'component_tool.py'

* Refactor import statements for better readability in `base.py`

* [autofix.ci] apply automated fixes

* Add noqa comment to suppress import warning and re-add Any import in base.py

---------

Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

src/backend/tests/conftest.py

@@ -367,7 +367,8 @@ def added_flow_with_prompt_and_history(client, json_flow_with_prompt_and_history
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture
@@ -379,7 +380,8 @@ def added_flow_chat_input(client, json_chat_input, logged_in_headers):
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture
@@ -391,7 +393,8 @@ def added_flow_two_outputs(client, json_two_outputs, logged_in_headers):
     assert response.status_code == 201
     assert response.json()["name"] == flow.name
     assert response.json()["data"] == flow.data
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture
@@ -403,7 +406,8 @@ def added_vector_store(client, json_vector_store, logged_in_headers):
     assert response.status_code == 201
     assert response.json()["name"] == vector_store.name
     assert response.json()["data"] == vector_store.data
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture
@@ -429,7 +433,8 @@ def flow_component(client: TestClient, logged_in_headers):
     flow = FlowCreate(**graph_dict)
     response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture
@@ -460,7 +465,8 @@ def get_simple_api_test(client, logged_in_headers, json_simple_api_test):
     flow = FlowCreate(name="Simple API Test", data=data, description="Simple API Test")
     response = client.post("api/v1/flows/", json=flow.model_dump(), headers=logged_in_headers)
     assert response.status_code == 201
-    return response.json()
+    yield response.json()
+    client.delete(f"api/v1/flows/{response.json()['id']}", headers=logged_in_headers)
 
 
 @pytest.fixture(name="starter_project")

src/backend/tests/unit/base/tools/test_component_tool.py

@@ -1,62 +0,0 @@
-import pytest
-
-from langflow.base.tools.component_tool import ComponentTool
-from langflow.components.inputs.ChatInput import ChatInput
-
-
-@pytest.fixture
-def client():
-    pass
-
-
-def test_component_tool():
-    chat_input = ChatInput()
-    component_tool = ComponentTool(component=chat_input)
-    assert component_tool.name == "ChatInput"
-    assert component_tool.description == chat_input.description
-    assert component_tool.args == {
-        "input_value": {
-            "default": "",
-            "description": "Message to be passed as input.",
-            "title": "Input Value",
-            "type": "string",
-        },
-        "should_store_message": {
-            "default": True,
-            "description": "Store the message in the history.",
-            "title": "Should Store Message",
-            "type": "boolean",
-        },
-        "sender": {
-            "default": "User",
-            "description": "Type of sender.",
-            "enum": ["Machine", "User"],
-            "title": "Sender",
-            "type": "string",
-        },
-        "sender_name": {
-            "default": "User",
-            "description": "Name of the sender.",
-            "title": "Sender Name",
-            "type": "string",
-        },
-        "session_id": {
-            "default": "",
-            "description": "The session ID of the chat. If empty, the current session ID parameter will be used.",
-            "title": "Session Id",
-            "type": "string",
-        },
-        "files": {
-            "default": "",
-            "description": "Files to be sent with the message.",
-            "items": {"type": "string"},
-            "title": "Files",
-            "type": "array",
-        },
-    }
-    assert component_tool.component == chat_input
-
-    result = component_tool.invoke(input=dict(input_value="test"))
-    assert isinstance(result, dict)
-    assert hasattr(result["message"], "get_text")
-    assert result["message"].get_text() == "test"

src/backend/tests/unit/base/tools/test_component_toolkit.py

@@ -0,0 +1,103 @@
+import os
+
+import pytest
+
+from langflow.base.tools.component_tool import ComponentToolkit
+from langflow.components.agents.ToolCallingAgent import ToolCallingAgentComponent
+from langflow.components.inputs.ChatInput import ChatInput
+from langflow.components.models.OpenAIModel import OpenAIModelComponent
+from langflow.components.outputs import ChatOutput
+from langflow.graph.graph.base import Graph
+from langflow.schema.message import Message
+from langflow.services.settings.feature_flags import FEATURE_FLAGS
+
+
+@pytest.fixture
+def client():
+    pass
+
+
+@pytest.fixture
+def add_toolkit_output():
+    FEATURE_FLAGS.add_toolkit_output = True
+    yield
+    FEATURE_FLAGS.add_toolkit_output = False
+
+
+def test_component_tool():
+    chat_input = ChatInput()
+    component_toolkit = ComponentToolkit(component=chat_input)
+    component_tool = component_toolkit.get_tools()[0]
+    assert component_tool.name == "ChatInput-message_response"
+    terms = [
+        "message_response",
+        "files",
+        "input_value",
+        "sender",
+        "sender_name",
+        "session_id",
+        "should_store_message",
+    ]
+    assert all(term in component_tool.description for term in terms)
+    assert component_tool.args == {
+        "input_value": {
+            "default": "",
+            "description": "Message to be passed as input.",
+            "title": "Input Value",
+            "type": "string",
+        },
+        "should_store_message": {
+            "default": True,
+            "description": "Store the message in the history.",
+            "title": "Should Store Message",
+            "type": "boolean",
+        },
+        "sender": {
+            "default": "User",
+            "description": "Type of sender.",
+            "enum": ["Machine", "User"],
+            "title": "Sender",
+            "type": "string",
+        },
+        "sender_name": {
+            "default": "User",
+            "description": "Name of the sender.",
+            "title": "Sender Name",
+            "type": "string",
+        },
+        "session_id": {
+            "default": "",
+            "description": "The session ID of the chat. If empty, the current session ID parameter will be used.",
+            "title": "Session Id",
+            "type": "string",
+        },
+        "files": {
+            "default": "",
+            "description": "Files to be sent with the message.",
+            "items": {"type": "string"},
+            "title": "Files",
+            "type": "array",
+        },
+    }
+    assert component_toolkit.component == chat_input
+
+    result = component_tool.invoke(input=dict(input_value="test"))
+    assert isinstance(result, Message)
+    assert result.get_text() == "test"
+
+
+@pytest.mark.api_key_required
+def test_component_tool_with_api_key(client, add_toolkit_output):
+    chat_output = ChatOutput()
+    openai_llm = OpenAIModelComponent()
+    openai_llm.set(api_key=os.environ["OPENAI_API_KEY"])
+    tool_calling_agent = ToolCallingAgentComponent()
+    tool_calling_agent.set(
+        llm=openai_llm.build_model, tools=[chat_output], input_value="Which tools are available? Please tell its name."
+    )
+
+    g = Graph(start=tool_calling_agent, end=tool_calling_agent)
+    assert g is not None
+    results = list(g.start())
+    assert len(results) == 4
+    assert "message_response" in tool_calling_agent._outputs_map["response"].value.get_text()

src/backend/tests/unit/custom/component/test_component_to_tool.py

@@ -1,3 +1,5 @@
+from collections.abc import Callable
+
 import pytest
 
 from langflow.components.inputs.ChatInput import ChatInput
@@ -8,9 +10,29 @@ def client():
     pass
 
 
-def test_component_to_tool():
+def test_component_to_toolkit():
     chat_input = ChatInput()
-    tool = chat_input.to_tool()
-    assert tool.name == "ChatInput"
-    assert tool.description == "Get chat inputs from the Playground."
-    assert tool.component._id == chat_input._id
+    tools = chat_input.to_toolkit()
+    assert len(tools) == 1
+    tool = tools[0]
+
+    assert tool.name == "ChatInput-message_response"
+    terms = [
+        "message_response",
+        "files",
+        "input_value",
+        "sender",
+        "sender_name",
+        "session_id",
+        "should_store_message",
+    ]
+    assert all(term in tool.description for term in terms)
+
+    assert isinstance(tool.func, Callable)
+    assert tool.args_schema is not None
+
+
+def test_component_to_tool_has_no_component_as_tool():
+    chat_input = ChatInput()
+    tools = chat_input.to_toolkit()
+    assert len(tools) == 1

src/backend/tests/unit/custom/custom_component/test_component.py

@@ -1,9 +1,12 @@
 import pytest
 
 from langflow.components.agents.CrewAIAgent import CrewAIAgentComponent
+from langflow.components.agents.ToolCallingAgent import ToolCallingAgentComponent
 from langflow.components.helpers.SequentialTask import SequentialTaskComponent
 from langflow.components.inputs.ChatInput import ChatInput
+from langflow.components.models.OpenAIModel import OpenAIModelComponent
 from langflow.components.outputs import ChatOutput
+from langflow.template.field.base import Output
 
 
 @pytest.fixture
@@ -24,3 +27,45 @@ def test_set_component():
     task.set(agent=crewai_agent)
     assert task._edges[0]["source"] == crewai_agent._id
     assert crewai_agent in task._components
+
+
+def _output_required_inputs_are_in_inputs(output: Output, inputs: list[str]):
+    return all(input_type in inputs for input_type in output.required_inputs)
+
+
+def _assert_all_outputs_have_different_required_inputs(outputs: list[Output]):
+    required_inputs = [tuple(output.required_inputs) for output in outputs]
+    assert len(required_inputs) == len(set(required_inputs)), "All outputs must have different required inputs"
+    return True
+
+
+def test_set_required_inputs():
+    chatinput = ChatInput()
+
+    assert all(_output_required_inputs_are_in_inputs(output, chatinput._inputs) for output in chatinput.outputs)
+    assert _assert_all_outputs_have_different_required_inputs(chatinput.outputs)
+
+
+def test_set_required_inputs_various_components():
+    chatinput = ChatInput()
+    chatoutput = ChatOutput()
+    task = SequentialTaskComponent()
+    tool_calling_agent = ToolCallingAgentComponent()
+    openai_component = OpenAIModelComponent()
+
+    assert all(_output_required_inputs_are_in_inputs(output, chatinput._inputs) for output in chatinput.outputs)
+    assert all(_output_required_inputs_are_in_inputs(output, chatoutput._inputs) for output in chatoutput.outputs)
+    assert all(_output_required_inputs_are_in_inputs(output, task._inputs) for output in task.outputs)
+    assert all(
+        _output_required_inputs_are_in_inputs(output, tool_calling_agent._inputs)
+        for output in tool_calling_agent.outputs
+    )
+    assert all(
+        _output_required_inputs_are_in_inputs(output, openai_component._inputs) for output in openai_component.outputs
+    )
+
+    assert _assert_all_outputs_have_different_required_inputs(chatinput.outputs)
+    assert _assert_all_outputs_have_different_required_inputs(chatoutput.outputs)
+    assert _assert_all_outputs_have_different_required_inputs(task.outputs)
+    assert _assert_all_outputs_have_different_required_inputs(tool_calling_agent.outputs)
+    assert _assert_all_outputs_have_different_required_inputs(openai_component.outputs)

src/backend/tests/unit/graph/graph/test_base.py

@@ -145,13 +145,6 @@ def test_graph_functional_start_end():
     assert results[-1] == Finish()
 
 
-def test_graph_set_with_invalid_component():
-    chat_input = ChatInput(_id="chat_input")
-    chat_output = ChatOutput(input_value="test", _id="chat_output")
-    with pytest.raises(ValueError, match="There are multiple outputs"):
-        chat_output.set(sender_name=chat_input)
-
-
 @pytest.mark.skip(reason="Temporarily disabled")
 def test_graph_set_with_valid_component():
     tool = YfinanceToolComponent()

src/backend/tests/unit/inputs/test_inputs.py

@@ -21,8 +21,8 @@ from langflow.inputs.inputs import (
     SecretStrInput,
     StrInput,
     TableInput,
-    instantiate_input,
 )
+from langflow.inputs.utils import instantiate_input
 from langflow.schema.message import Message
 
 

src/backend/tests/unit/io/test_io_schema.py

@@ -218,18 +218,6 @@ class TestCreateInputSchema:
         assert field_info.description == ""
 
     # Handling invalid field types
-    def test_invalid_field_types_handling(self):
-        from langflow.inputs.inputs import StrInput
-        from langflow.io.schema import create_input_schema
-
-        class InvalidFieldType:
-            pass
-
-        input_instance = StrInput(name="test_field")
-        input_instance.field_type = InvalidFieldType()
-
-        with pytest.raises(KeyError):
-            create_input_schema([input_instance])
 
     # Handling input types with None as default value
     def test_none_default_value_handling(self):

src/backend/tests/unit/test_custom_component.py

@@ -11,6 +11,7 @@ from langflow.custom.code_parser.code_parser import CodeParser, CodeSyntaxError
 from langflow.custom.custom_component.base_component import BaseComponent, ComponentCodeNullError
 from langflow.custom.utils import build_custom_component_template
 from langflow.services.database.models.flow import FlowCreate
+from langflow.services.settings.feature_flags import FEATURE_FLAGS
 
 
 @pytest.fixture
@@ -543,6 +544,14 @@ def test_custom_component_multiple_outputs(code_component_with_multiple_outputs,
     assert frontnd_node_dict["outputs"][0]["types"] == ["Text"]
 
 
+def test_feature_flags_add_toolkit_output(active_user, code_component_with_multiple_outputs):
+    frontnd_node_dict, _ = build_custom_component_template(code_component_with_multiple_outputs, active_user.id)
+    len_outputs = len(frontnd_node_dict["outputs"])
+    FEATURE_FLAGS.add_toolkit_output = True
+    frontnd_node_dict, _ = build_custom_component_template(code_component_with_multiple_outputs, active_user.id)
+    assert len(frontnd_node_dict["outputs"]) == len_outputs + 1
+
+
 def test_custom_component_subclass_from_lctoolcomponent():
     # Import LCToolComponent and create a subclass
     code = dedent("""