diff --git a/pyproject.toml b/pyproject.toml index a24d3e4f4..aabbe24a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langflow" -version = "0.5.0a4" +version = "0.5.0a5" description = "A Python package with a built-in web application" authors = ["Logspace "] maintainers = [ diff --git a/src/backend/langflow/interface/custom/constants.py b/src/backend/langflow/interface/custom/constants.py index f5911198e..58ecef637 100644 --- a/src/backend/langflow/interface/custom/constants.py +++ b/src/backend/langflow/interface/custom/constants.py @@ -1,4 +1,4 @@ -from langchain import PromptTemplate +from langchain.prompts import PromptTemplate from langchain.chains.base import Chain from langchain.document_loaders.base import BaseLoader from langchain.schema.embeddings import Embeddings @@ -45,7 +45,7 @@ DEFAULT_CUSTOM_COMPONENT_CODE = """from langflow import CustomComponent from langchain.llms.base import BaseLLM from langchain.chains import LLMChain -from langchain import PromptTemplate +from langchain.prompts import PromptTemplate from langchain.schema import Document import requests diff --git a/src/backend/langflow/interface/importing/utils.py b/src/backend/langflow/interface/importing/utils.py index 000eed205..44d72df42 100644 --- a/src/backend/langflow/interface/importing/utils.py +++ b/src/backend/langflow/interface/importing/utils.py @@ -3,7 +3,7 @@ import importlib from typing import Any, Type -from langchain import PromptTemplate +from langchain.prompts import PromptTemplate from langchain.agents import Agent from langchain.base_language import BaseLanguageModel from langchain.chains.base import Chain diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py index 3a388f694..8b9dde7bd 100644 --- a/src/backend/langflow/services/database/manager.py +++ b/src/backend/langflow/services/database/manager.py @@ -58,6 +58,27 @@ class DatabaseService(Service): with Session(self.engine) as session: yield session + def migrate_flows_if_auto_login(self): + # if auto_login is enabled, we need to migrate the flows + # to the default superuser if they don't have a user id + # associated with them + settings_service = get_settings_service() + if settings_service.auth_settings.AUTO_LOGIN: + with Session(self.engine) as session: + flows = ( + session.query(models.Flow) + .filter(models.Flow.user_id == None) # noqa + .all() + ) + if flows: + logger.debug("Migrating flows to default superuser") + username = settings_service.auth_settings.SUPERUSER + user = get_user_by_username(session, username) + for flow in flows: + flow.user_id = user.id + session.commit() + logger.debug("Flows migrated successfully") + def check_schema_health(self) -> bool: inspector = inspect(self.engine) diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 20c23f8c0..b7d93184f 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -7,7 +7,7 @@ from langflow.services.settings.constants import ( DEFAULT_SUPERUSER_PASSWORD, ) from sqlmodel import Session -from .getters import get_session, get_settings_service +from .getters import get_db_service, get_session, get_settings_service from loguru import logger @@ -54,8 +54,15 @@ def get_or_create_super_user(session: Session, username, password, is_default): logger.debug("Creating default superuser.") else: logger.debug("Creating superuser.") - - return create_super_user(username, password, db=session) + try: + return create_super_user(username, password, db=session) + except Exception as exc: + if "UNIQUE constraint failed: user.username" in str(exc): + # This is to deal with workers running this + # at startup and trying to create the superuser + # at the same time. + logger.debug("Superuser already exists.") + return None def setup_superuser(settings_service, session: Session): @@ -193,3 +200,8 @@ def initialize_services(): setup_superuser( service_manager.get(ServiceType.SETTINGS_SERVICE), next(get_session()) ) + try: + get_db_service().migrate_flows_if_auto_login() + except Exception as exc: + logger.error(f"Error migrating flows: {exc}") + raise RuntimeError("Error migrating flows") from exc diff --git a/tests/test_custom_component.py b/tests/test_custom_component.py index 1695cfd38..47c9cbfb2 100644 --- a/tests/test_custom_component.py +++ b/tests/test_custom_component.py @@ -21,7 +21,7 @@ from langflow.interface.custom.custom_component import CustomComponent from langchain.llms.base import BaseLLM from langchain.chains import LLMChain -from langchain import PromptTemplate +from langchain.prompts import PromptTemplate from langchain.schema import Document import requests