From 461238aa0e8cf6b8ae21ab5e63bae55951b15ea2 Mon Sep 17 00:00:00 2001
From: Edwin Jose <edwin.jose@datastax.com>
Date: Thu, 3 Oct 2024 11:43:31 -0400
Subject: [PATCH] fix: Add `session_id` Parameter to `run_flow_from_json`
 Function (#3989)

This pull request fixes the `run_flow_from_json` function by adding  `session_id` parameter. This update ensures that session management is properly handled during flow execution.

Changes include:
- Added `session_id` as an optional input parameter.
- Updated function documentation to reflect the inclusion of the new parameter.
- Added unit test for run_flow_from_json input parameters

This fix enhances the function's usability by allowing for better tracking of individual sessions.
---
 src/backend/base/langflow/load/load.py        |  3 +++
 .../base/langflow/processing/process.py       |  4 ++-
 src/backend/tests/unit/base/load/test_load.py | 26 +++++++++++++++++++
 .../modals/apiModal/utils/get-python-code.tsx |  1 +
 4 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 src/backend/tests/unit/base/load/test_load.py

diff --git a/src/backend/base/langflow/load/load.py b/src/backend/base/langflow/load/load.py
index d3b669c4b..f7294f136 100644
--- a/src/backend/base/langflow/load/load.py
+++ b/src/backend/base/langflow/load/load.py
@@ -72,6 +72,7 @@ def load_flow_from_json(
 def run_flow_from_json(
     flow: Path | str | dict,
     input_value: str,
+    session_id: str | None = None,
     tweaks: dict | None = None,
     input_type: str = "chat",
     output_type: str = "chat",
@@ -89,6 +90,7 @@ def run_flow_from_json(
     Args:
         flow (Union[Path, str, dict]): The path to the JSON file or the JSON dictionary representing the flow.
         input_value (str): The input value to be processed by the flow.
+        session_id (str | None, optional): The session ID to be used for the flow. Defaults to None.
         tweaks (Optional[dict], optional): Optional tweaks to be applied to the flow. Defaults to None.
         input_type (str, optional): The type of the input value. Defaults to "chat".
         output_type (str, optional): The type of the output value. Defaults to "chat".
@@ -125,6 +127,7 @@ def run_flow_from_json(
     )
     return run_graph(
         graph=graph,
+        session_id=session_id,
         input_value=input_value,
         input_type=input_type,
         output_type=output_type,
diff --git a/src/backend/base/langflow/processing/process.py b/src/backend/base/langflow/processing/process.py
index 48079ae6c..e3117df07 100644
--- a/src/backend/base/langflow/processing/process.py
+++ b/src/backend/base/langflow/processing/process.py
@@ -62,6 +62,7 @@ def run_graph(
     input_value: str,
     input_type: str,
     output_type: str,
+    session_id: str | None = None,
     fallback_to_env_vars: bool = False,
     output_component: str | None = None,
 ) -> list[RunOutputs]:
@@ -73,6 +74,7 @@ def run_graph(
         input_value (str): The input value to be passed to the graph.
         input_type (str): The type of the input value.
         output_type (str): The type of the desired output.
+        session_id (str | None, optional): The session ID to be used for the flow. Defaults to None.
         output_component (Optional[str], optional): The specific output component to retrieve. Defaults to None.
 
     Returns:
@@ -105,7 +107,7 @@ def run_graph(
         types,
         outputs or [],
         stream=False,
-        session_id="",
+        session_id=session_id,
         fallback_to_env_vars=fallback_to_env_vars,
     )
 
diff --git a/src/backend/tests/unit/base/load/test_load.py b/src/backend/tests/unit/base/load/test_load.py
new file mode 100644
index 000000000..cb58fae5d
--- /dev/null
+++ b/src/backend/tests/unit/base/load/test_load.py
@@ -0,0 +1,26 @@
+from langflow.load import run_flow_from_json
+
+
+def test_run_flow_from_json_params():
+    # Define the expected parameters
+    expected_params = {
+        "flow",
+        "input_value",
+        "session_id",
+        "tweaks",
+        "input_type",
+        "output_type",
+        "output_component",
+        "log_level",
+        "log_file",
+        "env_file",
+        "cache",
+        "disable_logs",
+        "fallback_to_env_vars",
+    }
+
+    # Check if the function accepts all expected parameters
+    params = run_flow_from_json.__code__.co_varnames[: run_flow_from_json.__code__.co_argcount]
+    assert expected_params.issubset(params), "Not all expected parameters are present in run_flow_from_json"
+
+    # TODO: Add tests by loading a flow and running it need to text with fake llm and check if it returns the correct output
diff --git a/src/frontend/src/modals/apiModal/utils/get-python-code.tsx b/src/frontend/src/modals/apiModal/utils/get-python-code.tsx
index 39fd1c154..407b7832d 100644
--- a/src/frontend/src/modals/apiModal/utils/get-python-code.tsx
+++ b/src/frontend/src/modals/apiModal/utils/get-python-code.tsx
@@ -22,6 +22,7 @@ TWEAKS = ${tweaksString}
 
 result = run_flow_from_json(flow="${flowName}.json",
                             input_value="message",
+                            session_id="", # provide a session id if you want to use session state
                             fallback_to_env_vars=True, # False by default
                             tweaks=TWEAKS)`;
 }