Skip to content

Commit

Permalink
refactor: Tool Output Management: Add Component-Specific Variable for…
Browse files Browse the repository at this point in the history
… 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`.
  • Loading branch information
edwinjosechittilappilly authored Oct 23, 2024
1 parent 4075212 commit 5811857
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,9 @@ def __init__(self, **kwargs) -> None:
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"):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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'."""
Expand Down
19 changes: 17 additions & 2 deletions src/backend/tests/unit/test_custom_component_with_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

0 comments on commit 5811857

Please sign in to comment.