From 50f9f72f5f67b0483e827b13438a94326eb50b21 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 17 Sep 2024 15:35:22 -0300 Subject: [PATCH] feat: Add run_id parameter to run_flow function (#3834) * Add run_id parameter to set run ID in graph within load_and_run_flow function * Add run_id parameter to run_flow call in FlowTool for tracking runs * Add run_id parameter to CustomComponent's run_flow method * fix: mypy error arg-type --------- Co-authored-by: italojohnny --- src/backend/base/langflow/base/tools/flow_tool.py | 2 ++ .../base/langflow/custom/custom_component/custom_component.py | 1 + src/backend/base/langflow/helpers/flow.py | 3 +++ 3 files changed, 6 insertions(+) diff --git a/src/backend/base/langflow/base/tools/flow_tool.py b/src/backend/base/langflow/base/tools/flow_tool.py index 03936ac5839b..298228b9c4e6 100644 --- a/src/backend/base/langflow/base/tools/flow_tool.py +++ b/src/backend/base/langflow/base/tools/flow_tool.py @@ -95,10 +95,12 @@ async def _arun( ) -> str: """Use the tool asynchronously.""" tweaks = self.build_tweaks_dict(args, kwargs) + run_id = self.graph.run_id if self.graph else None run_outputs = await run_flow( tweaks={key: {"input_value": value} for key, value in tweaks.items()}, flow_id=self.flow_id, user_id=self.user_id, + run_id=run_id, ) if not run_outputs: return "No output" 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 a6c5836b1cbe..781e24a0a2f1 100644 --- a/src/backend/base/langflow/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/custom/custom_component/custom_component.py @@ -484,6 +484,7 @@ async def run_flow( flow_name=flow_name, tweaks=tweaks, user_id=str(self._user_id), + run_id=self.graph.run_id, ) def list_flows(self) -> list[Data]: diff --git a/src/backend/base/langflow/helpers/flow.py b/src/backend/base/langflow/helpers/flow.py index b753620c5de5..89530504d264 100644 --- a/src/backend/base/langflow/helpers/flow.py +++ b/src/backend/base/langflow/helpers/flow.py @@ -74,10 +74,13 @@ async def run_flow( flow_name: Optional[str] = None, output_type: Optional[str] = "chat", user_id: Optional[str] = None, + run_id: Optional[str] = None, ) -> List[RunOutputs]: if user_id is None: raise ValueError("Session is invalid") graph = await load_flow(user_id, flow_id, flow_name, tweaks) + if run_id: + graph.set_run_id(UUID(run_id)) if inputs is None: inputs = []