diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index 116fc59ee8..37d351627b 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -39,25 +39,11 @@ jobs: - name: Install dependencies run: uv sync --project api --dev - - name: Run Unit tests - run: | - uv run --project api bash dev/pytest/pytest_unit_tests.sh - - name: Run pyrefly check run: | cd api uv add --dev pyrefly uv run pyrefly check || true - - name: Coverage Summary - run: | - set -x - # Extract coverage percentage and create a summary - TOTAL_COVERAGE=$(python -c 'import json; print(json.load(open("coverage.json"))["totals"]["percent_covered_display"])') - - # Create a detailed coverage summary - echo "### Test Coverage Summary :test_tube:" >> $GITHUB_STEP_SUMMARY - echo "Total Coverage: ${TOTAL_COVERAGE}%" >> $GITHUB_STEP_SUMMARY - uv run --project api coverage report --format=markdown >> $GITHUB_STEP_SUMMARY - name: Run dify config tests run: uv run --project api dev/pytest/pytest_config_tests.py @@ -93,3 +79,19 @@ jobs: - name: Run TestContainers run: uv run --project api bash dev/pytest/pytest_testcontainers.sh + + - name: Run Unit tests + run: | + uv run --project api bash dev/pytest/pytest_unit_tests.sh + + - name: Coverage Summary + run: | + set -x + # Extract coverage percentage and create a summary + TOTAL_COVERAGE=$(python -c 'import json; print(json.load(open("coverage.json"))["totals"]["percent_covered_display"])') + + # Create a detailed coverage summary + echo "### Test Coverage Summary :test_tube:" >> $GITHUB_STEP_SUMMARY + echo "Total Coverage: ${TOTAL_COVERAGE}%" >> $GITHUB_STEP_SUMMARY + uv run --project api coverage report --format=markdown >> $GITHUB_STEP_SUMMARY + diff --git a/api/controllers/console/apikey.py b/api/controllers/console/apikey.py index 7c39b04464..4f04af7932 100644 --- a/api/controllers/console/apikey.py +++ b/api/controllers/console/apikey.py @@ -12,7 +12,7 @@ from models.dataset import Dataset from models.model import ApiToken, App from . import api, console_ns -from .wraps import account_initialization_required, setup_required +from .wraps import account_initialization_required, edit_permission_required, setup_required api_key_fields = { "id": fields.String, @@ -67,14 +67,12 @@ class BaseApiKeyListResource(Resource): return {"items": keys} @marshal_with(api_key_fields) + @edit_permission_required def post(self, resource_id): assert self.resource_id_field is not None, "resource_id_field must be set" resource_id = str(resource_id) - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() _get_resource(resource_id, current_tenant_id, self.resource_model) - if not current_user.has_edit_permission: - raise Forbidden() - current_key_count = ( db.session.query(ApiToken) .where(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id) diff --git a/api/controllers/console/app/annotation.py b/api/controllers/console/app/annotation.py index 589955738c..3e549d869e 100644 --- a/api/controllers/console/app/annotation.py +++ b/api/controllers/console/app/annotation.py @@ -2,13 +2,13 @@ from typing import Literal from flask import request from flask_restx import Resource, fields, marshal, marshal_with, reqparse -from werkzeug.exceptions import Forbidden from controllers.common.errors import NoFileUploadedError, TooManyFilesError from controllers.console import api, console_ns from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, + edit_permission_required, setup_required, ) from extensions.ext_redis import redis_client @@ -16,7 +16,7 @@ from fields.annotation_fields import ( annotation_fields, annotation_hit_history_fields, ) -from libs.login import current_account_with_tenant, login_required +from libs.login import login_required from services.annotation_service import AppAnnotationService @@ -41,12 +41,8 @@ class AnnotationReplyActionApi(Resource): @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") + @edit_permission_required def post(self, app_id, action: Literal["enable", "disable"]): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) parser = reqparse.RequestParser() parser.add_argument("score_threshold", required=True, type=float, location="json") @@ -70,12 +66,8 @@ class AppAnnotationSettingDetailApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) result = AppAnnotationService.get_app_annotation_setting_by_app_id(app_id) return result, 200 @@ -101,12 +93,8 @@ class AppAnnotationSettingUpdateApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, app_id, annotation_setting_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) annotation_setting_id = str(annotation_setting_id) @@ -129,12 +117,8 @@ class AnnotationReplyActionStatusApi(Resource): @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") + @edit_permission_required def get(self, app_id, job_id, action): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - job_id = str(job_id) app_annotation_job_key = f"{action}_app_annotation_job_{str(job_id)}" cache_result = redis_client.get(app_annotation_job_key) @@ -166,12 +150,8 @@ class AnnotationApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) keyword = request.args.get("keyword", default="", type=str) @@ -207,12 +187,8 @@ class AnnotationApi(Resource): @account_initialization_required @cloud_edition_billing_resource_check("annotation") @marshal_with(annotation_fields) + @edit_permission_required def post(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) parser = reqparse.RequestParser() parser.add_argument("question", required=True, type=str, location="json") @@ -224,12 +200,8 @@ class AnnotationApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def delete(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) # Use request.args.getlist to get annotation_ids array directly @@ -262,12 +234,8 @@ class AnnotationExportApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) annotation_list = AppAnnotationService.export_annotation_list_by_app_id(app_id) response = {"data": marshal(annotation_list, annotation_fields)} @@ -286,13 +254,9 @@ class AnnotationUpdateDeleteApi(Resource): @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") + @edit_permission_required @marshal_with(annotation_fields) def post(self, app_id, annotation_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) annotation_id = str(annotation_id) parser = reqparse.RequestParser() @@ -305,12 +269,8 @@ class AnnotationUpdateDeleteApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def delete(self, app_id, annotation_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) annotation_id = str(annotation_id) AppAnnotationService.delete_app_annotation(app_id, annotation_id) @@ -329,12 +289,8 @@ class AnnotationBatchImportApi(Resource): @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") + @edit_permission_required def post(self, app_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - app_id = str(app_id) # check file if "file" not in request.files: @@ -362,12 +318,8 @@ class AnnotationBatchImportStatusApi(Resource): @login_required @account_initialization_required @cloud_edition_billing_resource_check("annotation") + @edit_permission_required def get(self, app_id, job_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - job_id = str(job_id) indexing_cache_key = f"app_annotation_batch_import_{str(job_id)}" cache_result = redis_client.get(indexing_cache_key) @@ -399,12 +351,8 @@ class AnnotationHitHistoryListApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, app_id, annotation_id): - current_user, _ = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() - page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) app_id = str(app_id) diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index 3927685af3..3900f5a6eb 100644 --- a/api/controllers/console/app/app.py +++ b/api/controllers/console/app/app.py @@ -1,7 +1,5 @@ import uuid -from typing import cast -from flask_login import current_user from flask_restx import Resource, fields, inputs, marshal, marshal_with, reqparse from sqlalchemy import select from sqlalchemy.orm import Session @@ -12,15 +10,16 @@ from controllers.console.app.wraps import get_app_model from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, + edit_permission_required, enterprise_license_required, setup_required, ) from core.ops.ops_trace_manager import OpsTraceManager from extensions.ext_database import db from fields.app_fields import app_detail_fields, app_detail_fields_with_site, app_pagination_fields -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from libs.validators import validate_description_length -from models import Account, App +from models import App from services.app_dsl_service import AppDslService, ImportMode from services.app_service import AppService from services.enterprise.enterprise_service import EnterpriseService @@ -56,6 +55,7 @@ class AppListApi(Resource): @enterprise_license_required def get(self): """Get app list""" + current_user, current_tenant_id = current_account_with_tenant() def uuid_list(value): try: @@ -90,7 +90,7 @@ class AppListApi(Resource): # get app list app_service = AppService() - app_pagination = app_service.get_paginate_apps(current_user.id, current_user.current_tenant_id, args) + app_pagination = app_service.get_paginate_apps(current_user.id, current_tenant_id, args) if not app_pagination: return {"data": [], "total": 0, "page": 1, "limit": 20, "has_more": False} @@ -129,8 +129,10 @@ class AppListApi(Resource): @account_initialization_required @marshal_with(app_detail_fields) @cloud_edition_billing_resource_check("apps") + @edit_permission_required def post(self): """Create app""" + current_user, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, location="json") parser.add_argument("description", type=validate_description_length, location="json") @@ -140,19 +142,11 @@ class AppListApi(Resource): parser.add_argument("icon_background", type=str, location="json") args = parser.parse_args() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - if "mode" not in args or args["mode"] is None: raise BadRequest("mode is required") app_service = AppService() - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") - if current_user.current_tenant_id is None: - raise ValueError("current_user.current_tenant_id cannot be None") - app = app_service.create_app(current_user.current_tenant_id, args, current_user) + app = app_service.create_app(current_tenant_id, args, current_user) return app, 201 @@ -205,13 +199,10 @@ class AppApi(Resource): @login_required @account_initialization_required @get_app_model + @edit_permission_required @marshal_with(app_detail_fields_with_site) def put(self, app_model): """Update app""" - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, nullable=False, location="json") parser.add_argument("description", type=validate_description_length, location="json") @@ -248,12 +239,9 @@ class AppApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def delete(self, app_model): """Delete app""" - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - app_service = AppService() app_service.delete_app(app_model) @@ -283,12 +271,12 @@ class AppCopyApi(Resource): @login_required @account_initialization_required @get_app_model + @edit_permission_required @marshal_with(app_detail_fields_with_site) def post(self, app_model): """Copy app""" # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, location="json") @@ -301,9 +289,8 @@ class AppCopyApi(Resource): with Session(db.engine) as session: import_service = AppDslService(session) yaml_content = import_service.export_dsl(app_model=app_model, include_secret=True) - account = cast(Account, current_user) result = import_service.import_app( - account=account, + account=current_user, import_mode=ImportMode.YAML_CONTENT, yaml_content=yaml_content, name=args.get("name"), @@ -340,12 +327,9 @@ class AppExportApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, app_model): """Export app""" - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - # Add include_secret params parser = reqparse.RequestParser() parser.add_argument("include_secret", type=inputs.boolean, default=False, location="args") @@ -371,11 +355,8 @@ class AppNameApi(Resource): @account_initialization_required @get_app_model @marshal_with(app_detail_fields) + @edit_permission_required def post(self, app_model): - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, location="json") args = parser.parse_args() @@ -408,11 +389,8 @@ class AppIconApi(Resource): @account_initialization_required @get_app_model @marshal_with(app_detail_fields) + @edit_permission_required def post(self, app_model): - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("icon", type=str, location="json") parser.add_argument("icon_background", type=str, location="json") @@ -441,11 +419,8 @@ class AppSiteStatus(Resource): @account_initialization_required @get_app_model @marshal_with(app_detail_fields) + @edit_permission_required def post(self, app_model): - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("enable_site", type=bool, required=True, location="json") args = parser.parse_args() @@ -475,6 +450,7 @@ class AppApiStatus(Resource): @marshal_with(app_detail_fields) def post(self, app_model): # The role of the current user in the ta table must be admin or owner + current_user, _ = current_account_with_tenant() if not current_user.is_admin_or_owner: raise Forbidden() @@ -520,10 +496,9 @@ class AppTraceApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, app_id): # add app trace - if not current_user.is_editor: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("enabled", type=bool, required=True, location="json") parser.add_argument("tracing_provider", type=str, required=True, location="json") diff --git a/api/controllers/console/app/app_import.py b/api/controllers/console/app/app_import.py index 5751ff1f86..5e7ea6d481 100644 --- a/api/controllers/console/app/app_import.py +++ b/api/controllers/console/app/app_import.py @@ -1,11 +1,11 @@ from flask_restx import Resource, marshal_with, reqparse from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden from controllers.console.app.wraps import get_app_model from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, + edit_permission_required, setup_required, ) from extensions.ext_database import db @@ -26,12 +26,10 @@ class AppImportApi(Resource): @account_initialization_required @marshal_with(app_import_fields) @cloud_edition_billing_resource_check("apps") + @edit_permission_required def post(self): # Check user role first current_user, _ = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("mode", type=str, required=True, location="json") parser.add_argument("yaml_content", type=str, location="json") @@ -80,11 +78,10 @@ class AppImportConfirmApi(Resource): @login_required @account_initialization_required @marshal_with(app_import_fields) + @edit_permission_required def post(self, import_id): # Check user role first current_user, _ = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() # Create service with session with Session(db.engine) as session: @@ -107,11 +104,8 @@ class AppImportCheckDependenciesApi(Resource): @get_app_model @account_initialization_required @marshal_with(app_import_check_dependencies_fields) + @edit_permission_required def get(self, app_model: App): - current_user, _ = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() - with Session(db.engine) as session: import_service = AppDslService(session) result = import_service.check_dependencies(app_model=app_model) diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index 2f7b90e7fb..d69f05f23e 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -2,7 +2,7 @@ import logging from flask import request from flask_restx import Resource, fields, reqparse -from werkzeug.exceptions import Forbidden, InternalServerError, NotFound +from werkzeug.exceptions import InternalServerError, NotFound import services from controllers.console import api, console_ns @@ -15,7 +15,7 @@ from controllers.console.app.error import ( ProviderQuotaExceededError, ) from controllers.console.app.wraps import get_app_model -from controllers.console.wraps import account_initialization_required, setup_required +from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom @@ -151,13 +151,8 @@ class ChatMessageApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT]) + @edit_permission_required def post(self, app_model): - if not isinstance(current_user, Account): - raise Forbidden() - - if not current_user.has_edit_permission: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, required=True, location="json") parser.add_argument("query", type=str, required=True, location="json") diff --git a/api/controllers/console/app/conversation.py b/api/controllers/console/app/conversation.py index 3b8dff613b..779be62973 100644 --- a/api/controllers/console/app/conversation.py +++ b/api/controllers/console/app/conversation.py @@ -1,17 +1,16 @@ from datetime import datetime -import pytz # pip install pytz +import pytz import sqlalchemy as sa -from flask_login import current_user from flask_restx import Resource, marshal_with, reqparse from flask_restx.inputs import int_range from sqlalchemy import func, or_ from sqlalchemy.orm import joinedload -from werkzeug.exceptions import Forbidden, NotFound +from werkzeug.exceptions import NotFound from controllers.console import api, console_ns from controllers.console.app.wraps import get_app_model -from controllers.console.wraps import account_initialization_required, setup_required +from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from core.app.entities.app_invoke_entities import InvokeFrom from extensions.ext_database import db from fields.conversation_fields import ( @@ -22,8 +21,8 @@ from fields.conversation_fields import ( ) from libs.datetime_utils import naive_utc_now from libs.helper import DatetimeString -from libs.login import login_required -from models import Account, Conversation, EndUser, Message, MessageAnnotation +from libs.login import current_account_with_tenant, login_required +from models import Conversation, EndUser, Message, MessageAnnotation from models.model import AppMode from services.conversation_service import ConversationService from services.errors.conversation import ConversationNotExistsError @@ -57,9 +56,9 @@ class CompletionConversationApi(Resource): @account_initialization_required @get_app_model(mode=AppMode.COMPLETION) @marshal_with(conversation_pagination_fields) + @edit_permission_required def get(self, app_model): - if not current_user.is_editor: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("keyword", type=str, location="args") parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -84,6 +83,7 @@ class CompletionConversationApi(Resource): ) account = current_user + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -137,9 +137,8 @@ class CompletionConversationDetailApi(Resource): @account_initialization_required @get_app_model(mode=AppMode.COMPLETION) @marshal_with(conversation_message_detail_fields) + @edit_permission_required def get(self, app_model, conversation_id): - if not current_user.is_editor: - raise Forbidden() conversation_id = str(conversation_id) return _get_conversation(app_model, conversation_id) @@ -154,14 +153,12 @@ class CompletionConversationDetailApi(Resource): @login_required @account_initialization_required @get_app_model(mode=AppMode.COMPLETION) + @edit_permission_required def delete(self, app_model, conversation_id): - if not current_user.is_editor: - raise Forbidden() + current_user, _ = current_account_with_tenant() conversation_id = str(conversation_id) try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") ConversationService.delete(app_model, conversation_id, current_user) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -206,9 +203,9 @@ class ChatConversationApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) @marshal_with(conversation_with_summary_pagination_fields) + @edit_permission_required def get(self, app_model): - if not current_user.is_editor: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("keyword", type=str, location="args") parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -260,6 +257,7 @@ class ChatConversationApi(Resource): ) account = current_user + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -341,9 +339,8 @@ class ChatConversationDetailApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) @marshal_with(conversation_detail_fields) + @edit_permission_required def get(self, app_model, conversation_id): - if not current_user.is_editor: - raise Forbidden() conversation_id = str(conversation_id) return _get_conversation(app_model, conversation_id) @@ -358,14 +355,12 @@ class ChatConversationDetailApi(Resource): @login_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) @account_initialization_required + @edit_permission_required def delete(self, app_model, conversation_id): - if not current_user.is_editor: - raise Forbidden() + current_user, _ = current_account_with_tenant() conversation_id = str(conversation_id) try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") ConversationService.delete(app_model, conversation_id, current_user) except ConversationNotExistsError: raise NotFound("Conversation Not Exists.") @@ -374,6 +369,7 @@ class ChatConversationDetailApi(Resource): def _get_conversation(app_model, conversation_id): + current_user, _ = current_account_with_tenant() conversation = ( db.session.query(Conversation) .where(Conversation.id == conversation_id, Conversation.app_id == app_model.id) diff --git a/api/controllers/console/app/mcp_server.py b/api/controllers/console/app/mcp_server.py index b9a383ee61..599f5adb34 100644 --- a/api/controllers/console/app/mcp_server.py +++ b/api/controllers/console/app/mcp_server.py @@ -1,16 +1,15 @@ import json from enum import StrEnum -from flask_login import current_user from flask_restx import Resource, fields, marshal_with, reqparse from werkzeug.exceptions import NotFound from controllers.console import api, console_ns from controllers.console.app.wraps import get_app_model -from controllers.console.wraps import account_initialization_required, setup_required +from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from extensions.ext_database import db from fields.app_fields import app_server_fields -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from models.model import AppMCPServer @@ -25,9 +24,9 @@ class AppMCPServerController(Resource): @api.doc(description="Get MCP server configuration for an application") @api.doc(params={"app_id": "Application ID"}) @api.response(200, "MCP server configuration retrieved successfully", app_server_fields) - @setup_required @login_required @account_initialization_required + @setup_required @get_app_model @marshal_with(app_server_fields) def get(self, app_model): @@ -48,14 +47,14 @@ class AppMCPServerController(Resource): ) @api.response(201, "MCP server configuration created successfully", app_server_fields) @api.response(403, "Insufficient permissions") - @setup_required - @login_required @account_initialization_required @get_app_model + @login_required + @setup_required @marshal_with(app_server_fields) + @edit_permission_required def post(self, app_model): - if not current_user.is_editor: - raise NotFound() + _, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("description", type=str, required=False, location="json") parser.add_argument("parameters", type=dict, required=True, location="json") @@ -71,7 +70,7 @@ class AppMCPServerController(Resource): parameters=json.dumps(args["parameters"], ensure_ascii=False), status=AppMCPServerStatus.ACTIVE, app_id=app_model.id, - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, server_code=AppMCPServer.generate_server_code(16), ) db.session.add(server) @@ -95,14 +94,13 @@ class AppMCPServerController(Resource): @api.response(200, "MCP server configuration updated successfully", app_server_fields) @api.response(403, "Insufficient permissions") @api.response(404, "Server not found") - @setup_required - @login_required - @account_initialization_required @get_app_model + @login_required + @setup_required + @account_initialization_required @marshal_with(app_server_fields) + @edit_permission_required def put(self, app_model): - if not current_user.is_editor: - raise NotFound() parser = reqparse.RequestParser() parser.add_argument("id", type=str, required=True, location="json") parser.add_argument("description", type=str, required=False, location="json") @@ -142,13 +140,13 @@ class AppMCPServerRefreshController(Resource): @login_required @account_initialization_required @marshal_with(app_server_fields) + @edit_permission_required def get(self, server_id): - if not current_user.is_editor: - raise NotFound() + _, current_tenant_id = current_account_with_tenant() server = ( db.session.query(AppMCPServer) .where(AppMCPServer.id == server_id) - .where(AppMCPServer.tenant_id == current_user.current_tenant_id) + .where(AppMCPServer.tenant_id == current_tenant_id) .first() ) if not server: diff --git a/api/controllers/console/app/message.py b/api/controllers/console/app/message.py index 46523feccc..005cff75fc 100644 --- a/api/controllers/console/app/message.py +++ b/api/controllers/console/app/message.py @@ -3,7 +3,7 @@ import logging from flask_restx import Resource, fields, marshal_with, reqparse from flask_restx.inputs import int_range from sqlalchemy import exists, select -from werkzeug.exceptions import Forbidden, InternalServerError, NotFound +from werkzeug.exceptions import InternalServerError, NotFound from controllers.console import api, console_ns from controllers.console.app.error import ( @@ -17,6 +17,7 @@ from controllers.console.explore.error import AppSuggestedQuestionsAfterAnswerDi from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, + edit_permission_required, setup_required, ) from core.app.entities.app_invoke_entities import InvokeFrom @@ -26,8 +27,7 @@ from extensions.ext_database import db from fields.conversation_fields import annotation_fields, message_detail_fields from libs.helper import uuid_value from libs.infinite_scroll_pagination import InfiniteScrollPagination -from libs.login import current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from models.model import AppMode, Conversation, Message, MessageAnnotation, MessageFeedback from services.annotation_service import AppAnnotationService from services.errors.conversation import ConversationNotExistsError @@ -56,15 +56,13 @@ class ChatMessageListApi(Resource): ) @api.response(200, "Success", message_infinite_scroll_pagination_fields) @api.response(404, "Conversation not found") - @setup_required @login_required - @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) @account_initialization_required + @setup_required + @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) @marshal_with(message_infinite_scroll_pagination_fields) + @edit_permission_required def get(self, app_model): - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("conversation_id", required=True, type=uuid_value, location="args") parser.add_argument("first_id", type=uuid_value, location="args") @@ -154,8 +152,7 @@ class MessageFeedbackApi(Resource): @login_required @account_initialization_required def post(self, app_model): - if current_user is None: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("message_id", required=True, type=uuid_value, location="json") @@ -211,18 +208,14 @@ class MessageAnnotationApi(Resource): ) @api.response(200, "Annotation created successfully", annotation_fields) @api.response(403, "Insufficient permissions") + @marshal_with(annotation_fields) + @get_app_model @setup_required @login_required - @account_initialization_required @cloud_edition_billing_resource_check("annotation") - @get_app_model - @marshal_with(annotation_fields) + @account_initialization_required + @edit_permission_required def post(self, app_model): - if not isinstance(current_user, Account): - raise Forbidden() - if not current_user.has_edit_permission: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("message_id", required=False, type=uuid_value, location="json") parser.add_argument("question", required=True, type=str, location="json") @@ -270,6 +263,7 @@ class MessageSuggestedQuestionApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) def get(self, app_model, message_id): + current_user, _ = current_account_with_tenant() message_id = str(message_id) try: @@ -304,12 +298,12 @@ class MessageApi(Resource): @api.doc(params={"app_id": "Application ID", "message_id": "Message ID"}) @api.response(200, "Message retrieved successfully", message_detail_fields) @api.response(404, "Message not found") + @get_app_model @setup_required @login_required @account_initialization_required - @get_app_model @marshal_with(message_detail_fields) - def get(self, app_model, message_id): + def get(self, app_model, message_id: str): message_id = str(message_id) message = db.session.query(Message).where(Message.id == message_id, Message.app_id == app_model.id).first() diff --git a/api/controllers/console/app/site.py b/api/controllers/console/app/site.py index 6537ea5a50..1da704efcc 100644 --- a/api/controllers/console/app/site.py +++ b/api/controllers/console/app/site.py @@ -9,7 +9,7 @@ from extensions.ext_database import db from fields.app_fields import app_site_fields from libs.datetime_utils import naive_utc_now from libs.login import current_account_with_tenant, login_required -from models import Account, Site +from models import Site def parse_app_site_args(): @@ -107,8 +107,6 @@ class AppSite(Resource): if value is not None: setattr(site, attr_name, value) - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") site.updated_by = current_user.id site.updated_at = naive_utc_now() db.session.commit() @@ -142,8 +140,6 @@ class AppSiteAccessTokenReset(Resource): raise NotFound site.code = Site.generate_code(16) - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") site.updated_by = current_user.id site.updated_at = naive_utc_now() db.session.commit() diff --git a/api/controllers/console/app/statistic.py b/api/controllers/console/app/statistic.py index 5974395c6a..cfe5b3ff17 100644 --- a/api/controllers/console/app/statistic.py +++ b/api/controllers/console/app/statistic.py @@ -4,7 +4,6 @@ from decimal import Decimal import pytz import sqlalchemy as sa from flask import jsonify -from flask_login import current_user from flask_restx import Resource, fields, reqparse from controllers.console import api, console_ns @@ -13,7 +12,7 @@ from controllers.console.wraps import account_initialization_required, setup_req from core.app.entities.app_invoke_entities import InvokeFrom from extensions.ext_database import db from libs.helper import DatetimeString -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from models import AppMode, Message @@ -37,7 +36,7 @@ class DailyMessageStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -53,6 +52,7 @@ WHERE app_id = :app_id AND invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -109,13 +109,13 @@ class DailyConversationStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") parser.add_argument("end", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") args = parser.parse_args() - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -175,7 +175,7 @@ class DailyTerminalsStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -191,7 +191,7 @@ WHERE app_id = :app_id AND invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -247,7 +247,7 @@ class DailyTokenCostStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -264,7 +264,7 @@ WHERE app_id = :app_id AND invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -322,7 +322,7 @@ class AverageSessionInteractionStatistic(Resource): @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT]) def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -346,7 +346,7 @@ FROM c.app_id = :app_id AND m.invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -413,7 +413,7 @@ class UserSatisfactionRateStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -433,7 +433,7 @@ WHERE m.app_id = :app_id AND m.invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -494,7 +494,7 @@ class AverageResponseTimeStatistic(Resource): @account_initialization_required @get_app_model(mode=AppMode.COMPLETION) def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -510,7 +510,7 @@ WHERE app_id = :app_id AND invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -566,7 +566,7 @@ class TokensPerSecondStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -585,7 +585,7 @@ WHERE app_id = :app_id AND invoke_from != :invoke_from""" arg_dict = {"tz": account.timezone, "app_id": app_model.id, "invoke_from": InvokeFrom.DEBUGGER} - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index 578d864b80..172a80736f 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -12,7 +12,7 @@ import services from controllers.console import api, console_ns from controllers.console.app.error import ConversationCompletedError, DraftWorkflowNotExist, DraftWorkflowNotSync from controllers.console.app.wraps import get_app_model -from controllers.console.wraps import account_initialization_required, setup_required +from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError from core.app.app_config.features.file_upload.manager import FileUploadConfigManager from core.app.apps.base_app_queue_manager import AppQueueManager @@ -27,9 +27,8 @@ from fields.workflow_run_fields import workflow_run_node_execution_fields from libs import helper from libs.datetime_utils import naive_utc_now from libs.helper import TimestampField, uuid_value -from libs.login import current_user, login_required +from libs.login import current_account_with_tenant, login_required from models import App -from models.account import Account from models.model import AppMode from models.workflow import Workflow from services.app_generate_service import AppGenerateService @@ -70,15 +69,11 @@ class DraftWorkflowApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) @marshal_with(workflow_fields) + @edit_permission_required def get(self, app_model: App): """ Get draft workflow """ - # The role of the current user in the ta table must be admin, owner, or editor - assert isinstance(current_user, Account) - if not current_user.has_edit_permission: - raise Forbidden() - # fetch draft workflow by app_model workflow_service = WorkflowService() workflow = workflow_service.get_draft_workflow(app_model=app_model) @@ -110,14 +105,12 @@ class DraftWorkflowApi(Resource): @api.response(200, "Draft workflow synced successfully", workflow_fields) @api.response(400, "Invalid workflow configuration") @api.response(403, "Permission denied") + @edit_permission_required def post(self, app_model: App): """ Sync draft workflow """ - # The role of the current user in the ta table must be admin, owner, or editor - assert isinstance(current_user, Account) - if not current_user.has_edit_permission: - raise Forbidden() + current_user, _ = current_account_with_tenant() content_type = request.headers.get("Content-Type", "") @@ -149,10 +142,6 @@ class DraftWorkflowApi(Resource): return {"message": "Invalid JSON data"}, 400 else: abort(415) - - if not isinstance(current_user, Account): - raise Forbidden() - workflow_service = WorkflowService() try: @@ -206,17 +195,12 @@ class AdvancedChatDraftWorkflowRunApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT]) + @edit_permission_required def post(self, app_model: App): """ Run draft workflow """ - # The role of the current user in the ta table must be admin, owner, or editor - assert isinstance(current_user, Account) - if not current_user.has_edit_permission: - raise Forbidden() - - if not isinstance(current_user, Account): - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") @@ -271,16 +255,12 @@ class AdvancedChatDraftRunIterationNodeApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT]) + @edit_permission_required def post(self, app_model: App, node_id: str): """ Run draft workflow iteration node """ - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") args = parser.parse_args() @@ -323,16 +303,12 @@ class WorkflowDraftRunIterationNodeApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.WORKFLOW]) + @edit_permission_required def post(self, app_model: App, node_id: str): """ Run draft workflow iteration node """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account): - raise Forbidden() - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") args = parser.parse_args() @@ -375,17 +351,12 @@ class AdvancedChatDraftRunLoopNodeApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT]) + @edit_permission_required def post(self, app_model: App, node_id: str): """ Run draft workflow loop node """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") args = parser.parse_args() @@ -428,17 +399,12 @@ class WorkflowDraftRunLoopNodeApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.WORKFLOW]) + @edit_permission_required def post(self, app_model: App, node_id: str): """ Run draft workflow loop node """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") args = parser.parse_args() @@ -480,17 +446,12 @@ class DraftWorkflowRunApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.WORKFLOW]) + @edit_permission_required def post(self, app_model: App): """ Run draft workflow """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, required=True, nullable=False, location="json") parser.add_argument("files", type=list, required=False, location="json") @@ -526,17 +487,11 @@ class WorkflowTaskStopApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @edit_permission_required def post(self, app_model: App, task_id: str): """ Stop workflow task """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - # Stop using both mechanisms for backward compatibility # Legacy stop flag mechanism (without user check) AppQueueManager.set_stop_flag_no_user_check(task_id) @@ -568,17 +523,12 @@ class DraftWorkflowNodeRunApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) @marshal_with(workflow_run_node_execution_fields) + @edit_permission_required def post(self, app_model: App, node_id: str): """ Run draft workflow node """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, required=True, nullable=False, location="json") parser.add_argument("query", type=str, required=False, location="json", default="") @@ -622,17 +572,11 @@ class PublishedWorkflowApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) @marshal_with(workflow_fields) + @edit_permission_required def get(self, app_model: App): """ Get published workflow """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - # fetch published workflow by app_model workflow_service = WorkflowService() workflow = workflow_service.get_published_workflow(app_model=app_model) @@ -644,16 +588,12 @@ class PublishedWorkflowApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @edit_permission_required def post(self, app_model: App): """ Publish workflow """ - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("marked_name", type=str, required=False, default="", location="json") parser.add_argument("marked_comment", type=str, required=False, default="", location="json") @@ -702,17 +642,11 @@ class DefaultBlockConfigsApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @edit_permission_required def get(self, app_model: App): """ Get default block config """ - - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - # Get default block configs workflow_service = WorkflowService() return workflow_service.get_default_block_configs() @@ -729,16 +663,11 @@ class DefaultBlockConfigApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @edit_permission_required def get(self, app_model: App, block_type: str): """ Get default block config """ - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() - parser = reqparse.RequestParser() parser.add_argument("q", type=str, location="args") args = parser.parse_args() @@ -769,17 +698,14 @@ class ConvertToWorkflowApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.CHAT, AppMode.COMPLETION]) + @edit_permission_required def post(self, app_model: App): """ Convert basic mode of chatbot app to workflow mode Convert expert mode of chatbot app to workflow mode Convert Completion App to Workflow App """ - if not isinstance(current_user, Account): - raise Forbidden() - # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.has_edit_permission: - raise Forbidden() + current_user, _ = current_account_with_tenant() if request.data: parser = reqparse.RequestParser() @@ -812,15 +738,12 @@ class PublishedAllWorkflowApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) @marshal_with(workflow_pagination_fields) + @edit_permission_required def get(self, app_model: App): """ Get published workflows """ - - if not isinstance(current_user, Account): - raise Forbidden() - if not current_user.has_edit_permission: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("page", type=inputs.int_range(1, 99999), required=False, default=1, location="args") @@ -879,16 +802,12 @@ class WorkflowByIdApi(Resource): @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) @marshal_with(workflow_fields) + @edit_permission_required def patch(self, app_model: App, workflow_id: str): """ Update workflow attributes """ - if not isinstance(current_user, Account): - raise Forbidden() - # Check permission - if not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("marked_name", type=str, required=False, location="json") parser.add_argument("marked_comment", type=str, required=False, location="json") @@ -934,16 +853,11 @@ class WorkflowByIdApi(Resource): @login_required @account_initialization_required @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) + @edit_permission_required def delete(self, app_model: App, workflow_id: str): """ Delete workflow """ - if not isinstance(current_user, Account): - raise Forbidden() - # Check permission - if not current_user.has_edit_permission: - raise Forbidden() - workflow_service = WorkflowService() # Create a session and manage the transaction diff --git a/api/controllers/console/app/workflow_draft_variable.py b/api/controllers/console/app/workflow_draft_variable.py index da6b56d026..5e865dc4c1 100644 --- a/api/controllers/console/app/workflow_draft_variable.py +++ b/api/controllers/console/app/workflow_draft_variable.py @@ -22,8 +22,7 @@ from extensions.ext_database import db from factories.file_factory import build_from_mapping, build_from_mappings from factories.variable_factory import build_segment_with_type from libs.login import current_user, login_required -from models import App, AppMode -from models.account import Account +from models import Account, App, AppMode from models.workflow import WorkflowDraftVariable from services.workflow_draft_variable_service import WorkflowDraftVariableList, WorkflowDraftVariableService from services.workflow_service import WorkflowService diff --git a/api/controllers/console/app/workflow_run.py b/api/controllers/console/app/workflow_run.py index 23ba63845c..286ba65a7f 100644 --- a/api/controllers/console/app/workflow_run.py +++ b/api/controllers/console/app/workflow_run.py @@ -1,6 +1,5 @@ from typing import cast -from flask_login import current_user from flask_restx import Resource, marshal_with, reqparse from flask_restx.inputs import int_range @@ -14,7 +13,7 @@ from fields.workflow_run_fields import ( workflow_run_pagination_fields, ) from libs.helper import uuid_value -from libs.login import login_required +from libs.login import current_user, login_required from models import Account, App, AppMode, EndUser from services.workflow_run_service import WorkflowRunService diff --git a/api/controllers/console/app/workflow_statistic.py b/api/controllers/console/app/workflow_statistic.py index b8904bf3d9..8f7f936c9b 100644 --- a/api/controllers/console/app/workflow_statistic.py +++ b/api/controllers/console/app/workflow_statistic.py @@ -4,7 +4,6 @@ from decimal import Decimal import pytz import sqlalchemy as sa from flask import jsonify -from flask_login import current_user from flask_restx import Resource, reqparse from controllers.console import api, console_ns @@ -12,7 +11,7 @@ from controllers.console.app.wraps import get_app_model from controllers.console.wraps import account_initialization_required, setup_required from extensions.ext_database import db from libs.helper import DatetimeString -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from models.enums import WorkflowRunTriggeredFrom from models.model import AppMode @@ -29,7 +28,7 @@ class WorkflowDailyRunsStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -49,7 +48,7 @@ WHERE "app_id": app_model.id, "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, } - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -97,7 +96,7 @@ class WorkflowDailyTerminalsStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -117,7 +116,7 @@ WHERE "app_id": app_model.id, "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, } - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -165,7 +164,7 @@ class WorkflowDailyTokenCostStatistic(Resource): @login_required @account_initialization_required def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -185,7 +184,7 @@ WHERE "app_id": app_model.id, "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, } - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc @@ -238,7 +237,7 @@ class WorkflowAverageAppInteractionStatistic(Resource): @account_initialization_required @get_app_model(mode=[AppMode.WORKFLOW]) def get(self, app_model): - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("start", type=DatetimeString("%Y-%m-%d %H:%M"), location="args") @@ -271,7 +270,7 @@ GROUP BY "app_id": app_model.id, "triggered_from": WorkflowRunTriggeredFrom.APP_RUN, } - + assert account.timezone is not None timezone = pytz.timezone(account.timezone) utc_timezone = pytz.utc diff --git a/api/controllers/console/app/wraps.py b/api/controllers/console/app/wraps.py index 44aba01820..9bb2718f89 100644 --- a/api/controllers/console/app/wraps.py +++ b/api/controllers/console/app/wraps.py @@ -4,28 +4,29 @@ from typing import ParamSpec, TypeVar, Union from controllers.console.app.error import AppNotFoundError from extensions.ext_database import db -from libs.login import current_user +from libs.login import current_account_with_tenant from models import App, AppMode -from models.account import Account P = ParamSpec("P") R = TypeVar("R") +P1 = ParamSpec("P1") +R1 = TypeVar("R1") def _load_app_model(app_id: str) -> App | None: - assert isinstance(current_user, Account) + _, current_tenant_id = current_account_with_tenant() app_model = ( db.session.query(App) - .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") + .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal") .first() ) return app_model def get_app_model(view: Callable[P, R] | None = None, *, mode: Union[AppMode, list[AppMode], None] = None): - def decorator(view_func: Callable[P, R]): + def decorator(view_func: Callable[P1, R1]): @wraps(view_func) - def decorated_view(*args: P.args, **kwargs: P.kwargs): + def decorated_view(*args: P1.args, **kwargs: P1.kwargs): if not kwargs.get("app_id"): raise ValueError("missing app_id in path parameters") diff --git a/api/controllers/console/auth/activate.py b/api/controllers/console/auth/activate.py index 76171e3f8a..06d2b936b7 100644 --- a/api/controllers/console/auth/activate.py +++ b/api/controllers/console/auth/activate.py @@ -7,7 +7,7 @@ from controllers.console.error import AlreadyActivateError from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from libs.helper import StrLen, email, extract_remote_ip, timezone -from models.account import AccountStatus +from models import AccountStatus from services.account_service import AccountService, RegisterService active_check_parser = reqparse.RequestParser() diff --git a/api/controllers/console/auth/data_source_oauth.py b/api/controllers/console/auth/data_source_oauth.py index 6f1fd2f11a..0fd433d718 100644 --- a/api/controllers/console/auth/data_source_oauth.py +++ b/api/controllers/console/auth/data_source_oauth.py @@ -2,13 +2,12 @@ import logging import httpx from flask import current_app, redirect, request -from flask_login import current_user from flask_restx import Resource, fields from werkzeug.exceptions import Forbidden from configs import dify_config from controllers.console import api, console_ns -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from libs.oauth_data_source import NotionOAuth from ..wraps import account_initialization_required, setup_required @@ -45,6 +44,7 @@ class OAuthDataSource(Resource): @api.response(403, "Admin privileges required") def get(self, provider: str): # The role of the current user in the table must be admin or owner + current_user, _ = current_account_with_tenant() if not current_user.is_admin_or_owner: raise Forbidden() OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers() diff --git a/api/controllers/console/auth/email_register.py b/api/controllers/console/auth/email_register.py index d3613d9183..cabd118d23 100644 --- a/api/controllers/console/auth/email_register.py +++ b/api/controllers/console/auth/email_register.py @@ -19,7 +19,7 @@ from controllers.console.wraps import email_password_login_enabled, email_regist from extensions.ext_database import db from libs.helper import email, extract_remote_ip from libs.password import valid_password -from models.account import Account +from models import Account from services.account_service import AccountService from services.billing_service import BillingService from services.errors.account import AccountNotFoundError, AccountRegisterError diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index 704bcf8fb8..102d33966e 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/api/controllers/console/auth/forgot_password.py @@ -20,7 +20,7 @@ from events.tenant_event import tenant_was_created from extensions.ext_database import db from libs.helper import email, extract_remote_ip from libs.password import hash_password, valid_password -from models.account import Account +from models import Account from services.account_service import AccountService, TenantService from services.feature_service import FeatureService diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index ba614aa828..e4bbbf107b 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -1,5 +1,3 @@ -from typing import cast - import flask_login from flask import request from flask_restx import Resource, reqparse @@ -26,7 +24,7 @@ from controllers.console.error import ( from controllers.console.wraps import email_password_login_enabled, setup_required from events.tenant_event import tenant_was_created from libs.helper import email, extract_remote_ip -from models.account import Account +from libs.login import current_account_with_tenant from services.account_service import AccountService, RegisterService, TenantService from services.billing_service import BillingService from services.errors.account import AccountRegisterError @@ -96,7 +94,8 @@ class LoginApi(Resource): class LogoutApi(Resource): @setup_required def get(self): - account = cast(Account, flask_login.current_user) + current_user, _ = current_account_with_tenant() + account = current_user if isinstance(account, flask_login.AnonymousUserMixin): return {"result": "success"} AccountService.logout(account=account) diff --git a/api/controllers/console/auth/oauth.py b/api/controllers/console/auth/oauth.py index 4efeceb676..52459ad5eb 100644 --- a/api/controllers/console/auth/oauth.py +++ b/api/controllers/console/auth/oauth.py @@ -14,8 +14,7 @@ from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from libs.helper import extract_remote_ip from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo -from models import Account -from models.account import AccountStatus +from models import Account, AccountStatus from services.account_service import AccountService, RegisterService, TenantService from services.billing_service import BillingService from services.errors.account import AccountNotFoundError, AccountRegisterError diff --git a/api/controllers/console/auth/oauth_server.py b/api/controllers/console/auth/oauth_server.py index 46281860ae..188ef7f622 100644 --- a/api/controllers/console/auth/oauth_server.py +++ b/api/controllers/console/auth/oauth_server.py @@ -1,16 +1,15 @@ from collections.abc import Callable from functools import wraps -from typing import Concatenate, ParamSpec, TypeVar, cast +from typing import Concatenate, ParamSpec, TypeVar -import flask_login from flask import jsonify, request from flask_restx import Resource, reqparse from werkzeug.exceptions import BadRequest, NotFound from controllers.console.wraps import account_initialization_required, setup_required from core.model_runtime.utils.encoders import jsonable_encoder -from libs.login import login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required +from models import Account from models.model import OAuthProviderApp from services.oauth_server import OAUTH_ACCESS_TOKEN_EXPIRES_IN, OAuthGrantType, OAuthServerService @@ -116,7 +115,8 @@ class OAuthServerUserAuthorizeApi(Resource): @account_initialization_required @oauth_server_client_id_required def post(self, oauth_provider_app: OAuthProviderApp): - account = cast(Account, flask_login.current_user) + current_user, _ = current_account_with_tenant() + account = current_user user_account_id = account.id code = OAuthServerService.sign_oauth_authorization_code(oauth_provider_app.client_id, user_account_id) diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index fa89f45122..5c89b29057 100644 --- a/api/controllers/console/billing/billing.py +++ b/api/controllers/console/billing/billing.py @@ -2,8 +2,7 @@ from flask_restx import Resource, reqparse from controllers.console import console_ns from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required -from libs.login import current_user, login_required -from models.model import Account +from libs.login import current_account_with_tenant, login_required from services.billing_service import BillingService @@ -14,17 +13,13 @@ class Subscription(Resource): @account_initialization_required @only_edition_cloud def get(self): + current_user, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("plan", type=str, required=True, location="args", choices=["professional", "team"]) parser.add_argument("interval", type=str, required=True, location="args", choices=["month", "year"]) args = parser.parse_args() - assert isinstance(current_user, Account) - BillingService.is_tenant_owner_or_admin(current_user) - assert current_user.current_tenant_id is not None - return BillingService.get_subscription( - args["plan"], args["interval"], current_user.email, current_user.current_tenant_id - ) + return BillingService.get_subscription(args["plan"], args["interval"], current_user.email, current_tenant_id) @console_ns.route("/billing/invoices") @@ -34,7 +29,6 @@ class Invoices(Resource): @account_initialization_required @only_edition_cloud def get(self): - assert isinstance(current_user, Account) + current_user, current_tenant_id = current_account_with_tenant() BillingService.is_tenant_owner_or_admin(current_user) - assert current_user.current_tenant_id is not None - return BillingService.get_invoices(current_user.email, current_user.current_tenant_id) + return BillingService.get_invoices(current_user.email, current_tenant_id) diff --git a/api/controllers/console/billing/compliance.py b/api/controllers/console/billing/compliance.py index c0d104e0d4..3b32fe29a1 100644 --- a/api/controllers/console/billing/compliance.py +++ b/api/controllers/console/billing/compliance.py @@ -2,8 +2,7 @@ from flask import request from flask_restx import Resource, reqparse from libs.helper import extract_remote_ip -from libs.login import current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from services.billing_service import BillingService from .. import console_ns @@ -17,19 +16,17 @@ class ComplianceApi(Resource): @account_initialization_required @only_edition_cloud def get(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("doc_name", type=str, required=True, location="args") args = parser.parse_args() ip_address = extract_remote_ip(request) device_info = request.headers.get("User-Agent", "Unknown device") - return BillingService.get_compliance_download_link( doc_name=args.doc_name, account_id=current_user.id, - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, ip=ip_address, device_info=device_info, ) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index f86c5dfc3c..c03767d2e6 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -1,7 +1,6 @@ from typing import Any, cast from flask import request -from flask_login import current_user from flask_restx import Resource, fields, marshal, marshal_with, reqparse from sqlalchemy import select from werkzeug.exceptions import Forbidden, NotFound @@ -30,10 +29,9 @@ from extensions.ext_database import db from fields.app_fields import related_app_list from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields from fields.document_fields import document_status_fields -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from libs.validators import validate_description_length from models import ApiToken, Dataset, Document, DocumentSegment, UploadFile -from models.account import Account from models.dataset import DatasetPermissionEnum from models.provider_ids import ModelProviderID from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService @@ -138,6 +136,7 @@ class DatasetListApi(Resource): @account_initialization_required @enterprise_license_required def get(self): + current_user, current_tenant_id = current_account_with_tenant() page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) ids = request.args.getlist("ids") @@ -146,15 +145,15 @@ class DatasetListApi(Resource): tag_ids = request.args.getlist("tag_ids") include_all = request.args.get("include_all", default="false").lower() == "true" if ids: - datasets, total = DatasetService.get_datasets_by_ids(ids, current_user.current_tenant_id) + datasets, total = DatasetService.get_datasets_by_ids(ids, current_tenant_id) else: datasets, total = DatasetService.get_datasets( - page, limit, current_user.current_tenant_id, current_user, search, tag_ids, include_all + page, limit, current_tenant_id, current_user, search, tag_ids, include_all ) # check embedding setting provider_manager = ProviderManager() - configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id) + configurations = provider_manager.get_configurations(tenant_id=current_tenant_id) embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True) @@ -251,6 +250,7 @@ class DatasetListApi(Resource): required=False, ) args = parser.parse_args() + current_user, current_tenant_id = current_account_with_tenant() # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator if not current_user.is_dataset_editor: @@ -258,11 +258,11 @@ class DatasetListApi(Resource): try: dataset = DatasetService.create_empty_dataset( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, name=args["name"], description=args["description"], indexing_technique=args["indexing_technique"], - account=cast(Account, current_user), + account=current_user, permission=DatasetPermissionEnum.ONLY_ME, provider=args["provider"], external_knowledge_api_id=args["external_knowledge_api_id"], @@ -286,6 +286,7 @@ class DatasetApi(Resource): @login_required @account_initialization_required def get(self, dataset_id): + current_user, current_tenant_id = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: @@ -305,7 +306,7 @@ class DatasetApi(Resource): # check embedding setting provider_manager = ProviderManager() - configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id) + configurations = provider_manager.get_configurations(tenant_id=current_tenant_id) embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True) @@ -418,6 +419,7 @@ class DatasetApi(Resource): ) args = parser.parse_args() data = request.get_json() + current_user, current_tenant_id = current_account_with_tenant() # check embedding model setting if ( @@ -440,7 +442,7 @@ class DatasetApi(Resource): raise NotFound("Dataset not found.") result_data = cast(dict[str, Any], marshal(dataset, dataset_detail_fields)) - tenant_id = current_user.current_tenant_id + tenant_id = current_tenant_id if data.get("partial_member_list") and data.get("permission") == "partial_members": DatasetPermissionService.update_partial_member_list( @@ -464,9 +466,10 @@ class DatasetApi(Resource): @cloud_edition_billing_rate_limit_check("knowledge") def delete(self, dataset_id): dataset_id_str = str(dataset_id) + current_user, _ = current_account_with_tenant() # The role of the current user in the ta table must be admin, owner, or editor - if not (current_user.is_editor or current_user.is_dataset_operator): + if not (current_user.has_edit_permission or current_user.is_dataset_operator): raise Forbidden() try: @@ -505,6 +508,7 @@ class DatasetQueryApi(Resource): @login_required @account_initialization_required def get(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: @@ -556,15 +560,14 @@ class DatasetIndexingEstimateApi(Resource): "doc_language", type=str, default="English", required=False, nullable=False, location="json" ) args = parser.parse_args() + _, current_tenant_id = current_account_with_tenant() # validate args DocumentService.estimate_args_validate(args) extract_settings = [] if args["info_list"]["data_source_type"] == "upload_file": file_ids = args["info_list"]["file_info_list"]["file_ids"] file_details = db.session.scalars( - select(UploadFile).where( - UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids) - ) + select(UploadFile).where(UploadFile.tenant_id == current_tenant_id, UploadFile.id.in_(file_ids)) ).all() if file_details is None: @@ -592,7 +595,7 @@ class DatasetIndexingEstimateApi(Resource): "notion_workspace_id": workspace_id, "notion_obj_id": page["page_id"], "notion_page_type": page["type"], - "tenant_id": current_user.current_tenant_id, + "tenant_id": current_tenant_id, } ), document_model=args["doc_form"], @@ -608,7 +611,7 @@ class DatasetIndexingEstimateApi(Resource): "provider": website_info_list["provider"], "job_id": website_info_list["job_id"], "url": url, - "tenant_id": current_user.current_tenant_id, + "tenant_id": current_tenant_id, "mode": "crawl", "only_main_content": website_info_list["only_main_content"], } @@ -621,7 +624,7 @@ class DatasetIndexingEstimateApi(Resource): indexing_runner = IndexingRunner() try: response = indexing_runner.indexing_estimate( - current_user.current_tenant_id, + current_tenant_id, extract_settings, args["process_rule"], args["doc_form"], @@ -652,6 +655,7 @@ class DatasetRelatedAppListApi(Resource): @account_initialization_required @marshal_with(related_app_list) def get(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: @@ -683,11 +687,10 @@ class DatasetIndexingStatusApi(Resource): @login_required @account_initialization_required def get(self, dataset_id): + _, current_tenant_id = current_account_with_tenant() dataset_id = str(dataset_id) documents = db.session.scalars( - select(Document).where( - Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id - ) + select(Document).where(Document.dataset_id == dataset_id, Document.tenant_id == current_tenant_id) ).all() documents_status = [] for document in documents: @@ -739,10 +742,9 @@ class DatasetApiKeyApi(Resource): @account_initialization_required @marshal_with(api_key_list) def get(self): + _, current_tenant_id = current_account_with_tenant() keys = db.session.scalars( - select(ApiToken).where( - ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id - ) + select(ApiToken).where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id) ).all() return {"items": keys} @@ -752,12 +754,13 @@ class DatasetApiKeyApi(Resource): @marshal_with(api_key_fields) def post(self): # The role of the current user in the ta table must be admin or owner + current_user, current_tenant_id = current_account_with_tenant() if not current_user.is_admin_or_owner: raise Forbidden() current_key_count = ( db.session.query(ApiToken) - .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) + .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id) .count() ) @@ -770,7 +773,7 @@ class DatasetApiKeyApi(Resource): key = ApiToken.generate_api_key(self.token_prefix, 24) api_token = ApiToken() - api_token.tenant_id = current_user.current_tenant_id + api_token.tenant_id = current_tenant_id api_token.token = key api_token.type = self.resource_type db.session.add(api_token) @@ -790,6 +793,7 @@ class DatasetApiDeleteApi(Resource): @login_required @account_initialization_required def delete(self, api_key_id): + current_user, current_tenant_id = current_account_with_tenant() api_key_id = str(api_key_id) # The role of the current user in the ta table must be admin or owner @@ -799,7 +803,7 @@ class DatasetApiDeleteApi(Resource): key = ( db.session.query(ApiToken) .where( - ApiToken.tenant_id == current_user.current_tenant_id, + ApiToken.tenant_id == current_tenant_id, ApiToken.type == self.resource_type, ApiToken.id == api_key_id, ) @@ -898,6 +902,7 @@ class DatasetPermissionUserListApi(Resource): @login_required @account_initialization_required def get(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: diff --git a/api/controllers/console/datasets/datasets_document.py b/api/controllers/console/datasets/datasets_document.py index 011dacde76..9c0c54833e 100644 --- a/api/controllers/console/datasets/datasets_document.py +++ b/api/controllers/console/datasets/datasets_document.py @@ -6,7 +6,6 @@ from typing import Literal, cast import sqlalchemy as sa from flask import request -from flask_login import current_user from flask_restx import Resource, fields, marshal, marshal_with, reqparse from sqlalchemy import asc, desc, select from werkzeug.exceptions import Forbidden, NotFound @@ -53,9 +52,8 @@ from fields.document_fields import ( document_with_segments_fields, ) from libs.datetime_utils import naive_utc_now -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from models import Dataset, DatasetProcessRule, Document, DocumentSegment, UploadFile -from models.account import Account from models.dataset import DocumentPipelineExecutionLog from services.dataset_service import DatasetService, DocumentService from services.entities.knowledge_entities.knowledge_entities import KnowledgeConfig @@ -65,6 +63,7 @@ logger = logging.getLogger(__name__) class DocumentResource(Resource): def get_document(self, dataset_id: str, document_id: str) -> Document: + current_user, current_tenant_id = current_account_with_tenant() dataset = DatasetService.get_dataset(dataset_id) if not dataset: raise NotFound("Dataset not found.") @@ -79,12 +78,13 @@ class DocumentResource(Resource): if not document: raise NotFound("Document not found.") - if document.tenant_id != current_user.current_tenant_id: + if document.tenant_id != current_tenant_id: raise Forbidden("No permission.") return document def get_batch_documents(self, dataset_id: str, batch: str) -> Sequence[Document]: + current_user, _ = current_account_with_tenant() dataset = DatasetService.get_dataset(dataset_id) if not dataset: raise NotFound("Dataset not found.") @@ -112,6 +112,7 @@ class GetProcessRuleApi(Resource): @login_required @account_initialization_required def get(self): + current_user, _ = current_account_with_tenant() req_data = request.args document_id = req_data.get("document_id") @@ -168,6 +169,7 @@ class DatasetDocumentListApi(Resource): @login_required @account_initialization_required def get(self, dataset_id): + current_user, current_tenant_id = current_account_with_tenant() dataset_id = str(dataset_id) page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) @@ -199,7 +201,7 @@ class DatasetDocumentListApi(Resource): except services.errors.account.NoPermissionError as e: raise Forbidden(str(e)) - query = select(Document).filter_by(dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id) + query = select(Document).filter_by(dataset_id=str(dataset_id), tenant_id=current_tenant_id) if search: search = f"%{search}%" @@ -273,6 +275,7 @@ class DatasetDocumentListApi(Resource): @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") def post(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id) @@ -372,6 +375,7 @@ class DatasetInitApi(Resource): @cloud_edition_billing_rate_limit_check("knowledge") def post(self): # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor + current_user, current_tenant_id = current_account_with_tenant() if not current_user.is_dataset_editor: raise Forbidden() @@ -402,7 +406,7 @@ class DatasetInitApi(Resource): try: model_manager = ModelManager() model_manager.get_model_instance( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, provider=args["embedding_model_provider"], model_type=ModelType.TEXT_EMBEDDING, model=args["embedding_model"], @@ -419,9 +423,9 @@ class DatasetInitApi(Resource): try: dataset, documents, batch = DocumentService.save_document_without_dataset_id( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, knowledge_config=knowledge_config, - account=cast(Account, current_user), + account=current_user, ) except ProviderTokenNotInitError as ex: raise ProviderNotInitializeError(ex.description) @@ -447,6 +451,7 @@ class DocumentIndexingEstimateApi(DocumentResource): @login_required @account_initialization_required def get(self, dataset_id, document_id): + _, current_tenant_id = current_account_with_tenant() dataset_id = str(dataset_id) document_id = str(document_id) document = self.get_document(dataset_id, document_id) @@ -482,7 +487,7 @@ class DocumentIndexingEstimateApi(DocumentResource): try: estimate_response = indexing_runner.indexing_estimate( - current_user.current_tenant_id, + current_tenant_id, [extract_setting], data_process_rule_dict, document.doc_form, @@ -511,6 +516,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource): @login_required @account_initialization_required def get(self, dataset_id, batch): + _, current_tenant_id = current_account_with_tenant() dataset_id = str(dataset_id) batch = str(batch) documents = self.get_batch_documents(dataset_id, batch) @@ -530,7 +536,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource): file_id = data_source_info["upload_file_id"] file_detail = ( db.session.query(UploadFile) - .where(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id == file_id) + .where(UploadFile.tenant_id == current_tenant_id, UploadFile.id == file_id) .first() ) @@ -553,7 +559,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource): "notion_workspace_id": data_source_info["notion_workspace_id"], "notion_obj_id": data_source_info["notion_page_id"], "notion_page_type": data_source_info["type"], - "tenant_id": current_user.current_tenant_id, + "tenant_id": current_tenant_id, } ), document_model=document.doc_form, @@ -569,7 +575,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource): "provider": data_source_info["provider"], "job_id": data_source_info["job_id"], "url": data_source_info["url"], - "tenant_id": current_user.current_tenant_id, + "tenant_id": current_tenant_id, "mode": data_source_info["mode"], "only_main_content": data_source_info["only_main_content"], } @@ -583,7 +589,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource): indexing_runner = IndexingRunner() try: response = indexing_runner.indexing_estimate( - current_user.current_tenant_id, + current_tenant_id, extract_settings, data_process_rule_dict, document.doc_form, @@ -834,6 +840,7 @@ class DocumentProcessingApi(DocumentResource): @account_initialization_required @cloud_edition_billing_rate_limit_check("knowledge") def patch(self, dataset_id, document_id, action: Literal["pause", "resume"]): + current_user, _ = current_account_with_tenant() dataset_id = str(dataset_id) document_id = str(document_id) document = self.get_document(dataset_id, document_id) @@ -884,6 +891,7 @@ class DocumentMetadataApi(DocumentResource): @login_required @account_initialization_required def put(self, dataset_id, document_id): + current_user, _ = current_account_with_tenant() dataset_id = str(dataset_id) document_id = str(document_id) document = self.get_document(dataset_id, document_id) @@ -931,6 +939,7 @@ class DocumentStatusApi(DocumentResource): @cloud_edition_billing_resource_check("vector_space") @cloud_edition_billing_rate_limit_check("knowledge") def patch(self, dataset_id, action: Literal["enable", "disable", "archive", "un_archive"]): + current_user, _ = current_account_with_tenant() dataset_id = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id) if dataset is None: @@ -1077,12 +1086,13 @@ class DocumentRenameApi(DocumentResource): @marshal_with(document_fields) def post(self, dataset_id, document_id): # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator + current_user, _ = current_account_with_tenant() if not current_user.is_dataset_editor: raise Forbidden() dataset = DatasetService.get_dataset(dataset_id) if not dataset: raise NotFound("Dataset not found.") - DatasetService.check_dataset_operator_permission(cast(Account, current_user), dataset) + DatasetService.check_dataset_operator_permission(current_user, dataset) parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, nullable=False, location="json") args = parser.parse_args() @@ -1102,6 +1112,7 @@ class WebsiteDocumentSyncApi(DocumentResource): @account_initialization_required def get(self, dataset_id, document_id): """sync website document.""" + _, current_tenant_id = current_account_with_tenant() dataset_id = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id) if not dataset: @@ -1110,7 +1121,7 @@ class WebsiteDocumentSyncApi(DocumentResource): document = DocumentService.get_document(dataset.id, document_id) if not document: raise NotFound("Document not found.") - if document.tenant_id != current_user.current_tenant_id: + if document.tenant_id != current_tenant_id: raise Forbidden("No permission.") if document.data_source_type != "website_crawl": raise ValueError("Document is not a website document.") diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index adf9f53523..f590919180 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -1,7 +1,4 @@ -from typing import cast - from flask import request -from flask_login import current_user from flask_restx import Resource, fields, marshal, reqparse from werkzeug.exceptions import Forbidden, InternalServerError, NotFound @@ -10,8 +7,7 @@ from controllers.console import api, console_ns from controllers.console.datasets.error import DatasetNameDuplicateError from controllers.console.wraps import account_initialization_required, setup_required from fields.dataset_fields import dataset_detail_fields -from libs.login import login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from services.dataset_service import DatasetService from services.external_knowledge_service import ExternalDatasetService from services.hit_testing_service import HitTestingService @@ -40,12 +36,13 @@ class ExternalApiTemplateListApi(Resource): @login_required @account_initialization_required def get(self): + _, current_tenant_id = current_account_with_tenant() page = request.args.get("page", default=1, type=int) limit = request.args.get("limit", default=20, type=int) search = request.args.get("keyword", default=None, type=str) external_knowledge_apis, total = ExternalDatasetService.get_external_knowledge_apis( - page, limit, current_user.current_tenant_id, search + page, limit, current_tenant_id, search ) response = { "data": [item.to_dict() for item in external_knowledge_apis], @@ -60,6 +57,7 @@ class ExternalApiTemplateListApi(Resource): @login_required @account_initialization_required def post(self): + current_user, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument( "name", @@ -85,7 +83,7 @@ class ExternalApiTemplateListApi(Resource): try: external_knowledge_api = ExternalDatasetService.create_external_knowledge_api( - tenant_id=current_user.current_tenant_id, user_id=current_user.id, args=args + tenant_id=current_tenant_id, user_id=current_user.id, args=args ) except services.errors.dataset.DatasetNameDuplicateError: raise DatasetNameDuplicateError() @@ -115,6 +113,7 @@ class ExternalApiTemplateApi(Resource): @login_required @account_initialization_required def patch(self, external_knowledge_api_id): + current_user, current_tenant_id = current_account_with_tenant() external_knowledge_api_id = str(external_knowledge_api_id) parser = reqparse.RequestParser() @@ -136,7 +135,7 @@ class ExternalApiTemplateApi(Resource): ExternalDatasetService.validate_api_list(args["settings"]) external_knowledge_api = ExternalDatasetService.update_external_knowledge_api( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, user_id=current_user.id, external_knowledge_api_id=external_knowledge_api_id, args=args, @@ -148,13 +147,14 @@ class ExternalApiTemplateApi(Resource): @login_required @account_initialization_required def delete(self, external_knowledge_api_id): + current_user, current_tenant_id = current_account_with_tenant() external_knowledge_api_id = str(external_knowledge_api_id) # The role of the current user in the ta table must be admin, owner, or editor - if not (current_user.is_editor or current_user.is_dataset_operator): + if not (current_user.has_edit_permission or current_user.is_dataset_operator): raise Forbidden() - ExternalDatasetService.delete_external_knowledge_api(current_user.current_tenant_id, external_knowledge_api_id) + ExternalDatasetService.delete_external_knowledge_api(current_tenant_id, external_knowledge_api_id) return {"result": "success"}, 204 @@ -199,7 +199,8 @@ class ExternalDatasetCreateApi(Resource): @account_initialization_required def post(self): # The role of the current user in the ta table must be admin, owner, or editor - if not current_user.is_editor: + current_user, current_tenant_id = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -223,7 +224,7 @@ class ExternalDatasetCreateApi(Resource): try: dataset = ExternalDatasetService.create_external_dataset( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, user_id=current_user.id, args=args, ) @@ -255,6 +256,7 @@ class ExternalKnowledgeHitTestingApi(Resource): @login_required @account_initialization_required def post(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: @@ -277,7 +279,7 @@ class ExternalKnowledgeHitTestingApi(Resource): response = HitTestingService.external_retrieve( dataset=dataset, query=args["query"], - account=cast(Account, current_user), + account=current_user, external_retrieval_model=args["external_retrieval_model"], metadata_filtering_conditions=args["metadata_filtering_conditions"], ) diff --git a/api/controllers/console/datasets/metadata.py b/api/controllers/console/datasets/metadata.py index 8438458617..673bac1add 100644 --- a/api/controllers/console/datasets/metadata.py +++ b/api/controllers/console/datasets/metadata.py @@ -1,13 +1,12 @@ from typing import Literal -from flask_login import current_user from flask_restx import Resource, marshal_with, reqparse from werkzeug.exceptions import NotFound from controllers.console import console_ns from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required from fields.dataset_fields import dataset_metadata_fields -from libs.login import login_required +from libs.login import current_account_with_tenant, login_required from services.dataset_service import DatasetService from services.entities.knowledge_entities.knowledge_entities import ( MetadataArgs, @@ -24,6 +23,7 @@ class DatasetMetadataCreateApi(Resource): @enterprise_license_required @marshal_with(dataset_metadata_fields) def post(self, dataset_id): + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("type", type=str, required=True, nullable=False, location="json") parser.add_argument("name", type=str, required=True, nullable=False, location="json") @@ -59,6 +59,7 @@ class DatasetMetadataApi(Resource): @enterprise_license_required @marshal_with(dataset_metadata_fields) def patch(self, dataset_id, metadata_id): + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, nullable=False, location="json") args = parser.parse_args() @@ -79,6 +80,7 @@ class DatasetMetadataApi(Resource): @account_initialization_required @enterprise_license_required def delete(self, dataset_id, metadata_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) metadata_id_str = str(metadata_id) dataset = DatasetService.get_dataset(dataset_id_str) @@ -108,6 +110,7 @@ class DatasetMetadataBuiltInFieldActionApi(Resource): @account_initialization_required @enterprise_license_required def post(self, dataset_id, action: Literal["enable", "disable"]): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: @@ -128,6 +131,7 @@ class DocumentMetadataEditApi(Resource): @account_initialization_required @enterprise_license_required def post(self, dataset_id): + current_user, _ = current_account_with_tenant() dataset_id_str = str(dataset_id) dataset = DatasetService.get_dataset(dataset_id_str) if dataset is None: diff --git a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py index 4c34df7eff..194bd98fa3 100644 --- a/api/controllers/console/datasets/rag_pipeline/datasource_auth.py +++ b/api/controllers/console/datasets/rag_pipeline/datasource_auth.py @@ -4,10 +4,7 @@ from werkzeug.exceptions import Forbidden, NotFound from configs import dify_config from controllers.console import console_ns -from controllers.console.wraps import ( - account_initialization_required, - setup_required, -) +from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from core.model_runtime.errors.validate import CredentialsValidateFailedError from core.model_runtime.utils.encoders import jsonable_encoder from core.plugin.impl.oauth import OAuthHandler @@ -23,12 +20,11 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def get(self, provider_id: str): current_user, current_tenant_id = current_account_with_tenant() tenant_id = current_tenant_id - if not current_user.has_edit_permission: - raise Forbidden() credential_id = request.args.get("credential_id") datasource_provider_id = DatasourceProviderID(provider_id) @@ -130,11 +126,9 @@ class DatasourceAuth(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() - - if not current_user.has_edit_permission: - raise Forbidden() + _, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument( @@ -177,14 +171,14 @@ class DatasourceAuthDeleteApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() datasource_provider_id = DatasourceProviderID(provider_id) plugin_id = datasource_provider_id.plugin_id provider_name = datasource_provider_id.provider_name - if not current_user.has_edit_permission: - raise Forbidden() + parser = reqparse.RequestParser() parser.add_argument("credential_id", type=str, required=True, nullable=False, location="json") args = parser.parse_args() @@ -203,8 +197,9 @@ class DatasourceAuthUpdateApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() datasource_provider_id = DatasourceProviderID(provider_id) parser = reqparse.RequestParser() @@ -212,8 +207,7 @@ class DatasourceAuthUpdateApi(Resource): parser.add_argument("name", type=StrLen(max_length=100), required=False, nullable=True, location="json") parser.add_argument("credential_id", type=str, required=True, nullable=False, location="json") args = parser.parse_args() - if not current_user.has_edit_permission: - raise Forbidden() + datasource_provider_service = DatasourceProviderService() datasource_provider_service.update_datasource_credentials( tenant_id=current_tenant_id, @@ -257,11 +251,10 @@ class DatasourceAuthOauthCustomClient(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("client_params", type=dict, required=False, nullable=True, location="json") parser.add_argument("enable_oauth_custom_client", type=bool, required=False, nullable=True, location="json") @@ -296,11 +289,10 @@ class DatasourceAuthDefaultApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("id", type=str, required=True, nullable=False, location="json") args = parser.parse_args() @@ -319,11 +311,10 @@ class DatasourceUpdateProviderNameApi(Resource): @setup_required @login_required @account_initialization_required + @edit_permission_required def post(self, provider_id: str): - current_user, current_tenant_id = current_account_with_tenant() + _, current_tenant_id = current_account_with_tenant() - if not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("name", type=StrLen(max_length=100), required=True, nullable=False, location="json") parser.add_argument("credential_id", type=str, required=True, nullable=False, location="json") diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py index bef6bfd13e..2e8cc16dc1 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py @@ -23,7 +23,7 @@ from extensions.ext_database import db from factories.file_factory import build_from_mapping, build_from_mappings from factories.variable_factory import build_segment_with_type from libs.login import current_user, login_required -from models.account import Account +from models import Account from models.dataset import Pipeline from models.workflow import WorkflowDraftVariable from services.rag_pipeline.rag_pipeline import RagPipelineService diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_import.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_import.py index a82872ba2b..ca767dbb10 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_import.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_import.py @@ -1,6 +1,3 @@ -from typing import cast - -from flask_login import current_user # type: ignore from flask_restx import Resource, marshal_with, reqparse # type: ignore from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden @@ -13,8 +10,7 @@ from controllers.console.wraps import ( ) from extensions.ext_database import db from fields.rag_pipeline_fields import pipeline_import_check_dependencies_fields, pipeline_import_fields -from libs.login import login_required -from models import Account +from libs.login import current_account_with_tenant, login_required from models.dataset import Pipeline from services.app_dsl_service import ImportStatus from services.rag_pipeline.rag_pipeline_dsl_service import RagPipelineDslService @@ -28,7 +24,8 @@ class RagPipelineImportApi(Resource): @marshal_with(pipeline_import_fields) def post(self): # Check user role first - if not current_user.is_editor: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -47,7 +44,7 @@ class RagPipelineImportApi(Resource): with Session(db.engine) as session: import_service = RagPipelineDslService(session) # Import app - account = cast(Account, current_user) + account = current_user result = import_service.import_rag_pipeline( account=account, import_mode=args["mode"], @@ -74,15 +71,16 @@ class RagPipelineImportConfirmApi(Resource): @account_initialization_required @marshal_with(pipeline_import_fields) def post(self, import_id): + current_user, _ = current_account_with_tenant() # Check user role first - if not current_user.is_editor: + if not current_user.has_edit_permission: raise Forbidden() # Create service with session with Session(db.engine) as session: import_service = RagPipelineDslService(session) # Confirm import - account = cast(Account, current_user) + account = current_user result = import_service.confirm_import(import_id=import_id, account=account) session.commit() @@ -100,7 +98,8 @@ class RagPipelineImportCheckDependenciesApi(Resource): @account_initialization_required @marshal_with(pipeline_import_check_dependencies_fields) def get(self, pipeline: Pipeline): - if not current_user.is_editor: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() with Session(db.engine) as session: @@ -117,7 +116,8 @@ class RagPipelineExportApi(Resource): @get_rag_pipeline @account_initialization_required def get(self, pipeline: Pipeline): - if not current_user.is_editor: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() # Add include_secret params diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py index a75c121fbe..d4d6da7fe2 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py @@ -18,6 +18,7 @@ from controllers.console.app.error import ( from controllers.console.datasets.wraps import get_rag_pipeline from controllers.console.wraps import ( account_initialization_required, + edit_permission_required, setup_required, ) from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError @@ -36,8 +37,8 @@ from fields.workflow_run_fields import ( ) from libs import helper from libs.helper import TimestampField, uuid_value -from libs.login import current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, current_user, login_required +from models import Account from models.dataset import Pipeline from models.model import EndUser from services.errors.app import WorkflowHashNotEqualError @@ -56,15 +57,12 @@ class DraftRagPipelineApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required @marshal_with(workflow_fields) def get(self, pipeline: Pipeline): """ Get draft rag pipeline's workflow """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() - # fetch draft workflow by app_model rag_pipeline_service = RagPipelineService() workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) @@ -79,13 +77,13 @@ class DraftRagPipelineApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def post(self, pipeline: Pipeline): """ Sync draft workflow """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() + current_user, _ = current_account_with_tenant() content_type = request.headers.get("Content-Type", "") @@ -154,13 +152,13 @@ class RagPipelineDraftRunIterationNodeApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def post(self, pipeline: Pipeline, node_id: str): """ Run draft workflow iteration node """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("inputs", type=dict, location="json") @@ -194,7 +192,8 @@ class RagPipelineDraftRunLoopNodeApi(Resource): Run draft workflow loop node """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -229,7 +228,8 @@ class DraftRagPipelineRunApi(Resource): Run draft workflow """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -264,7 +264,8 @@ class PublishedRagPipelineRunApi(Resource): Run published workflow """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -303,7 +304,7 @@ class PublishedRagPipelineRunApi(Resource): # Run rag pipeline datasource # """ # # The role of the current user in the ta table must be admin, owner, or editor -# if not current_user.is_editor: +# if not current_user.has_edit_permission: # raise Forbidden() # # if not isinstance(current_user, Account): @@ -344,7 +345,7 @@ class PublishedRagPipelineRunApi(Resource): # Run rag pipeline datasource # """ # # The role of the current user in the ta table must be admin, owner, or editor -# if not current_user.is_editor: +# if not current_user.has_edit_permission: # raise Forbidden() # # if not isinstance(current_user, Account): @@ -385,7 +386,8 @@ class RagPipelinePublishedDatasourceNodeRunApi(Resource): Run rag pipeline datasource """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -428,7 +430,8 @@ class RagPipelineDraftDatasourceNodeRunApi(Resource): Run rag pipeline datasource """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -472,7 +475,8 @@ class RagPipelineDraftNodeRunApi(Resource): Run draft workflow node """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -505,7 +509,8 @@ class RagPipelineTaskStopApi(Resource): Stop workflow task """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() AppQueueManager.set_stop_flag(task_id, InvokeFrom.DEBUGGER, current_user.id) @@ -525,7 +530,8 @@ class PublishedRagPipelineApi(Resource): Get published pipeline """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() if not pipeline.is_published: return None @@ -545,7 +551,8 @@ class PublishedRagPipelineApi(Resource): Publish workflow """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() rag_pipeline_service = RagPipelineService() @@ -580,7 +587,8 @@ class DefaultRagPipelineBlockConfigsApi(Resource): Get default block config """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() # Get default block configs @@ -599,7 +607,8 @@ class DefaultRagPipelineBlockConfigApi(Resource): Get default block config """ # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -631,7 +640,8 @@ class PublishedAllRagPipelineApi(Resource): """ Get published workflows """ - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -681,7 +691,8 @@ class RagPipelineByIdApi(Resource): Update workflow attributes """ # Check permission - if not isinstance(current_user, Account) or not current_user.has_edit_permission: + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: raise Forbidden() parser = reqparse.RequestParser() @@ -733,13 +744,11 @@ class PublishedRagPipelineSecondStepApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def get(self, pipeline: Pipeline): """ Get second step parameters of rag pipeline """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("node_id", type=str, required=True, location="args") args = parser.parse_args() @@ -759,13 +768,11 @@ class PublishedRagPipelineFirstStepApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def get(self, pipeline: Pipeline): """ Get first step parameters of rag pipeline """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("node_id", type=str, required=True, location="args") args = parser.parse_args() @@ -785,13 +792,11 @@ class DraftRagPipelineFirstStepApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def get(self, pipeline: Pipeline): """ Get first step parameters of rag pipeline """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("node_id", type=str, required=True, location="args") args = parser.parse_args() @@ -811,13 +816,11 @@ class DraftRagPipelineSecondStepApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required def get(self, pipeline: Pipeline): """ Get second step parameters of rag pipeline """ - # The role of the current user in the ta table must be admin, owner, or editor - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() parser = reqparse.RequestParser() parser.add_argument("node_id", type=str, required=True, location="args") args = parser.parse_args() @@ -880,7 +883,7 @@ class RagPipelineWorkflowRunNodeExecutionListApi(Resource): @account_initialization_required @get_rag_pipeline @marshal_with(workflow_run_node_execution_list_fields) - def get(self, pipeline: Pipeline, run_id): + def get(self, pipeline: Pipeline, run_id: str): """ Get workflow run node execution list """ @@ -903,14 +906,8 @@ class DatasourceListApi(Resource): @login_required @account_initialization_required def get(self): - user = current_user - if not isinstance(user, Account): - raise Forbidden() - tenant_id = user.current_tenant_id - if not tenant_id: - raise Forbidden() - - return jsonable_encoder(RagPipelineManageService.list_rag_pipeline_datasources(tenant_id)) + _, current_tenant_id = current_account_with_tenant() + return jsonable_encoder(RagPipelineManageService.list_rag_pipeline_datasources(current_tenant_id)) @console_ns.route("/rag/pipelines//workflows/draft/nodes//last-run") @@ -940,11 +937,11 @@ class RagPipelineTransformApi(Resource): @setup_required @login_required @account_initialization_required - def post(self, dataset_id): - if not isinstance(current_user, Account): - raise Forbidden() + @edit_permission_required + def post(self, dataset_id: str): + current_user, _ = current_account_with_tenant() - if not (current_user.has_edit_permission or current_user.is_dataset_operator): + if not current_user.is_dataset_operator: raise Forbidden() dataset_id = str(dataset_id) @@ -959,14 +956,13 @@ class RagPipelineDatasourceVariableApi(Resource): @login_required @account_initialization_required @get_rag_pipeline + @edit_permission_required @marshal_with(workflow_run_node_execution_fields) def post(self, pipeline: Pipeline): """ Set datasource variables """ - if not isinstance(current_user, Account) or not current_user.has_edit_permission: - raise Forbidden() - + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("datasource_type", type=str, required=True, location="json") parser.add_argument("datasource_info", type=dict, required=True, location="json") diff --git a/api/controllers/console/datasets/wraps.py b/api/controllers/console/datasets/wraps.py index 98abb3ef8d..a8c1298e3e 100644 --- a/api/controllers/console/datasets/wraps.py +++ b/api/controllers/console/datasets/wraps.py @@ -3,8 +3,7 @@ from functools import wraps from controllers.console.datasets.error import PipelineNotFoundError from extensions.ext_database import db -from libs.login import current_user -from models.account import Account +from libs.login import current_account_with_tenant from models.dataset import Pipeline @@ -17,8 +16,7 @@ def get_rag_pipeline( if not kwargs.get("pipeline_id"): raise ValueError("missing pipeline_id in path parameters") - if not isinstance(current_user, Account): - raise ValueError("current_user is not an account") + _, current_tenant_id = current_account_with_tenant() pipeline_id = kwargs.get("pipeline_id") pipeline_id = str(pipeline_id) @@ -27,7 +25,7 @@ def get_rag_pipeline( pipeline = ( db.session.query(Pipeline) - .where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_user.current_tenant_id) + .where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id) .first() ) diff --git a/api/controllers/console/explore/message.py b/api/controllers/console/explore/message.py index b045e47846..064e026753 100644 --- a/api/controllers/console/explore/message.py +++ b/api/controllers/console/explore/message.py @@ -23,8 +23,7 @@ from core.model_runtime.errors.invoke import InvokeError from fields.message_fields import message_infinite_scroll_pagination_fields from libs import helper from libs.helper import uuid_value -from libs.login import current_user -from models import Account +from libs.login import current_account_with_tenant from models.model import AppMode from services.app_generate_service import AppGenerateService from services.errors.app import MoreLikeThisDisabledError @@ -48,6 +47,7 @@ logger = logging.getLogger(__name__) class MessageListApi(InstalledAppResource): @marshal_with(message_infinite_scroll_pagination_fields) def get(self, installed_app): + current_user, _ = current_account_with_tenant() app_model = installed_app.app app_mode = AppMode.value_of(app_model.mode) @@ -61,8 +61,6 @@ class MessageListApi(InstalledAppResource): args = parser.parse_args() try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") return MessageService.pagination_by_first_id( app_model, current_user, args["conversation_id"], args["first_id"], args["limit"] ) @@ -78,6 +76,7 @@ class MessageListApi(InstalledAppResource): ) class MessageFeedbackApi(InstalledAppResource): def post(self, installed_app, message_id): + current_user, _ = current_account_with_tenant() app_model = installed_app.app message_id = str(message_id) @@ -88,8 +87,6 @@ class MessageFeedbackApi(InstalledAppResource): args = parser.parse_args() try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") MessageService.create_feedback( app_model=app_model, message_id=message_id, @@ -109,6 +106,7 @@ class MessageFeedbackApi(InstalledAppResource): ) class MessageMoreLikeThisApi(InstalledAppResource): def get(self, installed_app, message_id): + current_user, _ = current_account_with_tenant() app_model = installed_app.app if app_model.mode != "completion": raise NotCompletionAppError() @@ -124,8 +122,6 @@ class MessageMoreLikeThisApi(InstalledAppResource): streaming = args["response_mode"] == "streaming" try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") response = AppGenerateService.generate_more_like_this( app_model=app_model, user=current_user, @@ -159,6 +155,7 @@ class MessageMoreLikeThisApi(InstalledAppResource): ) class MessageSuggestedQuestionApi(InstalledAppResource): def get(self, installed_app, message_id): + current_user, _ = current_account_with_tenant() app_model = installed_app.app app_mode = AppMode.value_of(app_model.mode) if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}: @@ -167,8 +164,6 @@ class MessageSuggestedQuestionApi(InstalledAppResource): message_id = str(message_id) try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") questions = MessageService.get_suggested_questions_after_answer( app_model=app_model, user=current_user, message_id=message_id, invoke_from=InvokeFrom.EXPLORE ) diff --git a/api/controllers/console/explore/saved_message.py b/api/controllers/console/explore/saved_message.py index 79e4a4339e..830685975b 100644 --- a/api/controllers/console/explore/saved_message.py +++ b/api/controllers/console/explore/saved_message.py @@ -7,8 +7,7 @@ from controllers.console.explore.error import NotCompletionAppError from controllers.console.explore.wraps import InstalledAppResource from fields.conversation_fields import message_file_fields from libs.helper import TimestampField, uuid_value -from libs.login import current_user -from models import Account +from libs.login import current_account_with_tenant from services.errors.message import MessageNotExistsError from services.saved_message_service import SavedMessageService @@ -35,6 +34,7 @@ class SavedMessageListApi(InstalledAppResource): @marshal_with(saved_message_infinite_scroll_pagination_fields) def get(self, installed_app): + current_user, _ = current_account_with_tenant() app_model = installed_app.app if app_model.mode != "completion": raise NotCompletionAppError() @@ -44,11 +44,10 @@ class SavedMessageListApi(InstalledAppResource): parser.add_argument("limit", type=int_range(1, 100), required=False, default=20, location="args") args = parser.parse_args() - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") return SavedMessageService.pagination_by_last_id(app_model, current_user, args["last_id"], args["limit"]) def post(self, installed_app): + current_user, _ = current_account_with_tenant() app_model = installed_app.app if app_model.mode != "completion": raise NotCompletionAppError() @@ -58,8 +57,6 @@ class SavedMessageListApi(InstalledAppResource): args = parser.parse_args() try: - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") SavedMessageService.save(app_model, current_user, args["message_id"]) except MessageNotExistsError: raise NotFound("Message Not Exists.") @@ -72,6 +69,7 @@ class SavedMessageListApi(InstalledAppResource): ) class SavedMessageApi(InstalledAppResource): def delete(self, installed_app, message_id): + current_user, _ = current_account_with_tenant() app_model = installed_app.app message_id = str(message_id) @@ -79,8 +77,6 @@ class SavedMessageApi(InstalledAppResource): if app_model.mode != "completion": raise NotCompletionAppError() - if not isinstance(current_user, Account): - raise ValueError("current_user must be an Account instance") SavedMessageService.delete(app_model, current_user, message_id) return {"result": "success"}, 204 diff --git a/api/controllers/console/explore/wraps.py b/api/controllers/console/explore/wraps.py index 5956eb52c4..df4eed18eb 100644 --- a/api/controllers/console/explore/wraps.py +++ b/api/controllers/console/explore/wraps.py @@ -8,9 +8,8 @@ from werkzeug.exceptions import NotFound from controllers.console.explore.error import AppAccessDeniedError from controllers.console.wraps import account_initialization_required from extensions.ext_database import db -from libs.login import current_user, login_required +from libs.login import current_account_with_tenant, login_required from models import InstalledApp -from models.account import Account from services.app_service import AppService from services.enterprise.enterprise_service import EnterpriseService from services.feature_service import FeatureService @@ -24,13 +23,10 @@ def installed_app_required(view: Callable[Concatenate[InstalledApp, P], R] | Non def decorator(view: Callable[Concatenate[InstalledApp, P], R]): @wraps(view) def decorated(installed_app_id: str, *args: P.args, **kwargs: P.kwargs): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + _, current_tenant_id = current_account_with_tenant() installed_app = ( db.session.query(InstalledApp) - .where( - InstalledApp.id == str(installed_app_id), InstalledApp.tenant_id == current_user.current_tenant_id - ) + .where(InstalledApp.id == str(installed_app_id), InstalledApp.tenant_id == current_tenant_id) .first() ) @@ -56,9 +52,9 @@ def user_allowed_to_access_app(view: Callable[Concatenate[InstalledApp, P], R] | def decorator(view: Callable[Concatenate[InstalledApp, P], R]): @wraps(view) def decorated(installed_app: InstalledApp, *args: P.args, **kwargs: P.kwargs): + current_user, _ = current_account_with_tenant() feature = FeatureService.get_system_features() if feature.webapp_auth.enabled: - assert isinstance(current_user, Account) app_id = installed_app.app_id app_code = AppService.get_app_code_by_id(app_id) res = EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp( diff --git a/api/controllers/console/extension.py b/api/controllers/console/extension.py index dd618307e9..f77996eb6a 100644 --- a/api/controllers/console/extension.py +++ b/api/controllers/console/extension.py @@ -4,8 +4,7 @@ from constants import HIDDEN_VALUE from controllers.console import api, console_ns from controllers.console.wraps import account_initialization_required, setup_required from fields.api_based_extension_fields import api_based_extension_fields -from libs.login import current_account_with_tenant, current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from models.api_based_extension import APIBasedExtension from services.api_based_extension_service import APIBasedExtensionService from services.code_based_extension_service import CodeBasedExtensionService @@ -68,8 +67,7 @@ class APIBasedExtensionAPI(Resource): @account_initialization_required @marshal_with(api_based_extension_fields) def post(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + _, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, location="json") parser.add_argument("api_endpoint", type=str, required=True, location="json") @@ -98,8 +96,6 @@ class APIBasedExtensionDetailAPI(Resource): @account_initialization_required @marshal_with(api_based_extension_fields) def get(self, id): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None api_based_extension_id = str(id) _, tenant_id = current_account_with_tenant() @@ -124,8 +120,6 @@ class APIBasedExtensionDetailAPI(Resource): @account_initialization_required @marshal_with(api_based_extension_fields) def post(self, id): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None api_based_extension_id = str(id) _, current_tenant_id = current_account_with_tenant() @@ -153,8 +147,6 @@ class APIBasedExtensionDetailAPI(Resource): @login_required @account_initialization_required def delete(self, id): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None api_based_extension_id = str(id) _, current_tenant_id = current_account_with_tenant() diff --git a/api/controllers/console/files.py b/api/controllers/console/files.py index 2b63f6febc..1cd193f7ad 100644 --- a/api/controllers/console/files.py +++ b/api/controllers/console/files.py @@ -1,7 +1,6 @@ from typing import Literal from flask import request -from flask_login import current_user from flask_restx import Resource, marshal_with from werkzeug.exceptions import Forbidden @@ -22,8 +21,7 @@ from controllers.console.wraps import ( ) from extensions.ext_database import db from fields.file_fields import file_fields, upload_config_fields -from libs.login import login_required -from models import Account +from libs.login import current_account_with_tenant, login_required from services.file_service import FileService from . import console_ns @@ -53,6 +51,7 @@ class FileApi(Resource): @marshal_with(file_fields) @cloud_edition_billing_resource_check("documents") def post(self): + current_user, _ = current_account_with_tenant() source_str = request.form.get("source") source: Literal["datasets"] | None = "datasets" if source_str == "datasets" else None @@ -65,16 +64,12 @@ class FileApi(Resource): if not file.filename: raise FilenameNotExistsError - if source == "datasets" and not current_user.is_dataset_editor: raise Forbidden() if source not in ("datasets", None): source = None - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - try: upload_file = FileService(db.engine).upload_file( filename=file.filename, diff --git a/api/controllers/console/tag/tags.py b/api/controllers/console/tag/tags.py index b6086c5766..5748ca110d 100644 --- a/api/controllers/console/tag/tags.py +++ b/api/controllers/console/tag/tags.py @@ -5,8 +5,7 @@ from werkzeug.exceptions import Forbidden from controllers.console import console_ns from controllers.console.wraps import account_initialization_required, setup_required from fields.tag_fields import dataset_tag_fields -from libs.login import current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from models.model import Tag from services.tag_service import TagService @@ -24,11 +23,10 @@ class TagListApi(Resource): @account_initialization_required @marshal_with(dataset_tag_fields) def get(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + _, current_tenant_id = current_account_with_tenant() tag_type = request.args.get("type", type=str, default="") keyword = request.args.get("keyword", default=None, type=str) - tags = TagService.get_tags(tag_type, current_user.current_tenant_id, keyword) + tags = TagService.get_tags(tag_type, current_tenant_id, keyword) return tags, 200 @@ -36,8 +34,7 @@ class TagListApi(Resource): @login_required @account_initialization_required def post(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, _ = current_account_with_tenant() # The role of the current user in the ta table must be admin, owner, or editor if not (current_user.has_edit_permission or current_user.is_dataset_editor): raise Forbidden() @@ -63,8 +60,7 @@ class TagUpdateDeleteApi(Resource): @login_required @account_initialization_required def patch(self, tag_id): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, _ = current_account_with_tenant() tag_id = str(tag_id) # The role of the current user in the ta table must be admin, owner, or editor if not (current_user.has_edit_permission or current_user.is_dataset_editor): @@ -87,8 +83,7 @@ class TagUpdateDeleteApi(Resource): @login_required @account_initialization_required def delete(self, tag_id): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, _ = current_account_with_tenant() tag_id = str(tag_id) # The role of the current user in the ta table must be admin, owner, or editor if not current_user.has_edit_permission: @@ -105,8 +100,7 @@ class TagBindingCreateApi(Resource): @login_required @account_initialization_required def post(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, _ = current_account_with_tenant() # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator if not (current_user.has_edit_permission or current_user.is_dataset_editor): raise Forbidden() @@ -133,8 +127,7 @@ class TagBindingDeleteApi(Resource): @login_required @account_initialization_required def post(self): - assert isinstance(current_user, Account) - assert current_user.current_tenant_id is not None + current_user, _ = current_account_with_tenant() # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator if not (current_user.has_edit_permission or current_user.is_dataset_editor): raise Forbidden() diff --git a/api/controllers/console/workspace/__init__.py b/api/controllers/console/workspace/__init__.py index 4a048f3c5e..876e2301f2 100644 --- a/api/controllers/console/workspace/__init__.py +++ b/api/controllers/console/workspace/__init__.py @@ -2,11 +2,11 @@ from collections.abc import Callable from functools import wraps from typing import ParamSpec, TypeVar -from flask_login import current_user from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden from extensions.ext_database import db +from libs.login import current_account_with_tenant from models.account import TenantPluginPermission P = ParamSpec("P") @@ -20,8 +20,9 @@ def plugin_permission_required( def interceptor(view: Callable[P, R]): @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): + current_user, current_tenant_id = current_account_with_tenant() user = current_user - tenant_id = user.current_tenant_id + tenant_id = current_tenant_id with Session(db.engine) as session: permission = ( diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index e2b0e3f84d..a5e6b8f473 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -2,7 +2,6 @@ from datetime import datetime import pytz from flask import request -from flask_login import current_user from flask_restx import Resource, fields, marshal_with, reqparse from sqlalchemy import select from sqlalchemy.orm import Session @@ -37,9 +36,8 @@ from extensions.ext_database import db from fields.member_fields import account_fields from libs.datetime_utils import naive_utc_now from libs.helper import TimestampField, email, extract_remote_ip, timezone -from libs.login import login_required -from models import AccountIntegrate, InvitationCode -from models.account import Account +from libs.login import current_account_with_tenant, login_required +from models import Account, AccountIntegrate, InvitationCode from services.account_service import AccountService from services.billing_service import BillingService from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError @@ -50,9 +48,7 @@ class AccountInitApi(Resource): @setup_required @login_required def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() if account.status == "active": raise AccountAlreadyInitedError() @@ -106,8 +102,7 @@ class AccountProfileApi(Resource): @marshal_with(account_fields) @enterprise_license_required def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() return current_user @@ -118,8 +113,7 @@ class AccountNameApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, location="json") args = parser.parse_args() @@ -140,8 +134,7 @@ class AccountAvatarApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("avatar", type=str, required=True, location="json") args = parser.parse_args() @@ -158,8 +151,7 @@ class AccountInterfaceLanguageApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("interface_language", type=supported_language, required=True, location="json") args = parser.parse_args() @@ -176,8 +168,7 @@ class AccountInterfaceThemeApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("interface_theme", type=str, choices=["light", "dark"], required=True, location="json") args = parser.parse_args() @@ -194,8 +185,7 @@ class AccountTimezoneApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("timezone", type=str, required=True, location="json") args = parser.parse_args() @@ -216,8 +206,7 @@ class AccountPasswordApi(Resource): @account_initialization_required @marshal_with(account_fields) def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("password", type=str, required=False, location="json") parser.add_argument("new_password", type=str, required=True, location="json") @@ -253,9 +242,7 @@ class AccountIntegrateApi(Resource): @account_initialization_required @marshal_with(integrate_list_fields) def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() account_integrates = db.session.scalars( select(AccountIntegrate).where(AccountIntegrate.account_id == account.id) @@ -298,9 +285,7 @@ class AccountDeleteVerifyApi(Resource): @login_required @account_initialization_required def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() token, code = AccountService.generate_account_deletion_verification_code(account) AccountService.send_account_deletion_verification_email(account, code) @@ -314,9 +299,7 @@ class AccountDeleteApi(Resource): @login_required @account_initialization_required def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("token", type=str, required=True, location="json") @@ -358,9 +341,7 @@ class EducationVerifyApi(Resource): @cloud_edition_billing_enabled @marshal_with(verify_fields) def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() return BillingService.EducationIdentity.verify(account.id, account.email) @@ -380,9 +361,7 @@ class EducationApi(Resource): @only_edition_cloud @cloud_edition_billing_enabled def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("token", type=str, required=True, location="json") @@ -399,9 +378,7 @@ class EducationApi(Resource): @cloud_edition_billing_enabled @marshal_with(status_fields) def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") - account = current_user + account, _ = current_account_with_tenant() res = BillingService.EducationIdentity.status(account.id) # convert expire_at to UTC timestamp from isoformat @@ -441,6 +418,7 @@ class ChangeEmailSendEmailApi(Resource): @login_required @account_initialization_required def post(self): + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("email", type=email, required=True, location="json") parser.add_argument("language", type=str, required=False, location="json") @@ -467,8 +445,6 @@ class ChangeEmailSendEmailApi(Resource): raise InvalidTokenError() user_email = reset_data.get("email", "") - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") if user_email != current_user.email: raise InvalidEmailError() else: @@ -551,8 +527,7 @@ class ChangeEmailResetApi(Resource): AccountService.revoke_change_email_token(args["token"]) old_email = reset_data.get("old_email", "") - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() if current_user.email != old_email: raise AccountNotFound() diff --git a/api/controllers/console/workspace/agent_providers.py b/api/controllers/console/workspace/agent_providers.py index e044b2db5b..0a8f49d2e5 100644 --- a/api/controllers/console/workspace/agent_providers.py +++ b/api/controllers/console/workspace/agent_providers.py @@ -3,8 +3,7 @@ from flask_restx import Resource, fields from controllers.console import api, console_ns from controllers.console.wraps import account_initialization_required, setup_required from core.model_runtime.utils.encoders import jsonable_encoder -from libs.login import current_user, login_required -from models.account import Account +from libs.login import current_account_with_tenant, login_required from services.agent_service import AgentService @@ -21,12 +20,11 @@ class AgentProviderListApi(Resource): @login_required @account_initialization_required def get(self): - assert isinstance(current_user, Account) + current_user, current_tenant_id = current_account_with_tenant() user = current_user - assert user.current_tenant_id is not None user_id = user.id - tenant_id = user.current_tenant_id + tenant_id = current_tenant_id return jsonable_encoder(AgentService.list_agent_providers(user_id, tenant_id)) @@ -45,9 +43,5 @@ class AgentProviderApi(Resource): @login_required @account_initialization_required def get(self, provider_name: str): - assert isinstance(current_user, Account) - user = current_user - assert user.current_tenant_id is not None - user_id = user.id - tenant_id = user.current_tenant_id - return jsonable_encoder(AgentService.get_agent_provider(user_id, tenant_id, provider_name)) + current_user, current_tenant_id = current_account_with_tenant() + return jsonable_encoder(AgentService.get_agent_provider(current_user.id, current_tenant_id, provider_name)) diff --git a/api/controllers/console/workspace/load_balancing_config.py b/api/controllers/console/workspace/load_balancing_config.py index 99a1c1f032..4e6f1fa3a5 100644 --- a/api/controllers/console/workspace/load_balancing_config.py +++ b/api/controllers/console/workspace/load_balancing_config.py @@ -5,8 +5,8 @@ from controllers.console import console_ns from controllers.console.wraps import account_initialization_required, setup_required from core.model_runtime.entities.model_entities import ModelType from core.model_runtime.errors.validate import CredentialsValidateFailedError -from libs.login import current_user, login_required -from models.account import Account, TenantAccountRole +from libs.login import current_account_with_tenant, login_required +from models import TenantAccountRole from services.model_load_balancing_service import ModelLoadBalancingService @@ -18,12 +18,11 @@ class LoadBalancingCredentialsValidateApi(Resource): @login_required @account_initialization_required def post(self, provider: str): - assert isinstance(current_user, Account) + current_user, current_tenant_id = current_account_with_tenant() if not TenantAccountRole.is_privileged_role(current_user.current_role): raise Forbidden() - tenant_id = current_user.current_tenant_id - assert tenant_id is not None + tenant_id = current_tenant_id parser = reqparse.RequestParser() parser.add_argument("model", type=str, required=True, nullable=False, location="json") @@ -72,12 +71,11 @@ class LoadBalancingConfigCredentialsValidateApi(Resource): @login_required @account_initialization_required def post(self, provider: str, config_id: str): - assert isinstance(current_user, Account) + current_user, current_tenant_id = current_account_with_tenant() if not TenantAccountRole.is_privileged_role(current_user.current_role): raise Forbidden() - tenant_id = current_user.current_tenant_id - assert tenant_id is not None + tenant_id = current_tenant_id parser = reqparse.RequestParser() parser.add_argument("model", type=str, required=True, nullable=False, location="json") diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index 4a0539785a..5be427e9bb 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -23,8 +23,8 @@ from controllers.console.wraps import ( ) from extensions.ext_database import db from libs.helper import TimestampField -from libs.login import current_user, login_required -from models.account import Account, Tenant, TenantStatus +from libs.login import current_account_with_tenant, login_required +from models.account import Tenant, TenantStatus from services.account_service import TenantService from services.feature_service import FeatureService from services.file_service import FileService @@ -70,8 +70,7 @@ class TenantListApi(Resource): @login_required @account_initialization_required def get(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, current_tenant_id = current_account_with_tenant() tenants = TenantService.get_join_tenants(current_user) tenant_dicts = [] @@ -85,7 +84,7 @@ class TenantListApi(Resource): "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 if current_user.current_tenant_id else False, + "current": tenant.id == current_tenant_id if current_tenant_id else False, } tenant_dicts.append(tenant_dict) @@ -130,8 +129,7 @@ class TenantApi(Resource): if request.path == "/info": logger.warning("Deprecated URL /info was used.") - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() tenant = current_user.current_tenant if not tenant: raise ValueError("No current tenant") @@ -155,8 +153,7 @@ class SwitchWorkspaceApi(Resource): @login_required @account_initialization_required def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("tenant_id", type=str, required=True, location="json") args = parser.parse_args() @@ -181,16 +178,12 @@ class CustomConfigWorkspaceApi(Resource): @account_initialization_required @cloud_edition_billing_resource_check("workspace_custom") def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + _, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("remove_webapp_brand", type=bool, location="json") parser.add_argument("replace_webapp_logo", type=str, location="json") args = parser.parse_args() - - if not current_user.current_tenant_id: - raise ValueError("No current tenant") - tenant = db.get_or_404(Tenant, current_user.current_tenant_id) + tenant = db.get_or_404(Tenant, current_tenant_id) custom_config_dict = { "remove_webapp_brand": args["remove_webapp_brand"], @@ -212,8 +205,7 @@ class WebappLogoWorkspaceApi(Resource): @account_initialization_required @cloud_edition_billing_resource_check("workspace_custom") def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + current_user, _ = current_account_with_tenant() # check file if "file" not in request.files: raise NoFileUploadedError() @@ -253,15 +245,14 @@ class WorkspaceInfoApi(Resource): @account_initialization_required # Change workspace name def post(self): - if not isinstance(current_user, Account): - raise ValueError("Invalid user account") + _, current_tenant_id = current_account_with_tenant() parser = reqparse.RequestParser() parser.add_argument("name", type=str, required=True, location="json") args = parser.parse_args() - if not current_user.current_tenant_id: + if not current_tenant_id: raise ValueError("No current tenant") - tenant = db.get_or_404(Tenant, current_user.current_tenant_id) + tenant = db.get_or_404(Tenant, current_tenant_id) tenant.name = args["name"] db.session.commit() diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 4158f0524f..2fa28711c3 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -30,10 +30,7 @@ def account_initialization_required(view: Callable[P, R]): def decorated(*args: P.args, **kwargs: P.kwargs): # check account initialization current_user, _ = current_account_with_tenant() - - account = current_user - - if account.status == AccountStatus.UNINITIALIZED: + if current_user.status == AccountStatus.UNINITIALIZED: raise AccountNotInitializedError() return view(*args, **kwargs) @@ -249,9 +246,9 @@ def email_password_login_enabled(view: Callable[P, R]): return decorated -def email_register_enabled(view): +def email_register_enabled(view: Callable[P, R]): @wraps(view) - def decorated(*args, **kwargs): + def decorated(*args: P.args, **kwargs: P.kwargs): features = FeatureService.get_system_features() if features.is_allow_register: return view(*args, **kwargs) @@ -299,3 +296,16 @@ def knowledge_pipeline_publish_enabled(view: Callable[P, R]): abort(403) return decorated + + +def edit_permission_required(f: Callable[P, R]): + @wraps(f) + def decorated_function(*args: P.args, **kwargs: P.kwargs): + from werkzeug.exceptions import Forbidden + + current_user, _ = current_account_with_tenant() + if not current_user.has_edit_permission: + raise Forbidden() + return f(*args, **kwargs) + + return decorated_function diff --git a/api/controllers/inner_api/mail.py b/api/controllers/inner_api/mail.py index 0b2be03e43..39411a077a 100644 --- a/api/controllers/inner_api/mail.py +++ b/api/controllers/inner_api/mail.py @@ -17,7 +17,7 @@ class BaseMail(Resource): def post(self): args = _mail_parser.parse_args() - send_inner_email_task.delay( + send_inner_email_task.delay( # type: ignore to=args["to"], subject=args["subject"], body=args["body"], diff --git a/api/controllers/inner_api/plugin/plugin.py b/api/controllers/inner_api/plugin/plugin.py index deab50076d..e4fe8d44bf 100644 --- a/api/controllers/inner_api/plugin/plugin.py +++ b/api/controllers/inner_api/plugin/plugin.py @@ -31,7 +31,7 @@ from core.plugin.entities.request import ( ) from core.tools.entities.tool_entities import ToolProviderType from libs.helper import length_prefixed_response -from models.account import Account, Tenant +from models import Account, Tenant from models.model import EndUser diff --git a/api/controllers/inner_api/workspace/workspace.py b/api/controllers/inner_api/workspace/workspace.py index 47f0240cd2..861da57708 100644 --- a/api/controllers/inner_api/workspace/workspace.py +++ b/api/controllers/inner_api/workspace/workspace.py @@ -7,7 +7,7 @@ from controllers.inner_api import inner_api_ns from controllers.inner_api.wraps import enterprise_inner_api_only from events.tenant_event import tenant_was_created from extensions.ext_database import db -from models.account import Account +from models import Account from services.account_service import TenantService diff --git a/api/controllers/service_api/app/annotation.py b/api/controllers/service_api/app/annotation.py index ad1bdc7334..0521f1537c 100644 --- a/api/controllers/service_api/app/annotation.py +++ b/api/controllers/service_api/app/annotation.py @@ -10,7 +10,7 @@ from controllers.service_api.wraps import validate_app_token from extensions.ext_redis import redis_client from fields.annotation_fields import annotation_fields, build_annotation_model from libs.login import current_user -from models.account import Account +from models import Account from models.model import App from services.annotation_service import AppAnnotationService diff --git a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py index 13ef8abc2d..38891f0180 100644 --- a/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py +++ b/api/controllers/service_api/dataset/rag_pipeline/rag_pipeline_workflow.py @@ -17,7 +17,7 @@ from core.app.apps.pipeline.pipeline_generator import PipelineGenerator from core.app.entities.app_invoke_entities import InvokeFrom from libs import helper from libs.login import current_user -from models.account import Account +from models import Account from models.dataset import Pipeline from models.engine import db from services.errors.file import FileTooLargeError, UnsupportedFileTypeError diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py index 2c9be4e887..638ab528f3 100644 --- a/api/controllers/service_api/wraps.py +++ b/api/controllers/service_api/wraps.py @@ -17,7 +17,7 @@ from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from libs.login import current_user -from models.account import Account, Tenant, TenantAccountJoin, TenantStatus +from models import Account, Tenant, TenantAccountJoin, TenantStatus from models.dataset import Dataset, RateLimitLog from models.model import ApiToken, App, DefaultEndUserSessionID, EndUser from services.feature_service import FeatureService diff --git a/api/controllers/web/forgot_password.py b/api/controllers/web/forgot_password.py index c743d0f52b..cbafd70e99 100644 --- a/api/controllers/web/forgot_password.py +++ b/api/controllers/web/forgot_password.py @@ -20,7 +20,7 @@ from controllers.web import web_ns from extensions.ext_database import db from libs.helper import email, extract_remote_ip from libs.password import hash_password, valid_password -from models.account import Account +from models import Account from services.account_service import AccountService diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py index e021b0aca7..b5af6382e8 100644 --- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py +++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py @@ -70,8 +70,7 @@ from core.workflow.system_variable import SystemVariable from core.workflow.workflow_cycle_manager import CycleManagerWorkflowInfo, WorkflowCycleManager from extensions.ext_database import db from libs.datetime_utils import naive_utc_now -from models import Conversation, EndUser, Message, MessageFile -from models.account import Account +from models import Account, Conversation, EndUser, Message, MessageFile from models.enums import CreatorUserRole from models.workflow import Workflow diff --git a/api/core/app/apps/chat/app_generator.py b/api/core/app/apps/chat/app_generator.py index 8bd956b314..c1251d2feb 100644 --- a/api/core/app/apps/chat/app_generator.py +++ b/api/core/app/apps/chat/app_generator.py @@ -23,7 +23,7 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db from factories import file_factory -from models.account import Account +from models import Account from models.model import App, EndUser from services.conversation_service import ConversationService diff --git a/api/core/app/apps/workflow/generate_task_pipeline.py b/api/core/app/apps/workflow/generate_task_pipeline.py index 56b0d91141..ec4dc87643 100644 --- a/api/core/app/apps/workflow/generate_task_pipeline.py +++ b/api/core/app/apps/workflow/generate_task_pipeline.py @@ -61,7 +61,7 @@ from core.workflow.repositories.workflow_node_execution_repository import Workfl from core.workflow.system_variable import SystemVariable from core.workflow.workflow_cycle_manager import CycleManagerWorkflowInfo, WorkflowCycleManager from extensions.ext_database import db -from models.account import Account +from models import Account from models.enums import CreatorUserRole from models.model import EndUser from models.workflow import ( diff --git a/api/core/ops/ops_trace_manager.py b/api/core/ops/ops_trace_manager.py index 526e78f1d0..7db9b076d2 100644 --- a/api/core/ops/ops_trace_manager.py +++ b/api/core/ops/ops_trace_manager.py @@ -913,4 +913,4 @@ class TraceQueueManager: "file_id": file_id, "app_id": task.app_id, } - process_trace_tasks.delay(file_info) + process_trace_tasks.delay(file_info) # type: ignore diff --git a/api/core/plugin/backwards_invocation/app.py b/api/core/plugin/backwards_invocation/app.py index 8b08b09eb9..32ac132e1e 100644 --- a/api/core/plugin/backwards_invocation/app.py +++ b/api/core/plugin/backwards_invocation/app.py @@ -14,7 +14,7 @@ from core.app.apps.workflow.app_generator import WorkflowAppGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.plugin.backwards_invocation.base import BaseBackwardsInvocation from extensions.ext_database import db -from models.account import Account +from models import Account from models.model import App, AppMode, EndUser diff --git a/api/core/repositories/celery_workflow_execution_repository.py b/api/core/repositories/celery_workflow_execution_repository.py index eda7b54d6a..460bb75722 100644 --- a/api/core/repositories/celery_workflow_execution_repository.py +++ b/api/core/repositories/celery_workflow_execution_repository.py @@ -108,7 +108,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository): execution_data = execution.model_dump() # Queue the save operation as a Celery task (fire and forget) - save_workflow_execution_task.delay( + save_workflow_execution_task.delay( # type: ignore execution_data=execution_data, tenant_id=self._tenant_id, app_id=self._app_id or "", diff --git a/api/core/tools/utils/message_transformer.py b/api/core/tools/utils/message_transformer.py index 0851a54338..ca2aa39861 100644 --- a/api/core/tools/utils/message_transformer.py +++ b/api/core/tools/utils/message_transformer.py @@ -12,7 +12,7 @@ from core.file import File, FileTransferMethod, FileType from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.tool_file_manager import ToolFileManager from libs.login import current_user -from models.account import Account +from models import Account logger = logging.getLogger(__name__) diff --git a/api/events/event_handlers/clean_when_dataset_deleted.py b/api/events/event_handlers/clean_when_dataset_deleted.py index 7caa2d1cc9..0f6aa0e778 100644 --- a/api/events/event_handlers/clean_when_dataset_deleted.py +++ b/api/events/event_handlers/clean_when_dataset_deleted.py @@ -1,10 +1,13 @@ from events.dataset_event import dataset_was_deleted +from models import Dataset from tasks.clean_dataset_task import clean_dataset_task @dataset_was_deleted.connect -def handle(sender, **kwargs): +def handle(sender: Dataset, **kwargs): dataset = sender + assert dataset.doc_form + assert dataset.indexing_technique clean_dataset_task.delay( dataset.id, dataset.tenant_id, diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 5571c0d9ba..836a5d938c 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -9,7 +9,7 @@ from configs import dify_config from dify_app import DifyApp from extensions.ext_database import db from libs.passport import PassportService -from models.account import Account, Tenant, TenantAccountJoin +from models import Account, Tenant, TenantAccountJoin from models.model import AppMCPServer, EndUser from services.account_service import AccountService diff --git a/api/libs/external_api.py b/api/libs/external_api.py index 25a82f8a96..a59230caaa 100644 --- a/api/libs/external_api.py +++ b/api/libs/external_api.py @@ -22,7 +22,7 @@ def register_external_error_handlers(api: Api): got_request_exception.send(current_app, exception=e) # If Werkzeug already prepared a Response, just use it. - if getattr(e, "response", None) is not None: + if e.response is not None: return e.response status_code = getattr(e, "code", 500) or 500 @@ -106,7 +106,7 @@ def register_external_error_handlers(api: Api): # Log stack exc_info: Any = sys.exc_info() if exc_info[1] is None: - exc_info = None + exc_info = (None, None, None) current_app.log_exception(exc_info) return data, status_code diff --git a/api/libs/helper.py b/api/libs/helper.py index 0551470f65..b878141d8e 100644 --- a/api/libs/helper.py +++ b/api/libs/helper.py @@ -24,7 +24,7 @@ from core.model_runtime.utils.encoders import jsonable_encoder from extensions.ext_redis import redis_client if TYPE_CHECKING: - from models.account import Account + from models import Account from models.model import EndUser logger = logging.getLogger(__name__) @@ -43,7 +43,7 @@ def extract_tenant_id(user: Union["Account", "EndUser"]) -> str | None: Raises: ValueError: If user is neither Account nor EndUser """ - from models.account import Account + from models import Account from models.model import EndUser if isinstance(user, Account): @@ -78,7 +78,7 @@ class AvatarUrlField(fields.Raw): if obj is None: return None - from models.account import Account + from models import Account if isinstance(obj, Account) and obj.avatar is not None: return file_helpers.get_signed_file_url(obj.avatar) diff --git a/api/libs/login.py b/api/libs/login.py index 24b8c4011a..2c75ef9297 100644 --- a/api/libs/login.py +++ b/api/libs/login.py @@ -7,7 +7,7 @@ from flask_login.config import EXEMPT_METHODS # type: ignore from werkzeug.local import LocalProxy from configs import dify_config -from models.account import Account +from models import Account from models.model import EndUser #: A proxy for the current user. If no user is logged in, this will be an diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py index ef6edd6709..b70707b17e 100644 --- a/api/schedule/mail_clean_document_notify_task.py +++ b/api/schedule/mail_clean_document_notify_task.py @@ -10,7 +10,7 @@ from configs import dify_config from extensions.ext_database import db from extensions.ext_mail import mail from libs.email_i18n import EmailType, get_email_i18n_service -from models.account import Account, Tenant, TenantAccountJoin +from models import Account, Tenant, TenantAccountJoin from models.dataset import Dataset, DatasetAutoDisableLog from services.feature_service import FeatureService diff --git a/api/services/agent_service.py b/api/services/agent_service.py index d631ce812f..b2db895a5a 100644 --- a/api/services/agent_service.py +++ b/api/services/agent_service.py @@ -10,7 +10,7 @@ from core.plugin.impl.exc import PluginDaemonClientSideError from core.tools.tool_manager import ToolManager from extensions.ext_database import db from libs.login import current_user -from models.account import Account +from models import Account from models.model import App, Conversation, EndUser, Message, MessageAgentThought diff --git a/api/services/app_service.py b/api/services/app_service.py index 4fc6cf2494..5f8c5089c9 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -18,7 +18,7 @@ from events.app_event import app_was_created from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from libs.login import current_user -from models.account import Account +from models import Account from models.model import App, AppMode, AppModelConfig, Site from models.tools import ApiToolProvider from services.billing_service import BillingService diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 9d6c5b4b31..a6851d2638 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -7,7 +7,7 @@ from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fix from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.helper import RateLimiter -from models.account import Account, TenantAccountJoin, TenantAccountRole +from models import Account, TenantAccountJoin, TenantAccountRole class BillingService: diff --git a/api/services/conversation_service.py b/api/services/conversation_service.py index a8e51a426d..39d6c81621 100644 --- a/api/services/conversation_service.py +++ b/api/services/conversation_service.py @@ -14,8 +14,7 @@ from extensions.ext_database import db from factories import variable_factory from libs.datetime_utils import naive_utc_now from libs.infinite_scroll_pagination import InfiniteScrollPagination -from models import ConversationVariable -from models.account import Account +from models import Account, ConversationVariable from models.model import App, Conversation, EndUser, Message from services.errors.conversation import ( ConversationNotExistsError, diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 53216e4fdd..f4047da6b8 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -29,7 +29,7 @@ from extensions.ext_redis import redis_client from libs import helper from libs.datetime_utils import naive_utc_now from libs.login import current_user -from models.account import Account, TenantAccountRole +from models import Account, TenantAccountRole from models.dataset import ( AppDatasetJoin, ChildChunk, diff --git a/api/services/file_service.py b/api/services/file_service.py index f0bb68766d..dd6a829ea2 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -19,7 +19,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor from extensions.ext_storage import storage from libs.datetime_utils import naive_utc_now from libs.helper import extract_tenant_id -from models.account import Account +from models import Account from models.enums import CreatorUserRole from models.model import EndUser, UploadFile diff --git a/api/services/hit_testing_service.py b/api/services/hit_testing_service.py index c6ea35076e..7fa82c6d22 100644 --- a/api/services/hit_testing_service.py +++ b/api/services/hit_testing_service.py @@ -9,7 +9,7 @@ from core.rag.models.document import Document from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.ext_database import db -from models.account import Account +from models import Account from models.dataset import Dataset, DatasetQuery logger = logging.getLogger(__name__) diff --git a/api/services/message_service.py b/api/services/message_service.py index 5e356bf925..9fdff18622 100644 --- a/api/services/message_service.py +++ b/api/services/message_service.py @@ -12,7 +12,7 @@ from core.ops.ops_trace_manager import TraceQueueManager, TraceTask from core.ops.utils import measure_time from extensions.ext_database import db from libs.infinite_scroll_pagination import InfiniteScrollPagination -from models.account import Account +from models import Account from models.model import App, AppMode, AppModelConfig, EndUser, Message, MessageFeedback from services.conversation_service import ConversationService from services.errors.message import ( diff --git a/api/services/metadata_service.py b/api/services/metadata_service.py index 6add830813..5f280c9e57 100644 --- a/api/services/metadata_service.py +++ b/api/services/metadata_service.py @@ -1,12 +1,11 @@ import copy import logging -from flask_login import current_user - from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now +from libs.login import current_account_with_tenant from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding from services.dataset_service import DocumentService from services.entities.knowledge_entities.knowledge_entities import ( @@ -23,11 +22,11 @@ class MetadataService: # check if metadata name is too long if len(metadata_args.name) > 255: raise ValueError("Metadata name cannot exceed 255 characters.") - + current_user, current_tenant_id = current_account_with_tenant() # check if metadata name already exists if ( db.session.query(DatasetMetadata) - .filter_by(tenant_id=current_user.current_tenant_id, dataset_id=dataset_id, name=metadata_args.name) + .filter_by(tenant_id=current_tenant_id, dataset_id=dataset_id, name=metadata_args.name) .first() ): raise ValueError("Metadata name already exists.") @@ -35,7 +34,7 @@ class MetadataService: if field.value == metadata_args.name: raise ValueError("Metadata name already exists in Built-in fields.") metadata = DatasetMetadata( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, dataset_id=dataset_id, type=metadata_args.type, name=metadata_args.name, @@ -53,9 +52,10 @@ class MetadataService: lock_key = f"dataset_metadata_lock_{dataset_id}" # check if metadata name already exists + current_user, current_tenant_id = current_account_with_tenant() if ( db.session.query(DatasetMetadata) - .filter_by(tenant_id=current_user.current_tenant_id, dataset_id=dataset_id, name=name) + .filter_by(tenant_id=current_tenant_id, dataset_id=dataset_id, name=name) .first() ): raise ValueError("Metadata name already exists.") @@ -220,9 +220,10 @@ class MetadataService: db.session.commit() # deal metadata binding db.session.query(DatasetMetadataBinding).filter_by(document_id=operation.document_id).delete() + current_user, current_tenant_id = current_account_with_tenant() for metadata_value in operation.metadata_list: dataset_metadata_binding = DatasetMetadataBinding( - tenant_id=current_user.current_tenant_id, + tenant_id=current_tenant_id, dataset_id=dataset.id, document_id=operation.document_id, metadata_id=metadata_value.id, diff --git a/api/services/oauth_server.py b/api/services/oauth_server.py index b722dbee22..b05b43d76e 100644 --- a/api/services/oauth_server.py +++ b/api/services/oauth_server.py @@ -7,7 +7,7 @@ from werkzeug.exceptions import BadRequest from extensions.ext_database import db from extensions.ext_redis import redis_client -from models.account import Account +from models import Account from models.model import OAuthProviderApp from services.account_service import AccountService diff --git a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py index ca871bcaa1..4ac2e0792b 100644 --- a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py +++ b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py @@ -1,7 +1,7 @@ import yaml -from flask_login import current_user from extensions.ext_database import db +from libs.login import current_account_with_tenant from models.dataset import PipelineCustomizedTemplate from services.rag_pipeline.pipeline_template.pipeline_template_base import PipelineTemplateRetrievalBase from services.rag_pipeline.pipeline_template.pipeline_template_type import PipelineTemplateType @@ -13,9 +13,8 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase): """ def get_pipeline_templates(self, language: str) -> dict: - result = self.fetch_pipeline_templates_from_customized( - tenant_id=current_user.current_tenant_id, language=language - ) + _, current_tenant_id = current_account_with_tenant() + result = self.fetch_pipeline_templates_from_customized(tenant_id=current_tenant_id, language=language) return result def get_pipeline_template_detail(self, template_id: str): diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 13c0ca7392..d2ba462a37 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -54,7 +54,7 @@ from core.workflow.system_variable import SystemVariable from core.workflow.workflow_entry import WorkflowEntry from extensions.ext_database import db from libs.infinite_scroll_pagination import InfiniteScrollPagination -from models.account import Account +from models import Account from models.dataset import ( # type: ignore Dataset, Document, diff --git a/api/services/saved_message_service.py b/api/services/saved_message_service.py index 67a0106bbd..4dd6c8107b 100644 --- a/api/services/saved_message_service.py +++ b/api/services/saved_message_service.py @@ -2,7 +2,7 @@ from typing import Union from extensions.ext_database import db from libs.infinite_scroll_pagination import InfiniteScrollPagination -from models.account import Account +from models import Account from models.model import App, EndUser from models.web import SavedMessage from services.message_service import MessageService diff --git a/api/services/web_conversation_service.py b/api/services/web_conversation_service.py index 0f54e838f3..560aec2330 100644 --- a/api/services/web_conversation_service.py +++ b/api/services/web_conversation_service.py @@ -6,7 +6,7 @@ from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom from extensions.ext_database import db from libs.infinite_scroll_pagination import InfiniteScrollPagination -from models.account import Account +from models import Account from models.model import App, EndUser from models.web import PinnedConversation from services.conversation_service import ConversationService diff --git a/api/services/webapp_auth_service.py b/api/services/webapp_auth_service.py index d30e14f7a1..693bfb95b6 100644 --- a/api/services/webapp_auth_service.py +++ b/api/services/webapp_auth_service.py @@ -10,7 +10,7 @@ from extensions.ext_database import db from libs.helper import TokenManager from libs.passport import PassportService from libs.password import compare_password -from models.account import Account, AccountStatus +from models import Account, AccountStatus from models.model import App, EndUser, Site from services.app_service import AppService from services.enterprise.enterprise_service import EnterpriseService diff --git a/api/services/workflow/workflow_converter.py b/api/services/workflow/workflow_converter.py index 9c09f54bf5..e70b2b5c95 100644 --- a/api/services/workflow/workflow_converter.py +++ b/api/services/workflow/workflow_converter.py @@ -22,7 +22,7 @@ from core.prompt.utils.prompt_template_parser import PromptTemplateParser from core.workflow.nodes import NodeType from events.app_event import app_was_created from extensions.ext_database import db -from models.account import Account +from models import Account from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint from models.model import App, AppMode, AppModelConfig from models.workflow import Workflow, WorkflowType diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index 344b7486ee..5e63a83bb1 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -32,8 +32,7 @@ from factories.file_factory import StorageKeyLoader from factories.variable_factory import build_segment, segment_to_variable from libs.datetime_utils import naive_utc_now from libs.uuid_utils import uuidv7 -from models import App, Conversation -from models.account import Account +from models import Account, App, Conversation from models.enums import DraftVariableType from models.workflow import Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, is_system_variable_editable from repositories.factory import DifyAPIRepositoryFactory diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index dea6a657a4..f765c229ab 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -30,7 +30,7 @@ from extensions.ext_database import db from extensions.ext_storage import storage from factories.file_factory import build_from_mapping, build_from_mappings from libs.datetime_utils import naive_utc_now -from models.account import Account +from models import Account from models.model import App, AppMode from models.tools import WorkflowToolProvider from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowType diff --git a/api/tasks/delete_account_task.py b/api/tasks/delete_account_task.py index 611aef86ad..fb5eb1d691 100644 --- a/api/tasks/delete_account_task.py +++ b/api/tasks/delete_account_task.py @@ -3,7 +3,7 @@ import logging from celery import shared_task from extensions.ext_database import db -from models.account import Account +from models import Account from services.billing_service import BillingService from tasks.mail_account_deletion_task import send_deletion_success_task diff --git a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py index 4171656131..6de95a3b85 100644 --- a/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/priority_rag_pipeline_run_task.py @@ -16,7 +16,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom, RagPipelineGenerat from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity from core.repositories.factory import DifyCoreRepositoryFactory from extensions.ext_database import db -from models.account import Account, Tenant +from models import Account, Tenant from models.dataset import Pipeline from models.enums import WorkflowRunTriggeredFrom from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom diff --git a/api/tasks/rag_pipeline/rag_pipeline_run_task.py b/api/tasks/rag_pipeline/rag_pipeline_run_task.py index 90ebe80daf..f4a092d97e 100644 --- a/api/tasks/rag_pipeline/rag_pipeline_run_task.py +++ b/api/tasks/rag_pipeline/rag_pipeline_run_task.py @@ -17,7 +17,7 @@ from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEnti from core.repositories.factory import DifyCoreRepositoryFactory from extensions.ext_database import db from extensions.ext_redis import redis_client -from models.account import Account, Tenant +from models import Account, Tenant from models.dataset import Pipeline from models.enums import WorkflowRunTriggeredFrom from models.workflow import Workflow, WorkflowNodeExecutionTriggeredFrom diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py index 9c12696824..9d208647e6 100644 --- a/api/tasks/retry_document_indexing_task.py +++ b/api/tasks/retry_document_indexing_task.py @@ -10,7 +10,7 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now -from models.account import Account, Tenant +from models import Account, Tenant from models.dataset import Dataset, Document, DocumentSegment from services.feature_service import FeatureService from services.rag_pipeline.rag_pipeline import RagPipelineService diff --git a/api/tests/test_containers_integration_tests/services/test_account_service.py b/api/tests/test_containers_integration_tests/services/test_account_service.py index 6eff73a8f3..c59fc50f08 100644 --- a/api/tests/test_containers_integration_tests/services/test_account_service.py +++ b/api/tests/test_containers_integration_tests/services/test_account_service.py @@ -8,7 +8,7 @@ from werkzeug.exceptions import Unauthorized from configs import dify_config from controllers.console.error import AccountNotFound, NotAllowedCreateWorkspace -from models.account import AccountStatus, TenantAccountJoin +from models import AccountStatus, TenantAccountJoin from services.account_service import AccountService, RegisterService, TenantService, TokenPair from services.errors.account import ( AccountAlreadyInTenantError, @@ -470,7 +470,7 @@ class TestAccountService: # Verify integration was created from extensions.ext_database import db - from models.account import AccountIntegrate + from models import AccountIntegrate integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="new-google").first() assert integration is not None @@ -505,7 +505,7 @@ class TestAccountService: # Verify integration was updated from extensions.ext_database import db - from models.account import AccountIntegrate + from models import AccountIntegrate integration = ( db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="exists-google").first() @@ -2303,7 +2303,7 @@ class TestRegisterService: # Verify account was created from extensions.ext_database import db - from models.account import Account + from models import Account from models.model import DifySetup account = db.session.query(Account).filter_by(email=admin_email).first() @@ -2352,7 +2352,7 @@ class TestRegisterService: # Verify no entities were created (rollback worked) from extensions.ext_database import db - from models.account import Account, Tenant, TenantAccountJoin + from models import Account, Tenant, TenantAccountJoin from models.model import DifySetup account = db.session.query(Account).filter_by(email=admin_email).first() @@ -2446,7 +2446,7 @@ class TestRegisterService: # Verify OAuth integration was created from extensions.ext_database import db - from models.account import AccountIntegrate + from models import AccountIntegrate integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider=provider).first() assert integration is not None @@ -2472,7 +2472,7 @@ class TestRegisterService: mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False # Execute registration with pending status - from models.account import AccountStatus + from models import AccountStatus account = RegisterService.register( email=email, @@ -2661,7 +2661,7 @@ class TestRegisterService: # Verify new account was created with pending status from extensions.ext_database import db - from models.account import Account, TenantAccountJoin + from models import Account, TenantAccountJoin new_account = db.session.query(Account).filter_by(email=new_member_email).first() assert new_account is not None diff --git a/api/tests/test_containers_integration_tests/services/test_agent_service.py b/api/tests/test_containers_integration_tests/services/test_agent_service.py index c572ddc925..ca513319b2 100644 --- a/api/tests/test_containers_integration_tests/services/test_agent_service.py +++ b/api/tests/test_containers_integration_tests/services/test_agent_service.py @@ -5,7 +5,7 @@ import pytest from faker import Faker from core.plugin.impl.exc import PluginDaemonClientSideError -from models.account import Account +from models import Account from models.model import AppModelConfig, Conversation, EndUser, Message, MessageAgentThought from services.account_service import AccountService, TenantService from services.agent_service import AgentService diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/test_containers_integration_tests/services/test_annotation_service.py index 4768b981cc..2b03ec1c26 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_annotation_service.py @@ -4,7 +4,7 @@ import pytest from faker import Faker from werkzeug.exceptions import NotFound -from models.account import Account +from models import Account from models.model import MessageAnnotation from services.annotation_service import AppAnnotationService from services.app_service import AppService diff --git a/api/tests/test_containers_integration_tests/services/test_app_service.py b/api/tests/test_containers_integration_tests/services/test_app_service.py index cbbbbddb21..e53392bcef 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_service.py @@ -4,7 +4,7 @@ import pytest from faker import Faker from constants.model_template import default_app_templates -from models.account import Account +from models import Account from models.model import App, Site from services.account_service import AccountService, TenantService from services.app_service import AppService diff --git a/api/tests/test_containers_integration_tests/services/test_file_service.py b/api/tests/test_containers_integration_tests/services/test_file_service.py index e6bfc157c7..4c94e42f3e 100644 --- a/api/tests/test_containers_integration_tests/services/test_file_service.py +++ b/api/tests/test_containers_integration_tests/services/test_file_service.py @@ -8,7 +8,7 @@ from sqlalchemy import Engine from werkzeug.exceptions import NotFound from configs import dify_config -from models.account import Account, Tenant +from models import Account, Tenant from models.enums import CreatorUserRole from models.model import EndUser, UploadFile from services.errors.file import FileTooLargeError, UnsupportedFileTypeError diff --git a/api/tests/test_containers_integration_tests/services/test_metadata_service.py b/api/tests/test_containers_integration_tests/services/test_metadata_service.py index 253791cc2d..c8ced3f3a5 100644 --- a/api/tests/test_containers_integration_tests/services/test_metadata_service.py +++ b/api/tests/test_containers_integration_tests/services/test_metadata_service.py @@ -4,7 +4,7 @@ import pytest from faker import Faker from core.rag.index_processor.constant.built_in_field import BuiltInField -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, DatasetMetadata, DatasetMetadataBinding, Document from services.entities.knowledge_entities.knowledge_entities import MetadataArgs from services.metadata_service import MetadataService @@ -17,9 +17,7 @@ class TestMetadataService: def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( - patch( - "services.metadata_service.current_user", create_autospec(Account, instance=True) - ) as mock_current_user, + patch("libs.login.current_user", create_autospec(Account, instance=True)) as mock_current_user, patch("services.metadata_service.redis_client") as mock_redis_client, patch("services.dataset_service.DocumentService") as mock_document_service, ): diff --git a/api/tests/test_containers_integration_tests/services/test_model_provider_service.py b/api/tests/test_containers_integration_tests/services/test_model_provider_service.py index fb319a4963..8cb3572c47 100644 --- a/api/tests/test_containers_integration_tests/services/test_model_provider_service.py +++ b/api/tests/test_containers_integration_tests/services/test_model_provider_service.py @@ -5,7 +5,7 @@ from faker import Faker from core.entities.model_entities import ModelStatus from core.model_runtime.entities.model_entities import FetchFrom, ModelType -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.provider import Provider, ProviderModel, ProviderModelSetting, ProviderType from services.model_provider_service import ModelProviderService diff --git a/api/tests/test_containers_integration_tests/services/test_tag_service.py b/api/tests/test_containers_integration_tests/services/test_tag_service.py index 3d1226019b..6732b8d558 100644 --- a/api/tests/test_containers_integration_tests/services/test_tag_service.py +++ b/api/tests/test_containers_integration_tests/services/test_tag_service.py @@ -5,7 +5,7 @@ from faker import Faker from sqlalchemy import select from werkzeug.exceptions import NotFound -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset from models.model import App, Tag, TagBinding from services.tag_service import TagService diff --git a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py b/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py index 5db7901cbc..bbbf48ede9 100644 --- a/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_web_conversation_service.py @@ -5,7 +5,7 @@ from faker import Faker from sqlalchemy import select from core.app.entities.app_invoke_entities import InvokeFrom -from models.account import Account +from models import Account from models.model import Conversation, EndUser from models.web import PinnedConversation from services.account_service import AccountService, TenantService diff --git a/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py b/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py index 059767458a..9fc16d9eb7 100644 --- a/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py +++ b/api/tests/test_containers_integration_tests/services/test_webapp_auth_service.py @@ -7,7 +7,7 @@ from faker import Faker from werkzeug.exceptions import NotFound, Unauthorized from libs.password import hash_password -from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole from models.model import App, Site from services.errors.account import AccountLoginError, AccountNotFoundError, AccountPasswordError from services.webapp_auth_service import WebAppAuthService, WebAppAuthType diff --git a/api/tests/test_containers_integration_tests/services/test_workspace_service.py b/api/tests/test_containers_integration_tests/services/test_workspace_service.py index 814d1908bd..4249642bc9 100644 --- a/api/tests/test_containers_integration_tests/services/test_workspace_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workspace_service.py @@ -3,7 +3,7 @@ from unittest.mock import patch import pytest from faker import Faker -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from services.workspace_service import WorkspaceService diff --git a/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py b/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py index 7366b08439..0871467a05 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py +++ b/api/tests/test_containers_integration_tests/services/tools/test_api_tools_manage_service.py @@ -3,7 +3,7 @@ from unittest.mock import patch import pytest from faker import Faker -from models.account import Account, Tenant +from models import Account, Tenant from models.tools import ApiToolProvider from services.tools.api_tools_manage_service import ApiToolManageService diff --git a/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py b/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py index f7a4c53318..71d55c3ade 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py +++ b/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py @@ -4,7 +4,7 @@ import pytest from faker import Faker from core.tools.entities.tool_entities import ToolProviderType -from models.account import Account, Tenant +from models import Account, Tenant from models.tools import MCPToolProvider from services.tools.mcp_tools_manage_service import UNCHANGED_SERVER_URL_PLACEHOLDER, MCPToolManageService diff --git a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py index 88aa0b6e72..2c5e719a58 100644 --- a/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py +++ b/api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py @@ -15,7 +15,7 @@ from core.app.app_config.entities import ( ) from core.model_runtime.entities.llm_entities import LLMMode from core.prompt.utils.prompt_template_parser import PromptTemplateParser -from models.account import Account, Tenant +from models import Account, Tenant from models.api_based_extension import APIBasedExtension from models.model import App, AppMode, AppModelConfig from models.workflow import Workflow diff --git a/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py index 96e673d855..68e485107c 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py @@ -6,7 +6,7 @@ from faker import Faker from core.rag.index_processor.constant.index_type import IndexType from extensions.ext_database import db from extensions.ext_redis import redis_client -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, DatasetAutoDisableLog, Document, DocumentSegment from tasks.add_document_to_index_task import add_document_to_index_task diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py index 8628e2af7f..f94c5b19e6 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_batch_clean_document_task.py @@ -14,7 +14,7 @@ from faker import Faker from extensions.ext_database import db from libs.datetime_utils import naive_utc_now -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.model import UploadFile from tasks.batch_clean_document_task import batch_clean_document_task diff --git a/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py index a9cfb6ffd4..1b844d6357 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py @@ -18,7 +18,7 @@ from unittest.mock import MagicMock, patch import pytest from faker import Faker -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import CreatorUserRole from models.model import UploadFile diff --git a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py index 99061d215f..45eb9d4f78 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock, patch import pytest from faker import Faker -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import ( AppDatasetJoin, Dataset, diff --git a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py index 987ebf8aca..8004175b2d 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py @@ -13,7 +13,7 @@ import pytest from faker import Faker from extensions.ext_redis import redis_client -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from tasks.create_segment_to_index_task import create_segment_to_index_task diff --git a/api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py b/api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py index bc3701d098..8785c948d1 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_disable_segment_from_index_task.py @@ -16,7 +16,7 @@ from faker import Faker from extensions.ext_database import db from extensions.ext_redis import redis_client -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from tasks.disable_segment_from_index_task import disable_segment_from_index_task diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py index a315577b78..448f6da5ec 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py @@ -4,7 +4,7 @@ import pytest from faker import Faker from extensions.ext_database import db -from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document from tasks.document_indexing_task import document_indexing_task