refactor: Tool Output Management: Add Component-Specific Variable for Conditional Tool Output (#4259)

Implement conditional tool output handling

- Updated `component.py` to check for `FEATURE_FLAGS.add_toolkit_output` and call `_append_tool_output` if the component has the `add_tool_output` attribute set to True.
- Ensured that the tool output behavior is correctly managed based on feature flags and component attributes in `custom_component.py`.
This commit is contained in:
Edwin Jose 2024-10-23 18:44:00 -04:00 committed by GitHub
commit 5811857ca2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 21 additions and 4 deletions

View file

@ -91,9 +91,9 @@ class Component(CustomComponent):
self.__inputs = inputs
self.__config = config
self._reset_all_output_values()
if FEATURE_FLAGS.add_toolkit_output and hasattr(self, "_append_tool_output"):
self._append_tool_output()
super().__init__(**config)
if (FEATURE_FLAGS.add_toolkit_output) and hasattr(self, "_append_tool_output") and self.add_tool_output:
self._append_tool_output()
if hasattr(self, "_trace_type"):
self.trace_type = self._trace_type
if not hasattr(self, "trace_type"):

View file

@ -59,6 +59,8 @@ class CustomComponent(BaseComponent):
is_input: bool | None = None
"""The input state of the component. Defaults to None.
If True, the component must have a field named 'input_value'."""
add_tool_output: bool | None = False
"""Indicates whether the component will be treated as a tool. Defaults to False."""
is_output: bool | None = None
"""The output state of the component. Defaults to None.
If True, the component must have a field named 'input_value'."""

View file

@ -14,6 +14,16 @@ def code_component_with_multiple_outputs():
return Component(_code=code)
@pytest.fixture
def code_component_with_multiple_outputs_with_add_tool_output():
code = Path("src/backend/tests/data/component_multiple_outputs.py").read_text(encoding="utf-8")
code = code.replace(
"class MultipleOutputsComponent(Component):",
"class MultipleOutputsComponent(Component):\n add_tool_output = True",
)
return Component(_code=code)
@pytest.fixture
def component(
client, # noqa: ARG001
@ -43,9 +53,14 @@ def test_list_flows_return_type(component):
assert isinstance(flows, list)
def test_feature_flags_add_toolkit_output(active_user, code_component_with_multiple_outputs):
def test_feature_flags_add_toolkit_output(
active_user, code_component_with_multiple_outputs, code_component_with_multiple_outputs_with_add_tool_output
):
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)
code_component_with_multiple_outputs_with_add_tool_output.add_tool_output = True
frontnd_node_dict, _ = build_custom_component_template(
code_component_with_multiple_outputs_with_add_tool_output, active_user.id
)
assert len(frontnd_node_dict["outputs"]) == len_outputs + 1