diff --git a/src/backend/langflow/__main__.py b/src/backend/langflow/__main__.py index 515613c23..5eb92d133 100644 --- a/src/backend/langflow/__main__.py +++ b/src/backend/langflow/__main__.py @@ -313,7 +313,7 @@ def superuser( if create_super_user(db=session, username=username, password=password): # Verify that the superuser was created - from langflow.services.database.models.user.user import User + from langflow.services.database.models.user.model import User user: User = session.query(User).filter(User.username == username).first() if user is None or not user.is_superuser: diff --git a/src/backend/langflow/api/v1/api_key.py b/src/backend/langflow/api/v1/api_key.py index fcd9c934f..5c7c6ca0b 100644 --- a/src/backend/langflow/api/v1/api_key.py +++ b/src/backend/langflow/api/v1/api_key.py @@ -2,7 +2,7 @@ from uuid import UUID from fastapi import APIRouter, HTTPException, Depends from langflow.api.v1.schemas import ApiKeysResponse, ApiKeyCreateRequest from langflow.services.auth import utils as auth_utils -from langflow.services.database.models.api_key.api_key import ( +from langflow.services.database.models.api_key.model import ( ApiKeyCreate, UnmaskedApiKeyRead, ) @@ -13,7 +13,7 @@ from langflow.services.database.models.api_key.crud import ( create_api_key, delete_api_key, ) -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.deps import ( get_session, get_settings_service, diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py index 2d2ebd952..58224d9ef 100644 --- a/src/backend/langflow/api/v1/endpoints.py +++ b/src/backend/langflow/api/v1/endpoints.py @@ -18,7 +18,7 @@ from langflow.processing.process import process_graph_cached, process_tweaks from langflow.services.auth.utils import api_key_security, get_current_active_user from langflow.services.cache.utils import save_uploaded_file from langflow.services.database.models.flow import Flow -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.deps import ( get_session, get_session_service, diff --git a/src/backend/langflow/api/v1/flows.py b/src/backend/langflow/api/v1/flows.py index d35b2a8d5..939899553 100644 --- a/src/backend/langflow/api/v1/flows.py +++ b/src/backend/langflow/api/v1/flows.py @@ -11,7 +11,7 @@ from langflow.api.utils import remove_api_keys from langflow.api.v1.schemas import FlowListCreate, FlowListRead from langflow.services.auth.utils import get_current_active_user from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.deps import get_session, get_settings_service # build router diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py index 99a8b08cb..f1f002184 100644 --- a/src/backend/langflow/api/v1/schemas.py +++ b/src/backend/langflow/api/v1/schemas.py @@ -2,7 +2,7 @@ from enum import Enum from pathlib import Path from typing import Any, Dict, List, Optional, Union from uuid import UUID -from langflow.services.database.models.api_key.api_key import ApiKeyRead +from langflow.services.database.models.api_key.model import ApiKeyRead from langflow.services.database.models.flow import FlowCreate, FlowRead from langflow.services.database.models.user import UserRead from langflow.services.database.models.base import orjson_dumps diff --git a/src/backend/langflow/api/v1/store.py b/src/backend/langflow/api/v1/store.py index 273ea41e4..e3d6cee9d 100644 --- a/src/backend/langflow/api/v1/store.py +++ b/src/backend/langflow/api/v1/store.py @@ -4,7 +4,7 @@ from uuid import UUID from fastapi import APIRouter, Depends, HTTPException, Query from langflow.services.auth import utils as auth_utils -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.deps import get_settings_service, get_store_service from langflow.services.store.exceptions import CustomException from langflow.services.store.schema import ( diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index da0ee7396..912d1fbe8 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -8,10 +8,10 @@ from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from jose import JWTError, jwt from sqlmodel import Session -from langflow.services.database.models.api_key.api_key import ApiKey +from langflow.services.database.models.api_key.model import ApiKey from langflow.services.database.models.api_key.crud import check_key from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.deps import get_session, get_settings_service oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False) diff --git a/src/backend/langflow/services/database/models/api_key/__init__.py b/src/backend/langflow/services/database/models/api_key/__init__.py index fbb8265b9..001b0327e 100644 --- a/src/backend/langflow/services/database/models/api_key/__init__.py +++ b/src/backend/langflow/services/database/models/api_key/__init__.py @@ -1,3 +1,3 @@ -from .api_key import ApiKey, ApiKeyCreate, UnmaskedApiKeyRead, ApiKeyRead +from .model import ApiKey, ApiKeyCreate, UnmaskedApiKeyRead, ApiKeyRead __all__ = ["ApiKey", "ApiKeyCreate", "UnmaskedApiKeyRead", "ApiKeyRead"] diff --git a/src/backend/langflow/services/database/models/api_key/api_key.py b/src/backend/langflow/services/database/models/api_key/model.py similarity index 86% rename from src/backend/langflow/services/database/models/api_key/api_key.py rename to src/backend/langflow/services/database/models/api_key/model.py index 684027ee2..1d81a09c2 100644 --- a/src/backend/langflow/services/database/models/api_key/api_key.py +++ b/src/backend/langflow/services/database/models/api_key/model.py @@ -1,15 +1,15 @@ -from pydantic import validator -from sqlmodel import Field, Relationship -from uuid import UUID, uuid4 -from typing import Optional, TYPE_CHECKING from datetime import datetime -from langflow.services.database.models.base import SQLModelSerializable +from typing import TYPE_CHECKING, Optional +from uuid import UUID, uuid4 + +from pydantic import validator +from sqlmodel import Field, Relationship, SQLModel if TYPE_CHECKING: from langflow.services.database.models.user import User -class ApiKeyBase(SQLModelSerializable): +class ApiKeyBase(SQLModel): name: Optional[str] = Field(index=True) created_at: datetime = Field(default_factory=datetime.utcnow) last_used_at: Optional[datetime] = Field(default=None, nullable=True) diff --git a/src/backend/langflow/services/database/models/base.py b/src/backend/langflow/services/database/models/base.py index 9dff68c19..3a834d5a8 100644 --- a/src/backend/langflow/services/database/models/base.py +++ b/src/backend/langflow/services/database/models/base.py @@ -16,9 +16,3 @@ def orjson_dumps(v, *, default=None, sort_keys=False, indent_2=True): if default is None: return orjson.dumps(v, option=option).decode() return orjson.dumps(v, default=default, option=option).decode() - - -class SQLModelSerializable(SQLModel): - # TODO[pydantic]: The following keys were removed: `json_loads`, `json_dumps`. - # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information. - pass diff --git a/src/backend/langflow/services/database/models/component/__init__.py b/src/backend/langflow/services/database/models/component/__init__.py index c787c3e04..c5ad7d47f 100644 --- a/src/backend/langflow/services/database/models/component/__init__.py +++ b/src/backend/langflow/services/database/models/component/__init__.py @@ -1,3 +1,3 @@ -from .component import Component, ComponentModel +from .model import Component, ComponentModel __all__ = ["Component", "ComponentModel"] diff --git a/src/backend/langflow/services/database/models/component/component.py b/src/backend/langflow/services/database/models/component/model.py similarity index 85% rename from src/backend/langflow/services/database/models/component/component.py rename to src/backend/langflow/services/database/models/component/model.py index 5c4e6c13a..0a5afc440 100644 --- a/src/backend/langflow/services/database/models/component/component.py +++ b/src/backend/langflow/services/database/models/component/model.py @@ -1,11 +1,11 @@ -from langflow.services.database.models.base import SQLModelSerializable, SQLModel -from sqlmodel import Field -from typing import Optional -from datetime import datetime import uuid +from datetime import datetime +from typing import Optional + +from sqlmodel import Field, SQLModel -class Component(SQLModelSerializable, table=True): +class Component(SQLModel, table=True): id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True) frontend_node_id: uuid.UUID = Field(index=True) name: str = Field(index=True) diff --git a/src/backend/langflow/services/database/models/flow/__init__.py b/src/backend/langflow/services/database/models/flow/__init__.py index 7c7cc0172..3c861963c 100644 --- a/src/backend/langflow/services/database/models/flow/__init__.py +++ b/src/backend/langflow/services/database/models/flow/__init__.py @@ -1,3 +1,3 @@ -from .flow import Flow, FlowCreate, FlowRead, FlowUpdate +from .model import Flow, FlowCreate, FlowRead, FlowUpdate __all__ = ["Flow", "FlowCreate", "FlowRead", "FlowUpdate"] diff --git a/src/backend/langflow/services/database/models/flow/flow.py b/src/backend/langflow/services/database/models/flow/model.py similarity index 91% rename from src/backend/langflow/services/database/models/flow/flow.py rename to src/backend/langflow/services/database/models/flow/model.py index 5bd3802ba..51902a2b3 100644 --- a/src/backend/langflow/services/database/models/flow/flow.py +++ b/src/backend/langflow/services/database/models/flow/model.py @@ -5,15 +5,13 @@ from typing import TYPE_CHECKING, Dict, Optional from uuid import UUID, uuid4 from pydantic import field_serializer, field_validator -from sqlmodel import JSON, Column, Field, Relationship - -from langflow.services.database.models.base import SQLModelSerializable +from sqlmodel import JSON, Column, Field, Relationship, SQLModel if TYPE_CHECKING: from langflow.services.database.models.user import User -class FlowBase(SQLModelSerializable): +class FlowBase(SQLModel): name: str = Field(index=True) description: Optional[str] = Field(index=True, nullable=True, default=None) data: Optional[Dict] = Field(default=None, nullable=True) @@ -69,7 +67,7 @@ class FlowRead(FlowBase): user_id: UUID = Field() -class FlowUpdate(SQLModelSerializable): +class FlowUpdate(SQLModel): name: Optional[str] = None description: Optional[str] = None data: Optional[Dict] = None diff --git a/src/backend/langflow/services/database/models/user/__init__.py b/src/backend/langflow/services/database/models/user/__init__.py index da9170eb7..eecf6d0b1 100644 --- a/src/backend/langflow/services/database/models/user/__init__.py +++ b/src/backend/langflow/services/database/models/user/__init__.py @@ -1,4 +1,4 @@ -from .user import User, UserCreate, UserRead, UserUpdate +from .model import User, UserCreate, UserRead, UserUpdate __all__ = [ "User", diff --git a/src/backend/langflow/services/database/models/user/crud.py b/src/backend/langflow/services/database/models/user/crud.py index c7239bcee..75ce56038 100644 --- a/src/backend/langflow/services/database/models/user/crud.py +++ b/src/backend/langflow/services/database/models/user/crud.py @@ -2,7 +2,7 @@ from datetime import datetime, timezone from typing import Union from uuid import UUID from fastapi import Depends, HTTPException, status -from langflow.services.database.models.user.user import User, UserUpdate +from langflow.services.database.models.user.model import User, UserUpdate from langflow.services.deps import get_session from sqlalchemy.exc import IntegrityError from sqlmodel import Session diff --git a/src/backend/langflow/services/database/models/user/user.py b/src/backend/langflow/services/database/models/user/model.py similarity index 82% rename from src/backend/langflow/services/database/models/user/user.py rename to src/backend/langflow/services/database/models/user/model.py index 50525d025..8e83870f0 100644 --- a/src/backend/langflow/services/database/models/user/user.py +++ b/src/backend/langflow/services/database/models/user/model.py @@ -1,17 +1,18 @@ -from langflow.services.database.models.base import SQLModel, SQLModelSerializable +from datetime import datetime +from typing import TYPE_CHECKING, Optional +from uuid import UUID, uuid4 + from sqlmodel import Field, Relationship - -from datetime import datetime -from typing import Optional, TYPE_CHECKING -from uuid import UUID, uuid4 +from langflow.services.database.models.base import SQLModel if TYPE_CHECKING: from langflow.services.database.models.api_key import ApiKey + from langflow.services.database.models.credential import Credential from langflow.services.database.models.flow import Flow -class User(SQLModelSerializable, table=True): +class User(SQLModel, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) username: str = Field(index=True, unique=True) password: str = Field() @@ -27,6 +28,10 @@ class User(SQLModelSerializable, table=True): ) store_api_key: str = Field(default=None, nullable=True) flows: list["Flow"] = Relationship(back_populates="user") + credentials: list["Credential"] = Relationship( + back_populates="user", + sa_relationship_kwargs={"cascade": "delete"}, + ) class UserCreate(SQLModel): diff --git a/src/backend/langflow/services/deps.py b/src/backend/langflow/services/deps.py index 776c0b2c4..2cac77a17 100644 --- a/src/backend/langflow/services/deps.py +++ b/src/backend/langflow/services/deps.py @@ -12,6 +12,11 @@ if TYPE_CHECKING: from langflow.services.settings.service import SettingsService from langflow.services.store.service import StoreService from langflow.services.task.service import TaskService + from langflow.services.credentials.service import CredentialService + + +def get_credential_service() -> "CredentialService": + return service_manager.get(ServiceType.CREDENTIAL_SERVICE) def get_settings_service() -> "SettingsService": diff --git a/src/backend/langflow/services/schema.py b/src/backend/langflow/services/schema.py index b899923fe..e35df5dec 100644 --- a/src/backend/langflow/services/schema.py +++ b/src/backend/langflow/services/schema.py @@ -15,3 +15,4 @@ class ServiceType(str, Enum): SESSION_SERVICE = "session_service" TASK_SERVICE = "task_service" STORE_SERVICE = "store_service" + CREDENTIAL_SERVICE = "credential_service" diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 72807ee0c..31ae609fb 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -44,7 +44,7 @@ def get_factories_and_deps(): def get_or_create_super_user(session: Session, username, password, is_default): - from langflow.services.database.models.user.user import User + from langflow.services.database.models.user.model import User user = session.query(User).filter(User.username == username).first() @@ -127,7 +127,7 @@ def teardown_superuser(settings_service, session): try: logger.debug("AUTO_LOGIN is set to False. Removing default superuser if exists.") username = DEFAULT_SUPERUSER - from langflow.services.database.models.user.user import User + from langflow.services.database.models.user.model import User user = session.query(User).filter(User.username == username).first() if user and user.is_superuser: diff --git a/tests/conftest.py b/tests/conftest.py index 0acc3116f..5694c2a43 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,8 +12,8 @@ from fastapi.testclient import TestClient from httpx import AsyncClient from langflow.graph.graph.base import Graph from langflow.services.auth.utils import get_password_hash -from langflow.services.database.models.flow.flow import Flow, FlowCreate -from langflow.services.database.models.user.user import User, UserCreate +from langflow.services.database.models.flow.model import Flow, FlowCreate +from langflow.services.database.models.user.model import User, UserCreate from langflow.services.database.utils import session_getter from langflow.services.deps import get_db_service from sqlmodel import Session, SQLModel, create_engine @@ -260,7 +260,7 @@ def logged_in_headers(client, active_user): @pytest.fixture def flow(client, json_flow: str, active_user): - from langflow.services.database.models.flow.flow import FlowCreate + from langflow.services.database.models.flow.model import FlowCreate loaded_json = json.loads(json_flow) flow_data = FlowCreate( diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py index aba0a1a78..aacc009d0 100644 --- a/tests/test_endpoints.py +++ b/tests/test_endpoints.py @@ -2,7 +2,7 @@ from collections import namedtuple import uuid from langflow.processing.process import Result from langflow.services.auth.utils import get_password_hash -from langflow.services.database.models.api_key.api_key import ApiKey +from langflow.services.database.models.api_key.model import ApiKey from langflow.services.deps import get_settings_service from langflow.services.database.utils import session_getter from langflow.services.deps import get_db_service diff --git a/tests/test_setup_superuser.py b/tests/test_setup_superuser.py index 95a8deffb..d3033b728 100644 --- a/tests/test_setup_superuser.py +++ b/tests/test_setup_superuser.py @@ -1,6 +1,6 @@ from unittest.mock import MagicMock, patch -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD from langflow.services.utils import teardown_superuser diff --git a/tests/test_user.py b/tests/test_user.py index e93b594cc..f8dc47073 100644 --- a/tests/test_user.py +++ b/tests/test_user.py @@ -3,7 +3,7 @@ from datetime import datetime import pytest from langflow.services.auth.utils import create_super_user, get_password_hash from langflow.services.database.models.user import UserUpdate -from langflow.services.database.models.user.user import User +from langflow.services.database.models.user.model import User from langflow.services.database.utils import session_getter from langflow.services.deps import get_db_service, get_settings_service