diff --git a/api/controllers/console/apikey.py b/api/controllers/console/apikey.py index e576a4d848..c63a21aaea 100644 --- a/api/controllers/console/apikey.py +++ b/api/controllers/console/apikey.py @@ -1,4 +1,5 @@ -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required import flask_restful from flask_restful import Resource, fields, marshal_with from werkzeug.exceptions import Forbidden diff --git a/api/controllers/console/app/app.py b/api/controllers/console/app/app.py index 9e949865f2..8dff7f3abf 100644 --- a/api/controllers/console/app/app.py +++ b/api/controllers/console/app/app.py @@ -3,7 +3,9 @@ import json import logging from datetime import datetime -from flask_login import login_required, current_user +import flask +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal_with, abort, inputs from werkzeug.exceptions import Forbidden @@ -316,7 +318,7 @@ class AppApi(Resource): if current_user.current_tenant.current_role not in ['admin', 'owner']: raise Forbidden() - + app = _get_app(app_id, current_user.current_tenant_id) db.session.delete(app) diff --git a/api/controllers/console/app/audio.py b/api/controllers/console/app/audio.py index 16749e994e..03dd12da3e 100644 --- a/api/controllers/console/app/audio.py +++ b/api/controllers/console/app/audio.py @@ -2,7 +2,7 @@ import logging from flask import request -from flask_login import login_required +from core.login.login import login_required from werkzeug.exceptions import InternalServerError, NotFound import services diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index 0773abb202..866721775b 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -5,7 +5,7 @@ from typing import Generator, Union import flask_login from flask import Response, stream_with_context -from flask_login import login_required +from core.login.login import login_required from werkzeug.exceptions import InternalServerError, NotFound import services diff --git a/api/controllers/console/app/conversation.py b/api/controllers/console/app/conversation.py index fdca10f94a..5b24bb2882 100644 --- a/api/controllers/console/app/conversation.py +++ b/api/controllers/console/app/conversation.py @@ -1,7 +1,8 @@ from datetime import datetime import pytz -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal_with from flask_restful.inputs import int_range from sqlalchemy import or_, func diff --git a/api/controllers/console/app/generator.py b/api/controllers/console/app/generator.py index f572f855e2..50161cd22d 100644 --- a/api/controllers/console/app/generator.py +++ b/api/controllers/console/app/generator.py @@ -1,4 +1,5 @@ -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse from controllers.console import api diff --git a/api/controllers/console/app/message.py b/api/controllers/console/app/message.py index 9c527eddc9..9745e2f41a 100644 --- a/api/controllers/console/app/message.py +++ b/api/controllers/console/app/message.py @@ -3,7 +3,7 @@ import logging from typing import Union, Generator from flask import Response, stream_with_context -from flask_login import current_user, login_required +from flask_login import current_user from flask_restful import Resource, reqparse, marshal_with, fields from flask_restful.inputs import int_range from werkzeug.exceptions import InternalServerError, NotFound @@ -16,6 +16,7 @@ from controllers.console.setup import setup_required from controllers.console.wraps import account_initialization_required from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \ ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError +from core.login.login import login_required from libs.helper import uuid_value, TimestampField from libs.infinite_scroll_pagination import InfiniteScrollPagination from extensions.ext_database import db diff --git a/api/controllers/console/app/model_config.py b/api/controllers/console/app/model_config.py index c8392e521f..4ffb8dd13f 100644 --- a/api/controllers/console/app/model_config.py +++ b/api/controllers/console/app/model_config.py @@ -3,12 +3,13 @@ import json from flask import request from flask_restful import Resource -from flask_login import login_required, current_user +from flask_login import current_user from controllers.console import api from controllers.console.app import _get_app from controllers.console.setup import setup_required from controllers.console.wraps import account_initialization_required +from core.login.login import login_required from events.app_event import app_model_config_was_updated from extensions.ext_database import db from models.model import AppModelConfig diff --git a/api/controllers/console/app/site.py b/api/controllers/console/app/site.py index 952bf45303..128df38110 100644 --- a/api/controllers/console/app/site.py +++ b/api/controllers/console/app/site.py @@ -1,5 +1,6 @@ # -*- coding:utf-8 -*- -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal_with from werkzeug.exceptions import NotFound, Forbidden diff --git a/api/controllers/console/app/statistic.py b/api/controllers/console/app/statistic.py index f13721350f..c7fb8fa8eb 100644 --- a/api/controllers/console/app/statistic.py +++ b/api/controllers/console/app/statistic.py @@ -4,7 +4,8 @@ from datetime import datetime import pytz from flask import jsonify -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse from controllers.console import api diff --git a/api/controllers/console/auth/data_source_oauth.py b/api/controllers/console/auth/data_source_oauth.py index f8106129dd..17823f235f 100644 --- a/api/controllers/console/auth/data_source_oauth.py +++ b/api/controllers/console/auth/data_source_oauth.py @@ -5,9 +5,12 @@ from typing import Optional import flask_login import requests from flask import request, redirect, current_app, session -from flask_login import current_user, login_required +from flask_login import current_user + from flask_restful import Resource from werkzeug.exceptions import Forbidden + +from core.login.login import login_required from libs.oauth_data_source import NotionOAuth from controllers.console import api from ..setup import setup_required diff --git a/api/controllers/console/datasets/data_source.py b/api/controllers/console/datasets/data_source.py index d65ef851be..87aa263040 100644 --- a/api/controllers/console/datasets/data_source.py +++ b/api/controllers/console/datasets/data_source.py @@ -3,7 +3,8 @@ import json from cachetools import TTLCache from flask import request, current_app -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, marshal_with, fields, reqparse, marshal from werkzeug.exceptions import NotFound diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index a2bdf28356..d1bc960bd8 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -1,6 +1,7 @@ # -*- coding:utf-8 -*- from flask import request -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal, marshal_with from werkzeug.exceptions import NotFound, Forbidden import services diff --git a/api/controllers/console/datasets/datasets_document.py b/api/controllers/console/datasets/datasets_document.py index 5d67e5bfe0..f0a884cbec 100644 --- a/api/controllers/console/datasets/datasets_document.py +++ b/api/controllers/console/datasets/datasets_document.py @@ -4,7 +4,8 @@ from datetime import datetime from typing import List from flask import request -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, fields, marshal, marshal_with, reqparse from sqlalchemy import desc, asc from werkzeug.exceptions import NotFound, Forbidden @@ -764,11 +765,13 @@ class DocumentMetadataApi(DocumentResource): metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type] document.doc_metadata = {} - - for key, value_type in metadata_schema.items(): - value = doc_metadata.get(key) - if value is not None and isinstance(value, value_type): - document.doc_metadata[key] = value + if doc_type == 'others': + document.doc_metadata = doc_metadata + else: + for key, value_type in metadata_schema.items(): + value = doc_metadata.get(key) + if value is not None and isinstance(value, value_type): + document.doc_metadata[key] = value document.doc_type = doc_type document.updated_at = datetime.utcnow() diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index 7ad4c8a2d6..e31ea030e1 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/api/controllers/console/datasets/datasets_segments.py @@ -1,9 +1,8 @@ # -*- coding:utf-8 -*- import uuid from datetime import datetime - from flask import request -from flask_login import login_required, current_user +from flask_login import current_user from flask_restful import Resource, reqparse, fields, marshal from werkzeug.exceptions import NotFound, Forbidden @@ -15,6 +14,7 @@ from controllers.console.setup import setup_required from controllers.console.wraps import account_initialization_required from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError from core.model_providers.model_factory import ModelFactory +from core.login.login import login_required from extensions.ext_database import db from extensions.ext_redis import redis_client from models.dataset import DocumentSegment diff --git a/api/controllers/console/datasets/file.py b/api/controllers/console/datasets/file.py index d208d7a25c..aef25b24f8 100644 --- a/api/controllers/console/datasets/file.py +++ b/api/controllers/console/datasets/file.py @@ -8,7 +8,8 @@ from pathlib import Path from cachetools import TTLCache from flask import request, current_app -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, marshal_with, fields from werkzeug.exceptions import NotFound diff --git a/api/controllers/console/datasets/hit_testing.py b/api/controllers/console/datasets/hit_testing.py index 00d14d93dc..ac37520790 100644 --- a/api/controllers/console/datasets/hit_testing.py +++ b/api/controllers/console/datasets/hit_testing.py @@ -1,6 +1,7 @@ import logging -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, marshal, fields from werkzeug.exceptions import InternalServerError, NotFound, Forbidden diff --git a/api/controllers/console/explore/installed_app.py b/api/controllers/console/explore/installed_app.py index 3a2a1dbee9..b6592e94de 100644 --- a/api/controllers/console/explore/installed_app.py +++ b/api/controllers/console/explore/installed_app.py @@ -1,7 +1,8 @@ # -*- coding:utf-8 -*- from datetime import datetime -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal_with, inputs from sqlalchemy import and_ from werkzeug.exceptions import NotFound, Forbidden, BadRequest diff --git a/api/controllers/console/explore/recommended_app.py b/api/controllers/console/explore/recommended_app.py index 03a9ef4394..7f34c0e678 100644 --- a/api/controllers/console/explore/recommended_app.py +++ b/api/controllers/console/explore/recommended_app.py @@ -1,5 +1,6 @@ # -*- coding:utf-8 -*- -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, fields, marshal_with from sqlalchemy import and_ diff --git a/api/controllers/console/explore/wraps.py b/api/controllers/console/explore/wraps.py index 601e9352ea..9c88b131e9 100644 --- a/api/controllers/console/explore/wraps.py +++ b/api/controllers/console/explore/wraps.py @@ -1,4 +1,5 @@ -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource from functools import wraps diff --git a/api/controllers/console/universal_chat/wraps.py b/api/controllers/console/universal_chat/wraps.py index 8d9516b9fe..d51b87a13a 100644 --- a/api/controllers/console/universal_chat/wraps.py +++ b/api/controllers/console/universal_chat/wraps.py @@ -1,7 +1,8 @@ import json from functools import wraps -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource from controllers.console.setup import setup_required from controllers.console.wraps import account_initialization_required diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index 2ba004a6fa..238ba96cc6 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -3,7 +3,8 @@ from datetime import datetime import pytz from flask import current_app, request -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, fields, marshal_with from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index 7f4857f767..ab1b692927 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -1,6 +1,7 @@ # -*- coding:utf-8 -*- from flask import current_app -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse, marshal_with, abort, fields, marshal import services diff --git a/api/controllers/console/workspace/model_providers.py b/api/controllers/console/workspace/model_providers.py index 0fa01d28a3..08fdb4f117 100644 --- a/api/controllers/console/workspace/model_providers.py +++ b/api/controllers/console/workspace/model_providers.py @@ -1,4 +1,5 @@ -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse from werkzeug.exceptions import Forbidden diff --git a/api/controllers/console/workspace/models.py b/api/controllers/console/workspace/models.py index 33f8edfed7..08eb41e885 100644 --- a/api/controllers/console/workspace/models.py +++ b/api/controllers/console/workspace/models.py @@ -1,4 +1,5 @@ -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse from controllers.console import api diff --git a/api/controllers/console/workspace/providers.py b/api/controllers/console/workspace/providers.py index b6f9c3c697..df09aad9fd 100644 --- a/api/controllers/console/workspace/providers.py +++ b/api/controllers/console/workspace/providers.py @@ -1,5 +1,6 @@ # -*- coding:utf-8 -*- -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, reqparse from werkzeug.exceptions import Forbidden diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py index b24d5c0d55..692e84ac07 100644 --- a/api/controllers/console/workspace/tool_providers.py +++ b/api/controllers/console/workspace/tool_providers.py @@ -1,6 +1,7 @@ import json -from flask_login import login_required, current_user +from flask_login import current_user +from core.login.login import login_required from flask_restful import Resource, abort, reqparse from werkzeug.exceptions import Forbidden diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index 8b0237eb25..89bb33c053 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -2,10 +2,13 @@ import logging from flask import request -from flask_login import login_required, current_user -from flask_restful import Resource, fields, marshal_with, reqparse, marshal +from flask_login import current_user +from core.login.login import login_required +from flask_restful import Resource, fields, marshal_with, reqparse, marshal, inputs +from flask_restful.inputs import int_range from controllers.console import api +from controllers.console.admin import admin_required from controllers.console.setup import setup_required from controllers.console.error import AccountNotLinkTenantError from controllers.console.wraps import account_initialization_required @@ -43,6 +46,13 @@ tenants_fields = { 'current': fields.Boolean } +workspace_fields = { + 'id': fields.String, + 'name': fields.String, + 'status': fields.String, + 'created_at': TimestampField +} + class TenantListApi(Resource): @setup_required @@ -57,6 +67,38 @@ class TenantListApi(Resource): return {'workspaces': marshal(tenants, tenants_fields)}, 200 +class WorkspaceListApi(Resource): + @setup_required + @admin_required + def get(self): + parser = reqparse.RequestParser() + parser.add_argument('page', type=inputs.int_range(1, 99999), required=False, default=1, location='args') + parser.add_argument('limit', type=inputs.int_range(1, 100), required=False, default=20, location='args') + args = parser.parse_args() + + tenants = db.session.query(Tenant).order_by(Tenant.created_at.desc())\ + .paginate(page=args['page'], per_page=args['limit']) + + has_more = False + if len(tenants.items) == args['limit']: + current_page_first_tenant = tenants[-1] + rest_count = db.session.query(Tenant).filter( + Tenant.created_at < current_page_first_tenant.created_at, + Tenant.id != current_page_first_tenant.id + ).count() + + if rest_count > 0: + has_more = True + total = db.session.query(Tenant).count() + return { + 'data': marshal(tenants.items, workspace_fields), + 'has_more': has_more, + 'limit': args['limit'], + 'page': args['page'], + 'total': total + }, 200 + + class TenantApi(Resource): @setup_required @login_required @@ -92,6 +134,7 @@ class SwitchWorkspaceApi(Resource): api.add_resource(TenantListApi, '/workspaces') # GET for getting all tenants +api.add_resource(WorkspaceListApi, '/all-workspaces') # GET for getting all tenants api.add_resource(TenantApi, '/workspaces/current', endpoint='workspaces_current') # GET for getting current tenant info api.add_resource(TenantApi, '/info', endpoint='info') # Deprecated api.add_resource(SwitchWorkspaceApi, '/workspaces/switch') # POST for switching tenant diff --git a/api/core/login/login.py b/api/core/login/login.py new file mode 100644 index 0000000000..6391bc9e86 --- /dev/null +++ b/api/core/login/login.py @@ -0,0 +1,108 @@ +import os +from functools import wraps + +import flask_login +from flask import current_app +from flask import g +from flask import has_request_context +from flask import request +from flask_login import user_logged_in +from flask_login.config import EXEMPT_METHODS +from werkzeug.exceptions import Unauthorized +from werkzeug.local import LocalProxy + +from extensions.ext_database import db +from models.account import Account, Tenant, TenantAccountJoin + +#: A proxy for the current user. If no user is logged in, this will be an +#: anonymous user +current_user = LocalProxy(lambda: _get_user()) + + +def login_required(func): + """ + If you decorate a view with this, it will ensure that the current user is + logged in and authenticated before calling the actual view. (If they are + not, it calls the :attr:`LoginManager.unauthorized` callback.) For + example:: + + @app.route('/post') + @login_required + def post(): + pass + + If there are only certain times you need to require that your user is + logged in, you can do so with:: + + if not current_user.is_authenticated: + return current_app.login_manager.unauthorized() + + ...which is essentially the code that this function adds to your views. + + It can be convenient to globally turn off authentication when unit testing. + To enable this, if the application configuration variable `LOGIN_DISABLED` + is set to `True`, this decorator will be ignored. + + .. Note :: + + Per `W3 guidelines for CORS preflight requests + `_, + HTTP ``OPTIONS`` requests are exempt from login checks. + + :param func: The view function to decorate. + :type func: function + """ + + @wraps(func) + def decorated_view(*args, **kwargs): + auth_header = request.headers.get('Authorization') + admin_api_key_enable = os.getenv('ADMIN_API_KEY_ENABLE', default='False') + if admin_api_key_enable: + if auth_header: + if ' ' not in auth_header: + raise Unauthorized('Invalid Authorization header format. Expected \'Bearer \' format.') + auth_scheme, auth_token = auth_header.split(None, 1) + auth_scheme = auth_scheme.lower() + if auth_scheme != 'bearer': + raise Unauthorized('Invalid Authorization header format. Expected \'Bearer \' format.') + admin_api_key = os.getenv('ADMIN_API_KEY') + + if admin_api_key: + if os.getenv('ADMIN_API_KEY') == auth_token: + workspace_id = request.headers.get('X-WORKSPACE-ID') + if workspace_id: + tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \ + .filter(Tenant.id == workspace_id) \ + .filter(TenantAccountJoin.tenant_id == Tenant.id) \ + .filter(TenantAccountJoin.role == 'owner') \ + .one_or_none() + if tenant_account_join: + tenant, ta = tenant_account_join + account = Account.query.filter_by(id=ta.account_id).first() + # Login admin + if account: + account.current_tenant = tenant + current_app.login_manager._update_request_context_with_user(account) + user_logged_in.send(current_app._get_current_object(), user=_get_user()) + if request.method in EXEMPT_METHODS or current_app.config.get("LOGIN_DISABLED"): + pass + elif not current_user.is_authenticated: + return current_app.login_manager.unauthorized() + + # flask 1.x compatibility + # current_app.ensure_sync is only available in Flask >= 2.0 + if callable(getattr(current_app, "ensure_sync", None)): + return current_app.ensure_sync(func)(*args, **kwargs) + return func(*args, **kwargs) + + return decorated_view + + +def _get_user(): + if has_request_context(): + if "_login_user" not in g: + current_app.login_manager._load_user() + + return g._login_user + + return None diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index f17f956a34..3b1f6cf549 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -284,8 +284,9 @@ class DocumentService: "github_link": str, "open_source_license": str, "commit_date": str, - "commit_author": str - } + "commit_author": str, + }, + "others": dict } @staticmethod diff --git a/web/app/components/datasets/documents/detail/index.tsx b/web/app/components/datasets/documents/detail/index.tsx index 57db36abe9..a18bae9930 100644 --- a/web/app/components/datasets/documents/detail/index.tsx +++ b/web/app/components/datasets/documents/detail/index.tsx @@ -170,7 +170,7 @@ const DocumentDetail: FC = ({ datasetId, documentId }) => { } {showMetadata && } diff --git a/web/models/datasets.ts b/web/models/datasets.ts index fa52ae48cb..cfaa965e08 100644 --- a/web/models/datasets.ts +++ b/web/models/datasets.ts @@ -242,7 +242,7 @@ export type FullDocumentDetail = SimpleDocumentDetail & { archived_reason: 'rule_modified' | 're_upload' archived_by: string archived_at: number - doc_type?: DocType | null + doc_type?: DocType | null | 'others' doc_metadata?: DocMetadata | null segment_count: number [key: string]: any