refactor: use typed SQLAlchemy base model and fix type errors

- Replace declarative_base with DeclarativeBase for better type support
- Fix type annotations in Account model and TenantAccountRole
- Correct variable naming to avoid shadowing and improve clarity
- Fix incorrect property references and type handling
- Resolve type errors throughout the codebase

Signed-off-by: -LAN- <laipz8200@outlook.com>

api/controllers/console/auth/login.py

@@ -202,18 +202,18 @@ class EmailCodeLoginApi(Resource):
         except AccountRegisterError as are:
             raise AccountInFreezeError()
         if account:
-            tenant = TenantService.get_join_tenants(account)
+            tenants = TenantService.get_join_tenants(account)
-            if not tenant:
+            if not tenants:
                 workspaces = FeatureService.get_system_features().license.workspaces
                 if not workspaces.is_available():
                     raise WorkspacesLimitExceeded()
                 if not FeatureService.get_system_features().is_allow_create_workspace:
                     raise NotAllowedCreateWorkspace()
                 else:
-                    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+                    new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
-                    TenantService.create_tenant_member(tenant, account, role="owner")
+                    TenantService.create_tenant_member(new_tenant, account, role="owner")
-                    account.current_tenant = tenant
+                    account.current_tenant = new_tenant
-                    tenant_was_created.send(tenant)
+                    tenant_was_created.send(new_tenant)
 
         if account is None:
             try:

api/controllers/console/auth/oauth.py

@@ -148,15 +148,15 @@ def _generate_account(provider: str, user_info: OAuthUserInfo):
     account = _get_account_by_openid_or_email(provider, user_info)
 
     if account:
-        tenant = TenantService.get_join_tenants(account)
+        tenants = TenantService.get_join_tenants(account)
-        if not tenant:
+        if not tenants:
             if not FeatureService.get_system_features().is_allow_create_workspace:
                 raise WorkSpaceNotAllowedCreateError()
             else:
-                tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+                new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
-                TenantService.create_tenant_member(tenant, account, role="owner")
+                TenantService.create_tenant_member(new_tenant, account, role="owner")
-                account.current_tenant = tenant
+                account.current_tenant = new_tenant
-                tenant_was_created.send(tenant)
+                tenant_was_created.send(new_tenant)
 
     if not account:
         if not FeatureService.get_system_features().is_allow_register:

api/controllers/console/datasets/datasets.py

@@ -540,9 +540,22 @@ class DatasetIndexingStatusApi(Resource):
                 .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
                 .count()
             )
-            document.completed_segments = completed_segments
+            # Create a dictionary with document attributes and additional fields
-            document.total_segments = total_segments
+            document_dict = {
-            documents_status.append(marshal(document, document_status_fields))
+                "id": document.id,
+                "indexing_status": document.indexing_status,
+                "processing_started_at": document.processing_started_at,
+                "parsing_completed_at": document.parsing_completed_at,
+                "cleaning_completed_at": document.cleaning_completed_at,
+                "splitting_completed_at": document.splitting_completed_at,
+                "completed_at": document.completed_at,
+                "paused_at": document.paused_at,
+                "error": document.error,
+                "stopped_at": document.stopped_at,
+                "completed_segments": completed_segments,
+                "total_segments": total_segments,
+            }
+            documents_status.append(marshal(document_dict, document_status_fields))
         data = {"data": documents_status}
         return data
 

api/controllers/console/datasets/datasets_document.py

@@ -583,11 +583,22 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
                 .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
                 .count()
             )
-            document.completed_segments = completed_segments
+            # Create a dictionary with document attributes and additional fields
-            document.total_segments = total_segments
+            document_dict = {
-            if document.is_paused:
+                "id": document.id,
-                document.indexing_status = "paused"
+                "indexing_status": "paused" if document.is_paused else document.indexing_status,
-            documents_status.append(marshal(document, document_status_fields))
+                "processing_started_at": document.processing_started_at,
+                "parsing_completed_at": document.parsing_completed_at,
+                "cleaning_completed_at": document.cleaning_completed_at,
+                "splitting_completed_at": document.splitting_completed_at,
+                "completed_at": document.completed_at,
+                "paused_at": document.paused_at,
+                "error": document.error,
+                "stopped_at": document.stopped_at,
+                "completed_segments": completed_segments,
+                "total_segments": total_segments,
+            }
+            documents_status.append(marshal(document_dict, document_status_fields))
         data = {"data": documents_status}
         return data
 
@@ -616,11 +627,22 @@ class DocumentIndexingStatusApi(DocumentResource):
             .count()
         )
 
-        document.completed_segments = completed_segments
+        # Create a dictionary with document attributes and additional fields
-        document.total_segments = total_segments
+        document_dict = {
-        if document.is_paused:
+            "id": document.id,
-            document.indexing_status = "paused"
+            "indexing_status": "paused" if document.is_paused else document.indexing_status,
-        return marshal(document, document_status_fields)
+            "processing_started_at": document.processing_started_at,
+            "parsing_completed_at": document.parsing_completed_at,
+            "cleaning_completed_at": document.cleaning_completed_at,
+            "splitting_completed_at": document.splitting_completed_at,
+            "completed_at": document.completed_at,
+            "paused_at": document.paused_at,
+            "error": document.error,
+            "stopped_at": document.stopped_at,
+            "completed_segments": completed_segments,
+            "total_segments": total_segments,
+        }
+        return marshal(document_dict, document_status_fields)
 
 
 class DocumentDetailApi(DocumentResource):

api/controllers/console/workspace/workspace.py

@@ -68,16 +68,24 @@ class TenantListApi(Resource):
     @account_initialization_required
     def get(self):
         tenants = TenantService.get_join_tenants(current_user)
+        tenant_dicts = []
 
         for tenant in tenants:
             features = FeatureService.get_features(tenant.id)
-            if features.billing.enabled:
+
-                tenant.plan = features.billing.subscription.plan
+            # Create a dictionary with tenant attributes
-            else:
+            tenant_dict = {
-                tenant.plan = "sandbox"
+                "id": tenant.id,
-            if tenant.id == current_user.current_tenant_id:
+                "name": tenant.name,
-                tenant.current = True  # Set current=True for current tenant
+                "status": tenant.status,
-        return {"workspaces": marshal(tenants, tenants_fields)}, 200
+                "created_at": tenant.created_at,
+                "plan": features.billing.subscription.plan if features.billing.enabled else "sandbox",
+                "current": tenant.id == current_user.current_tenant_id,
+            }
+
+            tenant_dicts.append(tenant_dict)
+
+        return {"workspaces": marshal(tenant_dicts, tenants_fields)}, 200
 
 
 class WorkspaceListApi(Resource):

api/controllers/files/upload.py

@@ -64,9 +64,24 @@ class PluginUploadFileApi(Resource):
 
             extension = guess_extension(tool_file.mimetype) or ".bin"
             preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension)
-            tool_file.mime_type = mimetype
+
-            tool_file.extension = extension
+            # Create a dictionary with all the necessary attributes
-            tool_file.preview_url = preview_url
+            result = {
+                "id": tool_file.id,
+                "user_id": tool_file.user_id,
+                "tenant_id": tool_file.tenant_id,
+                "conversation_id": tool_file.conversation_id,
+                "file_key": tool_file.file_key,
+                "mimetype": tool_file.mimetype,
+                "original_url": tool_file.original_url,
+                "name": tool_file.name,
+                "size": tool_file.size,
+                "mime_type": mimetype,
+                "extension": extension,
+                "preview_url": preview_url,
+            }
+
+            return result, 201
         except services.errors.file.FileTooLargeError as file_too_large_error:
             raise FileTooLargeError(file_too_large_error.description)
         except services.errors.file.UnsupportedFileTypeError:

api/controllers/service_api/dataset/document.py

@@ -388,11 +388,22 @@ class DocumentIndexingStatusApi(DatasetApiResource):
                 .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
                 .count()
             )
-            document.completed_segments = completed_segments
+            # Create a dictionary with document attributes and additional fields
-            document.total_segments = total_segments
+            document_dict = {
-            if document.is_paused:
+                "id": document.id,
-                document.indexing_status = "paused"
+                "indexing_status": "paused" if document.is_paused else document.indexing_status,
-            documents_status.append(marshal(document, document_status_fields))
+                "processing_started_at": document.processing_started_at,
+                "parsing_completed_at": document.parsing_completed_at,
+                "cleaning_completed_at": document.cleaning_completed_at,
+                "splitting_completed_at": document.splitting_completed_at,
+                "completed_at": document.completed_at,
+                "paused_at": document.paused_at,
+                "error": document.error,
+                "stopped_at": document.stopped_at,
+                "completed_segments": completed_segments,
+                "total_segments": total_segments,
+            }
+            documents_status.append(marshal(document_dict, document_status_fields))
         data = {"data": documents_status}
         return data
 

api/core/rag/datasource/retrieval_service.py

@@ -405,7 +405,29 @@ class RetrievalService:
                     record["child_chunks"] = segment_child_map[record["segment"].id].get("child_chunks")  # type: ignore
                     record["score"] = segment_child_map[record["segment"].id]["max_score"]
 
-            return [RetrievalSegments(**record) for record in records]
+            result = []
+            for record in records:
+                # Extract segment
+                segment = record["segment"]
+
+                # Extract child_chunks, ensuring it's a list or None
+                child_chunks = record.get("child_chunks")
+                if not isinstance(child_chunks, list):
+                    child_chunks = None
+
+                # Extract score, ensuring it's a float or None
+                score_value = record.get("score")
+                score = (
+                    float(score_value)
+                    if score_value is not None and isinstance(score_value, int | float | str)
+                    else None
+                )
+
+                # Create RetrievalSegments object
+                retrieval_segment = RetrievalSegments(segment=segment, child_chunks=child_chunks, score=score)
+                result.append(retrieval_segment)
+
+            return result
         except Exception as e:
             db.session.rollback()
             raise e

api/core/tools/tool_manager.py

@@ -528,7 +528,7 @@ class ToolManager:
                     yield provider
 
                 except Exception:
-                    logger.exception(f"load builtin provider {provider}")
+                    logger.exception(f"load builtin provider {provider_path}")
                     continue
         # set builtin providers loaded
         cls._builtin_providers_loaded = True
@@ -644,10 +644,10 @@ class ToolManager:
                 )
 
                 workflow_provider_controllers: list[WorkflowToolProviderController] = []
-                for provider in workflow_providers:
+                for workflow_provider in workflow_providers:
                     try:
                         workflow_provider_controllers.append(
-                            ToolTransformService.workflow_provider_to_controller(db_provider=provider)
+                            ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider)
                         )
                     except Exception:
                         # app has been deleted

api/models/account.py

@@ -1,6 +1,6 @@
 import enum
 import json
-from typing import cast
+from typing import Optional, cast
 
 from flask_login import UserMixin  # type: ignore
 from sqlalchemy import func
@@ -12,6 +12,66 @@ from .engine import db
 from .types import StringUUID
 
 
+class TenantAccountRole(enum.StrEnum):
+    OWNER = "owner"
+    ADMIN = "admin"
+    EDITOR = "editor"
+    NORMAL = "normal"
+    DATASET_OPERATOR = "dataset_operator"
+
+    @staticmethod
+    def is_valid_role(role: str) -> bool:
+        if not role:
+            return False
+        return role in {
+            TenantAccountRole.OWNER,
+            TenantAccountRole.ADMIN,
+            TenantAccountRole.EDITOR,
+            TenantAccountRole.NORMAL,
+            TenantAccountRole.DATASET_OPERATOR,
+        }
+
+    @staticmethod
+    def is_privileged_role(role: Optional["TenantAccountRole"]) -> bool:
+        if not role:
+            return False
+        return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}
+
+    @staticmethod
+    def is_admin_role(role: Optional["TenantAccountRole"]) -> bool:
+        if not role:
+            return False
+        return role == TenantAccountRole.ADMIN
+
+    @staticmethod
+    def is_non_owner_role(role: Optional["TenantAccountRole"]) -> bool:
+        if not role:
+            return False
+        return role in {
+            TenantAccountRole.ADMIN,
+            TenantAccountRole.EDITOR,
+            TenantAccountRole.NORMAL,
+            TenantAccountRole.DATASET_OPERATOR,
+        }
+
+    @staticmethod
+    def is_editing_role(role: Optional["TenantAccountRole"]) -> bool:
+        if not role:
+            return False
+        return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}
+
+    @staticmethod
+    def is_dataset_edit_role(role: Optional["TenantAccountRole"]) -> bool:
+        if not role:
+            return False
+        return role in {
+            TenantAccountRole.OWNER,
+            TenantAccountRole.ADMIN,
+            TenantAccountRole.EDITOR,
+            TenantAccountRole.DATASET_OPERATOR,
+        }
+
+
 class AccountStatus(enum.StrEnum):
     PENDING = "pending"
     UNINITIALIZED = "uninitialized"
@@ -21,6 +81,8 @@ class AccountStatus(enum.StrEnum):
 
 
 class Account(UserMixin, Base):
+    role: Optional["TenantAccountRole"] = None
+    _current_tenant: Optional["Tenant"] = None
     __tablename__ = "accounts"
     __table_args__ = (db.PrimaryKeyConstraint("id", name="account_pkey"), db.Index("account_email_idx", "email"))
 
@@ -47,18 +109,16 @@ class Account(UserMixin, Base):
 
     @property
     def current_tenant(self):
-        return self._current_tenant  # type: ignore
+        return self._current_tenant
 
     @current_tenant.setter
-    def current_tenant(self, value: "Tenant"):
+    def current_tenant(self, tenant: "Tenant"):
-        tenant = value
         ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=self.id).first()
         if ta:
-            tenant.current_role = ta.role
+            self.role = TenantAccountRole(ta.role)
-        else:
+            self._current_tenant = tenant
-            tenant = None  # type: ignore
+            return
-
+        self._current_tenant = None
-        self._current_tenant = tenant
 
     @property
     def current_tenant_id(self) -> str | None:
@@ -80,12 +140,12 @@ class Account(UserMixin, Base):
             return
 
         tenant, join = tenant_account_join
-        tenant.current_role = join.role
+        self.role = join.role
         self._current_tenant = tenant
 
     @property
     def current_role(self):
-        return self._current_tenant.current_role
+        return self.role
 
     def get_status(self) -> AccountStatus:
         status_str = self.status
@@ -105,23 +165,23 @@ class Account(UserMixin, Base):
     # check current_user.current_tenant.current_role in ['admin', 'owner']
     @property
     def is_admin_or_owner(self):
-        return TenantAccountRole.is_privileged_role(self._current_tenant.current_role)
+        return TenantAccountRole.is_privileged_role(self.role)
 
     @property
     def is_admin(self):
-        return TenantAccountRole.is_admin_role(self._current_tenant.current_role)
+        return TenantAccountRole.is_admin_role(self.role)
 
     @property
     def is_editor(self):
-        return TenantAccountRole.is_editing_role(self._current_tenant.current_role)
+        return TenantAccountRole.is_editing_role(self.role)
 
     @property
     def is_dataset_editor(self):
-        return TenantAccountRole.is_dataset_edit_role(self._current_tenant.current_role)
+        return TenantAccountRole.is_dataset_edit_role(self.role)
 
     @property
     def is_dataset_operator(self):
-        return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR
+        return self.role == TenantAccountRole.DATASET_OPERATOR
 
 
 class TenantStatus(enum.StrEnum):
@@ -129,66 +189,6 @@ class TenantStatus(enum.StrEnum):
     ARCHIVE = "archive"
 
 
-class TenantAccountRole(enum.StrEnum):
-    OWNER = "owner"
-    ADMIN = "admin"
-    EDITOR = "editor"
-    NORMAL = "normal"
-    DATASET_OPERATOR = "dataset_operator"
-
-    @staticmethod
-    def is_valid_role(role: str) -> bool:
-        if not role:
-            return False
-        return role in {
-            TenantAccountRole.OWNER,
-            TenantAccountRole.ADMIN,
-            TenantAccountRole.EDITOR,
-            TenantAccountRole.NORMAL,
-            TenantAccountRole.DATASET_OPERATOR,
-        }
-
-    @staticmethod
-    def is_privileged_role(role: str) -> bool:
-        if not role:
-            return False
-        return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}
-
-    @staticmethod
-    def is_admin_role(role: str) -> bool:
-        if not role:
-            return False
-        return role == TenantAccountRole.ADMIN
-
-    @staticmethod
-    def is_non_owner_role(role: str) -> bool:
-        if not role:
-            return False
-        return role in {
-            TenantAccountRole.ADMIN,
-            TenantAccountRole.EDITOR,
-            TenantAccountRole.NORMAL,
-            TenantAccountRole.DATASET_OPERATOR,
-        }
-
-    @staticmethod
-    def is_editing_role(role: str) -> bool:
-        if not role:
-            return False
-        return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}
-
-    @staticmethod
-    def is_dataset_edit_role(role: str) -> bool:
-        if not role:
-            return False
-        return role in {
-            TenantAccountRole.OWNER,
-            TenantAccountRole.ADMIN,
-            TenantAccountRole.EDITOR,
-            TenantAccountRole.DATASET_OPERATOR,
-        }
-
-
 class Tenant(Base):
     __tablename__ = "tenants"
     __table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),)

api/models/base.py

@@ -1,5 +1,7 @@
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from models.engine import metadata
 
-Base = declarative_base(metadata=metadata)
+
+class Base(DeclarativeBase):
+    metadata = metadata

api/models/tools.py

@@ -172,10 +172,6 @@ class WorkflowToolProvider(Base):
         db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
     )
 
-    @property
-    def schema_type(self) -> ApiProviderSchemaType:
-        return ApiProviderSchemaType.value_of(self.schema_type_str)
-
     @property
     def user(self) -> Account | None:
         return db.session.query(Account).filter(Account.id == self.user_id).first()

api/models/workflow.py

@@ -3,7 +3,7 @@ import logging
 from collections.abc import Mapping, Sequence
 from datetime import UTC, datetime
 from enum import Enum, StrEnum
-from typing import TYPE_CHECKING, Any, Optional, Self, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 from uuid import uuid4
 
 from core.variables import utils as variable_utils
@@ -150,7 +150,7 @@ class Workflow(Base):
         conversation_variables: Sequence[Variable],
         marked_name: str = "",
         marked_comment: str = "",
-    ) -> Self:
+    ) -> "Workflow":
         workflow = Workflow()
         workflow.id = str(uuid4())
         workflow.tenant_id = tenant_id

api/services/vector_service.py

@@ -23,11 +23,10 @@ class VectorService:
     ):
         documents: list[Document] = []
 
-        document: Document | None = None
         for segment in segments:
             if doc_form == IndexType.PARENT_CHILD_INDEX:
-                document = db.session.query(DatasetDocument).filter_by(id=segment.document_id).first()
+                dataset_document = db.session.query(DatasetDocument).filter_by(id=segment.document_id).first()
-                if not document:
+                if not dataset_document:
                     _logger.warning(
                         "Expected DatasetDocument record to exist, but none was found, document_id=%s, segment_id=%s",
                         segment.document_id,
@@ -37,7 +36,7 @@ class VectorService:
                 # get the process rule
                 processing_rule = (
                     db.session.query(DatasetProcessRule)
-                    .filter(DatasetProcessRule.id == document.dataset_process_rule_id)
+                    .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
                     .first()
                 )
                 if not processing_rule:
@@ -61,9 +60,11 @@ class VectorService:
                         )
                 else:
                     raise ValueError("The knowledge base index technique is not high quality!")
-                cls.generate_child_chunks(segment, document, dataset, embedding_model_instance, processing_rule, False)
+                cls.generate_child_chunks(
+                    segment, dataset_document, dataset, embedding_model_instance, processing_rule, False
+                )
             else:
-                document = Document(
+                rag_document = Document(
                     page_content=segment.content,
                     metadata={
                         "doc_id": segment.index_node_id,
@@ -72,7 +73,7 @@ class VectorService:
                         "dataset_id": segment.dataset_id,
                     },
                 )
-                documents.append(document)
+                documents.append(rag_document)
         if len(documents) > 0:
             index_processor = IndexProcessorFactory(doc_form).init_index_processor()
             index_processor.load(dataset, documents, with_keywords=True, keywords_list=keywords_list)

api/services/workflow_service.py

@@ -508,11 +508,11 @@ class WorkflowService:
             raise DraftWorkflowDeletionError("Cannot delete draft workflow versions")
 
         # Check if this workflow is currently referenced by an app
-        stmt = select(App).where(App.workflow_id == workflow_id)
+        app_stmt = select(App).where(App.workflow_id == workflow_id)
-        app = session.scalar(stmt)
+        app = session.scalar(app_stmt)
         if app:
             # Cannot delete a workflow that's currently in use by an app
-            raise WorkflowInUseError(f"Cannot delete workflow that is currently in use by app '{app.name}'")
+            raise WorkflowInUseError(f"Cannot delete workflow that is currently in use by app '{app.id}'")
 
         # Don't use workflow.tool_published as it's not accurate for specific workflow versions
         # Check if there's a tool provider using this specific workflow version

api/tasks/add_document_to_index_task.py

@@ -111,7 +111,7 @@ def add_document_to_index_task(dataset_document_id: str):
         logging.exception("add document to index failed")
         dataset_document.enabled = False
         dataset_document.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
-        dataset_document.status = "error"
+        dataset_document.indexing_status = "error"
         dataset_document.error = str(e)
         db.session.commit()
     finally:

api/tasks/remove_app_and_related_data_task.py

@@ -193,7 +193,7 @@ def _delete_app_workflow_runs(tenant_id: str, app_id: str):
 def _delete_app_workflow_node_executions(tenant_id: str, app_id: str):
     # Get app's owner
     with Session(db.engine, expire_on_commit=False) as session:
-        stmt = select(Account).where(Account.id == App.owner_id).where(App.id == app_id)
+        stmt = select(Account).where(Account.id == App.created_by).where(App.id == app_id)
         user = session.scalar(stmt)
 
     if user is None: