refactor: Use typed SQLAlchemy base model and fix type errors (#19980)

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)
-            if not tenant:
+            tenants = TenantService.get_join_tenants(account)
+            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")
-                    TenantService.create_tenant_member(tenant, account, role="owner")
-                    account.current_tenant = tenant
-                    tenant_was_created.send(tenant)
+                    new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+                    TenantService.create_tenant_member(new_tenant, account, role="owner")
+                    account.current_tenant = new_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)
-        if not tenant:
+        tenants = TenantService.get_join_tenants(account)
+        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")
-                TenantService.create_tenant_member(tenant, account, role="owner")
-                account.current_tenant = tenant
-                tenant_was_created.send(tenant)
+                new_tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+                TenantService.create_tenant_member(new_tenant, account, role="owner")
+                account.current_tenant = new_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
-            document.total_segments = total_segments
-            documents_status.append(marshal(document, document_status_fields))
+            # Create a dictionary with document attributes and additional fields
+            document_dict = {
+                "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
-            document.total_segments = total_segments
-            if document.is_paused:
-                document.indexing_status = "paused"
-            documents_status.append(marshal(document, document_status_fields))
+            # Create a dictionary with document attributes and additional fields
+            document_dict = {
+                "id": document.id,
+                "indexing_status": "paused" if document.is_paused else 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
 
@@ -616,11 +627,22 @@ class DocumentIndexingStatusApi(DocumentResource):
             .count()
         )
 
-        document.completed_segments = completed_segments
-        document.total_segments = total_segments
-        if document.is_paused:
-            document.indexing_status = "paused"
-        return marshal(document, document_status_fields)
+        # Create a dictionary with document attributes and additional fields
+        document_dict = {
+            "id": document.id,
+            "indexing_status": "paused" if document.is_paused else 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,
+        }
+        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
-            else:
-                tenant.plan = "sandbox"
-            if tenant.id == current_user.current_tenant_id:
-                tenant.current = True  # Set current=True for current tenant
-        return {"workspaces": marshal(tenants, tenants_fields)}, 200
+
+            # Create a dictionary with tenant attributes
+            tenant_dict = {
+                "id": tenant.id,
+                "name": tenant.name,
+                "status": tenant.status,
+                "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
-            tool_file.preview_url = preview_url
+
+            # Create a dictionary with all the necessary attributes
+            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
-            document.total_segments = total_segments
-            if document.is_paused:
-                document.indexing_status = "paused"
-            documents_status.append(marshal(document, document_status_fields))
+            # Create a dictionary with document attributes and additional fields
+            document_dict = {
+                "id": document.id,
+                "indexing_status": "paused" if document.is_paused else 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