diff --git a/src/backend/base/langflow/api/v1/flows.py b/src/backend/base/langflow/api/v1/flows.py index 7a999dea9..da9fe1430 100644 --- a/src/backend/base/langflow/api/v1/flows.py +++ b/src/backend/base/langflow/api/v1/flows.py @@ -25,7 +25,7 @@ from langflow.api.utils import ( validate_is_component, ) from langflow.api.v1.schemas import FlowListCreate -from langflow.initial_setup.setup import STARTER_FOLDER_NAME +from langflow.initial_setup.constants import STARTER_FOLDER_NAME from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate from langflow.services.database.models.flow.model import FlowHeader from langflow.services.database.models.flow.utils import get_webhook_component_in_flow diff --git a/src/backend/base/langflow/api/v1/folders.py b/src/backend/base/langflow/api/v1/folders.py index 4b9a6b279..8707209c9 100644 --- a/src/backend/base/langflow/api/v1/folders.py +++ b/src/backend/base/langflow/api/v1/folders.py @@ -19,7 +19,7 @@ from langflow.api.v1.flows import create_flows from langflow.api.v1.schemas import FlowListCreate from langflow.helpers.flow import generate_unique_flow_name from langflow.helpers.folders import generate_unique_folder_name -from langflow.initial_setup.setup import STARTER_FOLDER_NAME +from langflow.initial_setup.constants import STARTER_FOLDER_NAME from langflow.services.database.models.flow.model import Flow, FlowCreate, FlowRead from langflow.services.database.models.folder.constants import DEFAULT_FOLDER_NAME from langflow.services.database.models.folder.model import ( diff --git a/src/backend/base/langflow/initial_setup/constants.py b/src/backend/base/langflow/initial_setup/constants.py new file mode 100644 index 000000000..b84cb90d1 --- /dev/null +++ b/src/backend/base/langflow/initial_setup/constants.py @@ -0,0 +1,2 @@ +STARTER_FOLDER_NAME = "Starter Projects" +STARTER_FOLDER_DESCRIPTION = "Starter projects to help you get started in Langflow." diff --git a/src/backend/base/langflow/initial_setup/setup.py b/src/backend/base/langflow/initial_setup/setup.py index c9de941a6..b75088d8d 100644 --- a/src/backend/base/langflow/initial_setup/setup.py +++ b/src/backend/base/langflow/initial_setup/setup.py @@ -14,11 +14,8 @@ from loguru import logger from sqlalchemy.exc import NoResultFound from sqlmodel import select -from langflow.base.constants import ( - FIELD_FORMAT_ATTRIBUTES, - NODE_FORMAT_ATTRIBUTES, - ORJSON_OPTIONS, -) +from langflow.base.constants import FIELD_FORMAT_ATTRIBUTES, NODE_FORMAT_ATTRIBUTES, ORJSON_OPTIONS +from langflow.initial_setup.constants import STARTER_FOLDER_DESCRIPTION, STARTER_FOLDER_NAME 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 @@ -37,9 +34,6 @@ from langflow.services.deps import ( from langflow.template.field.prompt import DEFAULT_PROMPT_INTUT_TYPES from langflow.utils.util import escape_json_dump -STARTER_FOLDER_NAME = "Starter Projects" -STARTER_FOLDER_DESCRIPTION = "Starter projects to help you get started in Langflow." - # In the folder ./starter_projects we have a few JSON files that represent # starter projects. We want to load these into the database so that users # can use them as a starting point for their own projects. diff --git a/src/backend/base/langflow/services/database/service.py b/src/backend/base/langflow/services/database/service.py index 643835aa7..ed582ec01 100644 --- a/src/backend/base/langflow/services/database/service.py +++ b/src/backend/base/langflow/services/database/service.py @@ -21,13 +21,11 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlmodel import Session, SQLModel, create_engine, select, text from sqlmodel.ext.asyncio.session import AsyncSession +from langflow.initial_setup.constants import STARTER_FOLDER_NAME from langflow.services.base import Service from langflow.services.database import models from langflow.services.database.models.user.crud import get_user_by_username -from langflow.services.database.utils import ( - Result, - TableResults, -) +from langflow.services.database.utils import Result, TableResults from langflow.services.deps import get_settings_service from langflow.services.utils import teardown_superuser @@ -147,44 +145,77 @@ class DatabaseService(Service): yield session async def assign_orphaned_flows_to_superuser(self) -> None: - """Assign flows without a user ID to the default superuser if auto_login is enabled.""" - # Get the settings service to check if auto_login is enabled + """Assign orphaned flows to the default superuser when auto login is enabled.""" settings_service = get_settings_service() - if settings_service.auth_settings.AUTO_LOGIN: - async with self.with_async_session() as session: - # Select flows where user_id is NULL (orphaned flows) - stmt = select(models.Flow).where(models.Flow.user_id == None) # noqa: E711 - orphaned_flows = (await session.exec(stmt)).all() - if orphaned_flows: - logger.debug("Assigning orphaned flows to the default superuser") + if not settings_service.auth_settings.AUTO_LOGIN: + return - # Get the default superuser - superuser_username = settings_service.auth_settings.SUPERUSER - superuser = await get_user_by_username(session, superuser_username) + async with self.with_async_session() as session: + # Fetch orphaned flows + stmt = ( + select(models.Flow) + .join(models.Folder) + .where( + models.Flow.user_id == None, # noqa: E711 + models.Folder.name != STARTER_FOLDER_NAME, + ) + ) + orphaned_flows = (await session.exec(stmt)).all() - if not superuser: - logger.error("Default superuser not found") - msg = "Default superuser not found" - raise RuntimeError(msg) + if not orphaned_flows: + return - stmt = select(models.Flow.name).where(models.Flow.user_id == superuser.id) - result = await session.exec(stmt) - superuser_flows_names = result.all() - # Assign each orphaned flow to the superuser - for flow in orphaned_flows: - flow.user_id = superuser.id - if flow.name in superuser_flows_names: - name_match = re.search(r"\((\d+)\)$", flow.name) - if not name_match: - flow.name = f"{flow.name} (1)" - else: - num = int(name_match.group(1)) + 1 - flow.name = re.sub(r"\(\d+\)$", f"({num})", flow.name) + logger.debug("Assigning orphaned flows to the default superuser") - # Commit the changes to the database - await session.commit() - logger.debug("Successfully assigned orphaned flows to the default superuser") + # Retrieve superuser + superuser_username = settings_service.auth_settings.SUPERUSER + superuser = await get_user_by_username(session, superuser_username) + + if not superuser: + error_message = "Default superuser not found" + logger.error(error_message) + raise RuntimeError(error_message) + + # Get existing flow names for the superuser + existing_names: set[str] = set( + (await session.exec(select(models.Flow.name).where(models.Flow.user_id == superuser.id))).all() + ) + + # Process orphaned flows + for flow in orphaned_flows: + flow.user_id = superuser.id + flow.name = self._generate_unique_flow_name(flow.name, existing_names) + existing_names.add(flow.name) + + # Commit changes + await session.commit() + logger.debug("Successfully assigned orphaned flows to the default superuser") + + def _generate_unique_flow_name(self, original_name: str, existing_names: set[str]) -> str: + """Generate a unique flow name by adding or incrementing a suffix.""" + if original_name not in existing_names: + return original_name + + match = re.search(r"^(.*) \((\d+)\)$", original_name) + if match: + base_name, current_number = match.groups() + new_name = f"{base_name} ({int(current_number) + 1})" + else: + new_name = f"{original_name} (1)" + + # Ensure unique name by incrementing suffix + while new_name in existing_names: + match = re.match(r"^(.*) \((\d+)\)$", new_name) + if match is not None: + base_name, current_number = match.groups() + else: + error_message = "Invalid format: match is None" + raise ValueError(error_message) + + new_name = f"{base_name} ({int(current_number) + 1})" + + return new_name def check_schema_health(self) -> bool: inspector = inspect(self.engine) diff --git a/src/backend/tests/conftest.py b/src/backend/tests/conftest.py index 4f33ab1f2..7605c349c 100644 --- a/src/backend/tests/conftest.py +++ b/src/backend/tests/conftest.py @@ -17,7 +17,7 @@ from dotenv import load_dotenv from fastapi.testclient import TestClient from httpx import ASGITransport, AsyncClient from langflow.graph import Graph -from langflow.initial_setup.setup import STARTER_FOLDER_NAME +from langflow.initial_setup.constants import STARTER_FOLDER_NAME from langflow.services.auth.utils import get_password_hash from langflow.services.database.models.api_key.model import ApiKey from langflow.services.database.models.flow.model import Flow, FlowCreate diff --git a/src/backend/tests/unit/test_initial_setup.py b/src/backend/tests/unit/test_initial_setup.py index 348102551..3ec0361de 100644 --- a/src/backend/tests/unit/test_initial_setup.py +++ b/src/backend/tests/unit/test_initial_setup.py @@ -3,11 +3,9 @@ from datetime import datetime from pathlib import Path import pytest -from langflow.custom.directory_reader.utils import ( - abuild_custom_component_list_from_path, -) +from langflow.custom.directory_reader.utils import abuild_custom_component_list_from_path +from langflow.initial_setup.constants import STARTER_FOLDER_NAME from langflow.initial_setup.setup import ( - STARTER_FOLDER_NAME, get_project_data, load_starter_projects, update_projects_components_with_latest_component_versions,