Feat/token support (#909)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: jyong <jyong@dify.ai>

api/controllers/console/apikey.py

-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

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)

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

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

api/controllers/console/app/conversation.py

 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

api/controllers/console/app/generator.py

-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

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

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

api/controllers/console/app/site.py

 # -*- 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
 

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

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

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
 

api/controllers/console/datasets/datasets.py

 # -*- 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

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()

api/controllers/console/datasets/datasets_segments.py

 # -*- 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

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
 

api/controllers/console/datasets/hit_testing.py

 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
 

api/controllers/console/explore/installed_app.py

 # -*- 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

api/controllers/console/explore/recommended_app.py

 # -*- 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_
 

api/controllers/console/explore/wraps.py

-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
 

api/controllers/console/universal_chat/wraps.py

 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

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

api/controllers/console/workspace/members.py

 # -*- 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

api/controllers/console/workspace/model_providers.py

-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
 

api/controllers/console/workspace/models.py

-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

api/controllers/console/workspace/providers.py

 # -*- 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
 

api/controllers/console/workspace/tool_providers.py

 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
 

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

api/core/login/login.py

+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://www.w3.org/TR/cors/#cross-origin-request-with-preflight-0>`_,
+        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 <api-key>\' 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 <api-key>\' 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

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

web/app/components/datasets/documents/detail/index.tsx

@@ -170,7 +170,7 @@ const DocumentDetail: FC<Props> = ({ datasetId, documentId }) => {
             </div>
           }
           {showMetadata && <Metadata
-            docDetail={{ ...documentDetail, ...documentMetadata } as any}
+            docDetail={{ ...documentDetail, ...documentMetadata, doc_type: documentDetail?.doc_type === 'others' ? '' : documentDetail?.doc_type } as any}
             loading={isMetadataLoading}
             onUpdate={metadataMutate}
           />}

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