vibesurf 0.2.40__py3-none-any.whl

vibe_surf/langflow/services/tracing/utils.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+from vibe_surf.langflow.schema.data import Data
+
+
+def convert_to_langchain_type(value):
+    from vibe_surf.langflow.schema.message import Message
+
+    if isinstance(value, dict):
+        value = {key: convert_to_langchain_type(val) for key, val in value.items()}
+    elif isinstance(value, list):
+        value = [convert_to_langchain_type(v) for v in value]
+    elif isinstance(value, Message):
+        if "prompt" in value:
+            value = value.load_lc_prompt()
+        elif value.sender:
+            value = value.to_lc_message()
+        else:
+            value = value.to_lc_document()
+    elif isinstance(value, Data):
+        value = value.to_lc_document() if "text" in value.data else value.data
+    return value
+
+
+def convert_to_langchain_types(io_dict: dict[str, Any]):
+    converted = {}
+    for key, value in io_dict.items():
+        converted[key] = convert_to_langchain_type(value)
+    return converted

vibe_surf/langflow/services/utils.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete
+from sqlalchemy import exc as sqlalchemy_exc
+from sqlmodel import col, select
+
+from vibe_surf.langflow.logging.logger import logger
+from vibe_surf.langflow.services.auth.utils import create_super_user, verify_password
+from vibe_surf.langflow.services.cache.base import ExternalAsyncBaseCacheService
+from vibe_surf.langflow.services.cache.factory import CacheServiceFactory
+from vibe_surf.langflow.services.database.models.transactions.model import TransactionTable
+from vibe_surf.langflow.services.database.models.vertex_builds.model import VertexBuildTable
+from vibe_surf.langflow.services.database.utils import initialize_database
+from vibe_surf.langflow.services.schema import ServiceType
+from vibe_surf.langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD
+
+from .deps import get_db_service, get_service, get_settings_service
+
+if TYPE_CHECKING:
+    from sqlmodel.ext.asyncio.session import AsyncSession
+
+    from vibe_surf.langflow.services.settings.manager import SettingsService
+
+
+async def get_or_create_super_user(session: AsyncSession, username, password, is_default):
+    from vibe_surf.langflow.services.database.models.user.model import User
+
+    stmt = select(User).where(User.username == username)
+    result = await session.exec(stmt)
+    user = result.first()
+
+    if user and user.is_superuser:
+        return None  # Superuser already exists
+
+    if user and is_default:
+        if user.is_superuser:
+            if verify_password(password, user.password):
+                return None
+            # Superuser exists but password is incorrect
+            # which means that the user has changed the
+            # base superuser credentials.
+            # This means that the user has already created
+            # a superuser and changed the password in the UI
+            # so we don't need to do anything.
+            await logger.adebug(
+                "Superuser exists but password is incorrect. "
+                "This means that the user has changed the "
+                "base superuser credentials."
+            )
+            return None
+        logger.debug("User with superuser credentials exists but is not a superuser.")
+        return None
+
+    if user:
+        if verify_password(password, user.password):
+            msg = "User with superuser credentials exists but is not a superuser."
+            raise ValueError(msg)
+        msg = "Incorrect superuser credentials"
+        raise ValueError(msg)
+
+    if is_default:
+        logger.debug("Creating default superuser.")
+    else:
+        logger.debug("Creating superuser.")
+    return await create_super_user(username, password, db=session)
+
+
+async def setup_superuser(settings_service: SettingsService, session: AsyncSession) -> None:
+    if settings_service.auth_settings.AUTO_LOGIN:
+        await logger.adebug("AUTO_LOGIN is set to True. Creating default superuser.")
+        username = DEFAULT_SUPERUSER
+        password = DEFAULT_SUPERUSER_PASSWORD.get_secret_value()
+    else:
+        # Remove the default superuser if it exists
+        await teardown_superuser(settings_service, session)
+        # If AUTO_LOGIN is disabled, attempt to use configured credentials
+        # or fall back to default credentials if none are provided.
+        username = settings_service.auth_settings.SUPERUSER or DEFAULT_SUPERUSER
+        password = (settings_service.auth_settings.SUPERUSER_PASSWORD or DEFAULT_SUPERUSER_PASSWORD).get_secret_value()
+
+    if not username or not password:
+        msg = "Username and password must be set"
+        raise ValueError(msg)
+
+    is_default = (username == DEFAULT_SUPERUSER) and (password == DEFAULT_SUPERUSER_PASSWORD.get_secret_value())
+
+    try:
+        user = await get_or_create_super_user(
+            session=session, username=username, password=password, is_default=is_default
+        )
+        if user is not None:
+            await logger.adebug("Superuser created successfully.")
+    except Exception as exc:
+        logger.exception(exc)
+        msg = "Could not create superuser. Please create a superuser manually."
+        raise RuntimeError(msg) from exc
+    finally:
+        # Scrub credentials from in-memory settings after setup
+        settings_service.auth_settings.reset_credentials()
+
+
+async def teardown_superuser(settings_service, session: AsyncSession) -> None:
+    """Teardown the superuser."""
+    # If AUTO_LOGIN is True, we will remove the default superuser
+    # from the database.
+
+    if not settings_service.auth_settings.AUTO_LOGIN:
+        try:
+            await logger.adebug("AUTO_LOGIN is set to False. Removing default superuser if exists.")
+            username = DEFAULT_SUPERUSER
+            from vibe_surf.langflow.services.database.models.user.model import User
+
+            stmt = select(User).where(User.username == username)
+            result = await session.exec(stmt)
+            user = result.first()
+            # Check if super was ever logged in, if not delete it
+            # if it has logged in, it means the user is using it to login
+            if user and user.is_superuser is True and not user.last_login_at:
+                await session.delete(user)
+                await session.commit()
+                await logger.adebug("Default superuser removed successfully.")
+
+        except Exception as exc:
+            logger.exception(exc)
+            await session.rollback()
+            msg = "Could not remove default superuser."
+            raise RuntimeError(msg) from exc
+
+
+async def teardown_services() -> None:
+    """Teardown all the services."""
+    async with get_db_service().with_session() as session:
+        await teardown_superuser(get_settings_service(), session)
+
+    from vibe_surf.langflow.services.manager import service_manager
+
+    await service_manager.teardown()
+
+
+def initialize_settings_service() -> None:
+    """Initialize the settings manager."""
+    from vibe_surf.langflow.services.settings import factory as settings_factory
+
+    get_service(ServiceType.SETTINGS_SERVICE, settings_factory.SettingsServiceFactory())
+
+
+def initialize_session_service() -> None:
+    """Initialize the session manager."""
+    from vibe_surf.langflow.services.cache import factory as cache_factory
+    from vibe_surf.langflow.services.session import factory as session_service_factory
+
+    initialize_settings_service()
+
+    get_service(
+        ServiceType.CACHE_SERVICE,
+        cache_factory.CacheServiceFactory(),
+    )
+
+    get_service(
+        ServiceType.SESSION_SERVICE,
+        session_service_factory.SessionServiceFactory(),
+    )
+
+
+async def clean_transactions(settings_service: SettingsService, session: AsyncSession) -> None:
+    """Clean up old transactions from the database.
+
+    This function deletes transactions that exceed the maximum number to keep (configured in settings).
+    It orders transactions by timestamp descending and removes the oldest ones beyond the limit.
+
+    Args:
+        settings_service: The settings service containing configuration like max_transactions_to_keep
+        session: The database session to use for the deletion
+    """
+    try:
+        # Delete transactions using bulk delete
+        delete_stmt = delete(TransactionTable).where(
+            col(TransactionTable.id).in_(
+                select(TransactionTable.id)
+                .order_by(col(TransactionTable.timestamp).desc())
+                .offset(settings_service.settings.max_transactions_to_keep)
+            )
+        )
+
+        await session.exec(delete_stmt)
+        await session.commit()
+        logger.debug("Successfully cleaned up old transactions")
+    except (sqlalchemy_exc.SQLAlchemyError, asyncio.TimeoutError) as exc:
+        logger.error(f"Error cleaning up transactions: {exc!s}")
+        await session.rollback()
+        # Don't re-raise since this is a cleanup task
+
+
+async def clean_vertex_builds(settings_service: SettingsService, session: AsyncSession) -> None:
+    """Clean up old vertex builds from the database.
+
+    This function deletes vertex builds that exceed the maximum number to keep (configured in settings).
+    It orders vertex builds by timestamp descending and removes the oldest ones beyond the limit.
+
+    Args:
+        settings_service: The settings service containing configuration like max_vertex_builds_to_keep
+        session: The database session to use for the deletion
+    """
+    try:
+        # Delete vertex builds using bulk delete
+        delete_stmt = delete(VertexBuildTable).where(
+            col(VertexBuildTable.id).in_(
+                select(VertexBuildTable.id)
+                .order_by(col(VertexBuildTable.timestamp).desc())
+                .offset(settings_service.settings.max_vertex_builds_to_keep)
+            )
+        )
+
+        await session.exec(delete_stmt)
+        await session.commit()
+        logger.debug("Successfully cleaned up old vertex builds")
+    except (sqlalchemy_exc.SQLAlchemyError, asyncio.TimeoutError) as exc:
+        logger.error(f"Error cleaning up vertex builds: {exc!s}")
+        await session.rollback()
+        # Don't re-raise since this is a cleanup task
+
+
+async def initialize_services(*, fix_migration: bool = False) -> None:
+    """Initialize all the services needed."""
+    cache_service = get_service(ServiceType.CACHE_SERVICE, default=CacheServiceFactory())
+    # Test external cache connection
+    if isinstance(cache_service, ExternalAsyncBaseCacheService) and not (await cache_service.is_connected()):
+        msg = "Cache service failed to connect to external database"
+        raise ConnectionError(msg)
+
+    # Setup the superuser
+    await initialize_database(fix_migration=fix_migration)
+    db_service = get_db_service()
+    await db_service.initialize_alembic_log_file()
+    async with db_service.with_session() as session:
+        settings_service = get_service(ServiceType.SETTINGS_SERVICE)
+        await setup_superuser(settings_service, session)
+    try:
+        await get_db_service().assign_orphaned_flows_to_superuser()
+    except sqlalchemy_exc.IntegrityError as exc:
+        await logger.awarning(f"Error assigning orphaned flows to the superuser: {exc!s}")
+    await clean_transactions(settings_service, session)
+    await clean_vertex_builds(settings_service, session)

vibe_surf/langflow/services/variable/base.py

@@ -0,0 +1,119 @@
+import abc
+from uuid import UUID
+
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from vibe_surf.langflow.services.base import Service
+from vibe_surf.langflow.services.database.models.variable.model import Variable, VariableRead
+
+
+class VariableService(Service):
+    """Abstract base class for a variable service."""
+
+    name = "variable_service"
+
+    @abc.abstractmethod
+    async def initialize_user_variables(self, user_id: UUID | str, session: AsyncSession) -> None:
+        """Initialize user variables.
+
+        Args:
+            user_id: The user ID.
+            session: The database session.
+        """
+
+    @abc.abstractmethod
+    async def get_variable(self, user_id: UUID | str, name: str, field: str, session: AsyncSession) -> str:
+        """Async get a variable value.
+
+        Args:
+            user_id: The user ID.
+            name: The name of the variable.
+            field: The field of the variable.
+            session: The database session.
+
+        Returns:
+            The value of the variable.
+        """
+
+    @abc.abstractmethod
+    async def list_variables(self, user_id: UUID | str, session: AsyncSession) -> list[str | None]:
+        """List all variables.
+
+        Args:
+            user_id: The user ID.
+            session: The database session.
+
+        Returns:
+            A list of variable names.
+        """
+
+    @abc.abstractmethod
+    async def update_variable(self, user_id: UUID | str, name: str, value: str, session: AsyncSession) -> Variable:
+        """Update a variable.
+
+        Args:
+            user_id: The user ID.
+            name: The name of the variable.
+            value: The value of the variable.
+            session: The database session.
+
+        Returns:
+            The updated variable.
+        """
+
+    @abc.abstractmethod
+    async def delete_variable(self, user_id: UUID | str, name: str, session: AsyncSession) -> None:
+        """Delete a variable.
+
+        Args:
+            user_id: The user ID.
+            name: The name of the variable.
+            session: The database session.
+
+        Returns:
+            The deleted variable.
+        """
+
+    @abc.abstractmethod
+    async def delete_variable_by_id(self, user_id: UUID | str, variable_id: UUID, session: AsyncSession) -> None:
+        """Delete a variable by ID.
+
+        Args:
+            user_id: The user ID.
+            variable_id: The ID of the variable.
+            session: The database session.
+        """
+
+    @abc.abstractmethod
+    async def create_variable(
+        self,
+        user_id: UUID | str,
+        name: str,
+        value: str,
+        *,
+        default_fields: list[str],
+        type_: str,
+        session: AsyncSession,
+    ) -> Variable:
+        """Create a variable.
+
+        Args:
+            user_id: The user ID.
+            name: The name of the variable.
+            value: The value of the variable.
+            default_fields: The default fields of the variable.
+            type_: The type of the variable.
+            session: The database session.
+
+        Returns:
+            The created variable.
+        """
+
+    @abc.abstractmethod
+    async def get_all(self, user_id: UUID | str, session: AsyncSession) -> list[VariableRead]:
+        """Get all variables.
+
+        Args:
+            user_id: The user ID.
+            session: The database session.
+        """

vibe_surf/langflow/services/variable/constants.py

@@ -0,0 +1,2 @@
+CREDENTIAL_TYPE = "Credential"
+GENERIC_TYPE = "Generic"

vibe_surf/langflow/services/variable/factory.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from vibe_surf.langflow.services.factory import ServiceFactory
+from vibe_surf.langflow.services.variable.service import DatabaseVariableService, VariableService
+
+if TYPE_CHECKING:
+    from vibe_surf.langflow.services.settings.service import SettingsService
+
+
+class VariableServiceFactory(ServiceFactory):
+    def __init__(self) -> None:
+        super().__init__(VariableService)
+
+    @override
+    def create(self, settings_service: SettingsService):
+        # here you would have logic to create and configure a VariableService
+        # based on the settings_service
+
+        if settings_service.settings.variable_store == "kubernetes":
+            # Keep it here to avoid import errors
+            from vibe_surf.langflow.services.variable.kubernetes import KubernetesSecretService
+
+            return KubernetesSecretService(settings_service)
+        return DatabaseVariableService(settings_service)

vibe_surf/langflow/services/variable/kubernetes.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+from vibe_surf.langflow.logging.logger import logger
+from vibe_surf.langflow.services.auth import utils as auth_utils
+from vibe_surf.langflow.services.base import Service
+from vibe_surf.langflow.services.database.models.variable.model import Variable, VariableCreate, VariableRead
+from vibe_surf.langflow.services.variable.base import VariableService
+from vibe_surf.langflow.services.variable.constants import CREDENTIAL_TYPE, GENERIC_TYPE
+from vibe_surf.langflow.services.variable.kubernetes_secrets import KubernetesSecretManager, encode_user_id
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from sqlmodel import Session
+    from sqlmodel.ext.asyncio.session import AsyncSession
+
+    from vibe_surf.langflow.services.settings.service import SettingsService
+
+
+class KubernetesSecretService(VariableService, Service):
+    def __init__(self, settings_service: SettingsService):
+        self.settings_service = settings_service
+        # TODO: settings_service to set kubernetes namespace
+        self.kubernetes_secrets = KubernetesSecretManager()
+
+    @override
+    async def initialize_user_variables(self, user_id: UUID | str, session: AsyncSession) -> None:
+        # Check for environment variables that should be stored in the database
+        should_or_should_not = "Should" if self.settings_service.settings.store_environment_variables else "Should not"
+        await logger.ainfo(f"{should_or_should_not} store environment variables in the kubernetes.")
+        if self.settings_service.settings.store_environment_variables:
+            variables = {}
+            for var in self.settings_service.settings.variables_to_get_from_environment:
+                if var in os.environ:
+                    await logger.adebug(f"Creating {var} variable from environment.")
+                    value = os.environ[var]
+                    if isinstance(value, str):
+                        value = value.strip()
+                    key = CREDENTIAL_TYPE + "_" + var
+                    variables[key] = str(value)
+
+            try:
+                secret_name = encode_user_id(user_id)
+                await asyncio.to_thread(
+                    self.kubernetes_secrets.create_secret,
+                    name=secret_name,
+                    data=variables,
+                )
+            except Exception:  # noqa: BLE001
+                logger.exception(f"Error creating {var} variable")
+
+        else:
+            logger.info("Skipping environment variable storage.")
+
+    # resolve_variable is a helper function that resolves the variable name to the actual key in the secret
+    def resolve_variable(
+        self,
+        secret_name: str,
+        user_id: UUID | str,
+        name: str,
+    ) -> tuple[str, str]:
+        variables = self.kubernetes_secrets.get_secret(name=secret_name)
+        if not variables:
+            msg = f"user_id {user_id} variable not found."
+            raise ValueError(msg)
+
+        if name in variables:
+            return name, variables[name]
+        credential_name = CREDENTIAL_TYPE + "_" + name
+        if credential_name in variables:
+            return credential_name, variables[credential_name]
+        msg = f"user_id {user_id} variable name {name} not found."
+        raise ValueError(msg)
+
+    @override
+    async def get_variable(self, user_id: UUID | str, name: str, field: str, session: AsyncSession) -> str:
+        secret_name = encode_user_id(user_id)
+        key, value = await asyncio.to_thread(self.resolve_variable, secret_name, user_id, name)
+        if key.startswith(CREDENTIAL_TYPE + "_") and field == "session_id":
+            msg = (
+                f"variable {name} of type 'Credential' cannot be used in a Session ID field "
+                "because its purpose is to prevent the exposure of values."
+            )
+            raise TypeError(msg)
+        return value
+
+    @override
+    async def list_variables(
+        self,
+        user_id: UUID | str,
+        session: Session,
+    ) -> list[str | None]:
+        variables = await asyncio.to_thread(self.kubernetes_secrets.get_secret, name=encode_user_id(user_id))
+        if not variables:
+            return []
+
+        names = []
+        for key in variables:
+            if key.startswith(CREDENTIAL_TYPE + "_"):
+                names.append(key[len(CREDENTIAL_TYPE) + 1 :])
+            else:
+                names.append(key)
+        return names
+
+    def _update_variable(
+        self,
+        user_id: UUID | str,
+        name: str,
+        value: str,
+    ):
+        secret_name = encode_user_id(user_id)
+        secret_key, _ = self.resolve_variable(secret_name, user_id, name)
+        return self.kubernetes_secrets.update_secret(name=secret_name, data={secret_key: value})
+
+    @override
+    async def update_variable(
+        self,
+        user_id: UUID | str,
+        name: str,
+        value: str,
+        session: AsyncSession,
+    ):
+        return await asyncio.to_thread(self._update_variable, user_id, name, value)
+
+    def _delete_variable(self, user_id: UUID | str, name: str) -> None:
+        secret_name = encode_user_id(user_id)
+        secret_key, _ = self.resolve_variable(secret_name, user_id, name)
+        self.kubernetes_secrets.delete_secret_key(name=secret_name, key=secret_key)
+
+    @override
+    async def delete_variable(self, user_id: UUID | str, name: str, session: AsyncSession) -> None:
+        await asyncio.to_thread(self._delete_variable, user_id, name)
+
+    @override
+    async def delete_variable_by_id(self, user_id: UUID | str, variable_id: UUID | str, session: AsyncSession) -> None:
+        await self.delete_variable(user_id, str(variable_id), session)
+
+    @override
+    async def create_variable(
+        self,
+        user_id: UUID | str,
+        name: str,
+        value: str,
+        *,
+        default_fields: list[str],
+        type_: str,
+        session: AsyncSession,
+    ) -> Variable:
+        secret_name = encode_user_id(user_id)
+        secret_key = name
+        if type_ == CREDENTIAL_TYPE:
+            secret_key = CREDENTIAL_TYPE + "_" + name
+        else:
+            type_ = GENERIC_TYPE
+
+        await asyncio.to_thread(
+            self.kubernetes_secrets.upsert_secret, secret_name=secret_name, data={secret_key: value}
+        )
+
+        variable_base = VariableCreate(
+            name=name,
+            type=type_,
+            value=auth_utils.encrypt_api_key(value, settings_service=self.settings_service),
+            default_fields=default_fields,
+        )
+        return Variable.model_validate(variable_base, from_attributes=True, update={"user_id": user_id})
+
+    @override
+    async def get_all(self, user_id: UUID | str, session: AsyncSession) -> list[VariableRead]:
+        secret_name = encode_user_id(user_id)
+        variables = await asyncio.to_thread(self.kubernetes_secrets.get_secret, name=secret_name)
+        if not variables:
+            return []
+
+        variables_read = []
+        for key, value in variables.items():
+            name = key
+            type_ = GENERIC_TYPE
+            if key.startswith(CREDENTIAL_TYPE + "_"):
+                name = key[len(CREDENTIAL_TYPE) + 1 :]
+                type_ = CREDENTIAL_TYPE
+
+            decrypted_value = None
+            if type_ == GENERIC_TYPE:
+                decrypted_value = value
+
+            variable_base = VariableCreate(
+                name=name,
+                type=type_,
+                value=auth_utils.encrypt_api_key(value, settings_service=self.settings_service),
+                default_fields=[],
+            )
+            variable = Variable.model_validate(variable_base, from_attributes=True, update={"user_id": user_id})
+            variable_read = VariableRead.model_validate(variable, from_attributes=True)
+            variable_read.value = decrypted_value
+            variables_read.append(variable_read)
+
+        return variables_read