From 7d5ccb324a790eabfdc595fe4faf95e00982a915 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 24 Jul 2024 21:33:05 -0300 Subject: [PATCH] fix: validate params assignment in custom_component_update endpoint (#2931) * feat(RunFlow.py): update input and output definitions for RunFlowComponent * refactor: update params assignment in custom_component_update endpoint Simplify the params assignment in the custom_component_update endpoint by using a dictionary comprehension. This improves code readability and reduces the number of lines. * feat(custom_component.py, flow.py): add support for specifying output type in run_flow method to filter outputs based on type --- src/backend/base/langflow/api/v1/endpoints.py | 4 +- .../langflow/components/prototypes/RunFlow.py | 50 +++++++++++-------- .../custom_component/custom_component.py | 10 +++- src/backend/base/langflow/helpers/flow.py | 16 +++++- 4 files changed, 55 insertions(+), 25 deletions(-) diff --git a/src/backend/base/langflow/api/v1/endpoints.py b/src/backend/base/langflow/api/v1/endpoints.py index 84f047d77..fadd1d7b3 100644 --- a/src/backend/base/langflow/api/v1/endpoints.py +++ b/src/backend/base/langflow/api/v1/endpoints.py @@ -590,7 +590,9 @@ async def custom_component_update( ) if hasattr(cc_instance, "set_attributes"): template = code_request.get_template() - params = {key: value_dict["value"] for key, value_dict in template.items() if isinstance(value_dict, dict)} + params = { + key: value_dict.get("value") for key, value_dict in template.items() if isinstance(value_dict, dict) + } load_from_db_fields = [ field_name for field_name, field_dict in template.items() diff --git a/src/backend/base/langflow/components/prototypes/RunFlow.py b/src/backend/base/langflow/components/prototypes/RunFlow.py index 2114cdc55..285b0ff81 100644 --- a/src/backend/base/langflow/components/prototypes/RunFlow.py +++ b/src/backend/base/langflow/components/prototypes/RunFlow.py @@ -1,13 +1,13 @@ from typing import Any, List, Optional from langflow.base.flow_processing.utils import build_data_from_run_outputs -from langflow.custom import CustomComponent -from langflow.field_typing import NestedDict, Text +from langflow.custom import Component from langflow.graph.schema import RunOutputs +from langflow.io import DropdownInput, MessageTextInput, NestedDictInput, Output from langflow.schema import Data, dotdict -class RunFlowComponent(CustomComponent): +class RunFlowComponent(Component): display_name = "Run Flow" description = "A component to run a flow." name = "RunFlow" @@ -23,27 +23,33 @@ class RunFlowComponent(CustomComponent): return build_config - def build_config(self): - return { - "input_value": { - "display_name": "Input Value", - "multiline": True, - }, - "flow_name": { - "display_name": "Flow Name", - "info": "The name of the flow to run.", - "options": [], - "refresh_button": True, - }, - "tweaks": { - "display_name": "Tweaks", - "info": "Tweaks to apply to the flow.", - }, - } + inputs = [ + MessageTextInput( + name="input_value", + display_name="Input Value", + info="The input value to be processed by the flow.", + ), + DropdownInput( + name="flow_name", + display_name="Flow Name", + info="The name of the flow to run.", + options=[], + refresh_button=True, + ), + NestedDictInput( + name="tweaks", + display_name="Tweaks", + info="Tweaks to apply to the flow.", + ), + ] - async def build(self, input_value: Text, flow_name: str, tweaks: NestedDict) -> List[Data]: + outputs = [ + Output(display_name="Run Outputs", name="run_outputs", method="generate_results"), + ] + + async def generate_results(self) -> List[Data]: results: List[Optional[RunOutputs]] = await self.run_flow( - inputs={"input_value": input_value}, flow_name=flow_name, tweaks=tweaks + inputs={"input_value": self.input_value}, flow_name=self.flow_name, tweaks=self.tweaks ) if isinstance(results, list): data = [] diff --git a/src/backend/base/langflow/custom/custom_component/custom_component.py b/src/backend/base/langflow/custom/custom_component/custom_component.py index 5847757c1..753328aa0 100644 --- a/src/backend/base/langflow/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/custom/custom_component/custom_component.py @@ -480,9 +480,17 @@ class CustomComponent(BaseComponent): inputs: Optional[Union[dict, List[dict]]] = None, flow_id: Optional[str] = None, flow_name: Optional[str] = None, + output_type: Optional[str] = None, tweaks: Optional[dict] = None, ) -> Any: - return await run_flow(inputs=inputs, flow_id=flow_id, flow_name=flow_name, tweaks=tweaks, user_id=self._user_id) + return await run_flow( + inputs=inputs, + output_type=output_type, + flow_id=flow_id, + flow_name=flow_name, + tweaks=tweaks, + user_id=self._user_id, + ) def list_flows(self) -> List[Data]: if not self._user_id: diff --git a/src/backend/base/langflow/helpers/flow.py b/src/backend/base/langflow/helpers/flow.py index 9507ff7be..de0a14d89 100644 --- a/src/backend/base/langflow/helpers/flow.py +++ b/src/backend/base/langflow/helpers/flow.py @@ -71,6 +71,7 @@ async def run_flow( tweaks: Optional[dict] = None, flow_id: Optional[str] = None, flow_name: Optional[str] = None, + output_type: Optional[str] = None, user_id: Optional[str] = None, ) -> List[RunOutputs]: if user_id is None: @@ -89,10 +90,23 @@ async def run_flow( inputs_components.append(input_dict.get("components", [])) types.append(input_dict.get("type", "chat")) + outputs = [ + vertex.id + for vertex in graph.vertices + if output_type == "debug" + or ( + vertex.is_output and (output_type == "any" or output_type in vertex.id.lower()) # type: ignore + ) + ] + fallback_to_env_vars = get_settings_service().settings.fallback_to_env_var return await graph.arun( - inputs_list, inputs_components=inputs_components, types=types, fallback_to_env_vars=fallback_to_env_vars + inputs_list, + outputs=outputs, + inputs_components=inputs_components, + types=types, + fallback_to_env_vars=fallback_to_env_vars, )