From 7fe478e4b190f6e2f2e8593d5a1f21bc59f4fbfc Mon Sep 17 00:00:00 2001 From: Jordan Frazier <122494242+jordanrfrazier@users.noreply.github.com> Date: Fri, 4 Oct 2024 10:19:08 -0700 Subject: [PATCH] fix: set folder id on flows imported on startup (#4018) * set folder id on flows imported on startup --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- .../base/langflow/initial_setup/setup.py | 28 ++++++++++++++++++- .../services/database/models/folder/utils.py | 13 ++++++++- src/backend/tests/unit/test_database.py | 2 ++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/initial_setup/setup.py b/src/backend/base/langflow/initial_setup/setup.py index 9f440b027..9ec6d5c28 100644 --- a/src/backend/base/langflow/initial_setup/setup.py +++ b/src/backend/base/langflow/initial_setup/setup.py @@ -20,7 +20,10 @@ from langflow.graph.graph.base import Graph from langflow.services.auth.utils import create_super_user from langflow.services.database.models.flow.model import Flow, FlowCreate from langflow.services.database.models.folder.model import Folder, FolderCreate -from langflow.services.database.models.folder.utils import create_default_folder_if_it_doesnt_exist +from langflow.services.database.models.folder.utils import ( + create_default_folder_if_it_doesnt_exist, + get_default_folder_id, +) from langflow.services.database.models.user.crud import get_user_by_username from langflow.services.deps import get_settings_service, get_storage_service, get_variable_service, session_scope from langflow.template.field.prompt import DEFAULT_PROMPT_INTUT_TYPES @@ -497,6 +500,12 @@ def _is_valid_uuid(val): def load_flows_from_directory(): + """ + On langflow startup, this loads all flows from the directory specified in the settings. + + All flows are uploaded into the default folder for the superuser. + Note that this feature currently only works if AUTO_LOGIN is enabled in the settings. + """ settings_service = get_settings_service() flows_path = settings_service.settings.load_flows_path if not flows_path: @@ -522,6 +531,7 @@ def load_flows_from_directory(): existing = find_existing_flow(session, flow_id, flow_endpoint_name) if existing: + logger.debug(f"Found existing flow: {existing.name}") logger.info(f"Updating existing flow: {flow_id} with endpoint name {flow_endpoint_name}") for key, value in flow.items(): if hasattr(existing, key): @@ -529,10 +539,23 @@ def load_flows_from_directory(): setattr(existing, key, value) existing.updated_at = datetime.now(tz=timezone.utc).astimezone() existing.user_id = user_id + + # Generally, folder_id should not be None, but we must check this due to the previous + # behavior where flows could be added and folder_id was None, orphaning + # them within Langflow. + if existing.folder_id is None: + folder_id = get_default_folder_id(session, user_id) + existing.folder_id = folder_id + session.add(existing) else: logger.info(f"Creating new flow: {flow_id} with endpoint name {flow_endpoint_name}") + + # Current behavior loads all new flows into default folder + folder_id = get_default_folder_id(session, user_id) + flow["user_id"] = user_id + flow["folder_id"] = folder_id flow = Flow.model_validate(flow, from_attributes=True) flow.updated_at = datetime.now(tz=timezone.utc).astimezone() session.add(flow) @@ -540,11 +563,14 @@ def load_flows_from_directory(): def find_existing_flow(session, flow_id, flow_endpoint_name): if flow_endpoint_name: + logger.debug(f"flow_endpoint_name: {flow_endpoint_name}") stmt = select(Flow).where(Flow.endpoint_name == flow_endpoint_name) if existing := session.exec(stmt).first(): + logger.debug(f"Found existing flow by endpoint name: {existing.name}") return existing stmt = select(Flow).where(Flow.id == flow_id) if existing := session.exec(stmt).first(): + logger.debug(f"Found existing flow by id: {flow_id}") return existing return None diff --git a/src/backend/base/langflow/services/database/models/folder/utils.py b/src/backend/base/langflow/services/database/models/folder/utils.py index e7403fb11..4ffdb4fdc 100644 --- a/src/backend/base/langflow/services/database/models/folder/utils.py +++ b/src/backend/base/langflow/services/database/models/folder/utils.py @@ -11,7 +11,11 @@ from .model import Folder def create_default_folder_if_it_doesnt_exist(session: Session, user_id: UUID): folder = session.exec(select(Folder).where(Folder.user_id == user_id)).first() if not folder: - folder = Folder(name=DEFAULT_FOLDER_NAME, user_id=user_id, description=DEFAULT_FOLDER_DESCRIPTION) + folder = Folder( + name=DEFAULT_FOLDER_NAME, + user_id=user_id, + description=DEFAULT_FOLDER_DESCRIPTION, + ) session.add(folder) session.commit() session.refresh(folder) @@ -27,3 +31,10 @@ def create_default_folder_if_it_doesnt_exist(session: Session, user_id: UUID): ) session.commit() return folder + + +def get_default_folder_id(session: Session, user_id: UUID): + folder = session.exec(select(Folder).where(Folder.name == DEFAULT_FOLDER_NAME, Folder.user_id == user_id)).first() + if not folder: + folder = create_default_folder_if_it_doesnt_exist(session, user_id) + return folder.id diff --git a/src/backend/tests/unit/test_database.py b/src/backend/tests/unit/test_database.py index 9705df15e..1b9668e50 100644 --- a/src/backend/tests/unit/test_database.py +++ b/src/backend/tests/unit/test_database.py @@ -423,11 +423,13 @@ async def test_load_flows(client: TestClient, load_flows_dir): response = await client.get("api/v1/flows/c54f9130-f2fa-4a3e-b22a-3856d946351b") assert response.status_code == 200 assert response.json()["name"] == "BasicExample" + assert response.json()["folder_id"] is not None # re-run to ensure updates work well load_flows_from_directory() response = await client.get("api/v1/flows/c54f9130-f2fa-4a3e-b22a-3856d946351b") assert response.status_code == 200 assert response.json()["name"] == "BasicExample" + assert response.json()["folder_id"] is not None def test_sqlite_pragmas():