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 <italojohnnydosanjos@gmail.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-09-17 15:35:22 -03:00 • committed by GitHub
commit 50f9f72f5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 6 additions and 0 deletions

View file

@ -95,10 +95,12 @@ class FlowTool(BaseTool):
) -> str: ) -> str:
"""Use the tool asynchronously.""" """Use the tool asynchronously."""
tweaks = self.build_tweaks_dict(args, kwargs) tweaks = self.build_tweaks_dict(args, kwargs)
run_id = self.graph.run_id if self.graph else None
run_outputs = await run_flow( run_outputs = await run_flow(
tweaks={key: {"input_value": value} for key, value in tweaks.items()}, tweaks={key: {"input_value": value} for key, value in tweaks.items()},
flow_id=self.flow_id, flow_id=self.flow_id,
user_id=self.user_id, user_id=self.user_id,
run_id=run_id,
) )
if not run_outputs: if not run_outputs:
return "No output" return "No output"

View file

@ -484,6 +484,7 @@ class CustomComponent(BaseComponent):
flow_name=flow_name, flow_name=flow_name,
tweaks=tweaks, tweaks=tweaks,
user_id=str(self._user_id), user_id=str(self._user_id),
run_id=self.graph.run_id,
) )
def list_flows(self) -> list[Data]: def list_flows(self) -> list[Data]:

View file

@ -74,10 +74,13 @@ async def run_flow(
flow_name: Optional[str] = None, flow_name: Optional[str] = None,
output_type: Optional[str] = "chat", output_type: Optional[str] = "chat",
user_id: Optional[str] = None, user_id: Optional[str] = None,
run_id: Optional[str] = None,
) -> List[RunOutputs]: ) -> List[RunOutputs]:
if user_id is None: if user_id is None:
raise ValueError("Session is invalid") raise ValueError("Session is invalid")
graph = await load_flow(user_id, flow_id, flow_name, tweaks) 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: if inputs is None:
inputs = [] inputs = []