Feat/token support (langgenius#909)

Co-authored-by: StyleZhang <[email protected]>
Co-authored-by: jyong <[email protected]>

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

api/controllers/console/app/app.py

@@ -3,7 +3,9 @@
 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 @@ def delete(self, app_id):
 
         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 @@
 
 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

@@ -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

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

api/controllers/console/app/message.py

@@ -3,7 +3,7 @@
 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.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 @@
 
 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

@@ -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
 

api/controllers/console/app/statistic.py

@@ -4,7 +4,8 @@
 
 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 @@
 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 @@
 
 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

@@ -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

api/controllers/console/datasets/datasets_document.py

@@ -4,7 +4,8 @@
 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 @@ def put(self, dataset_id, document_id):
         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

@@ -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.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 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

@@ -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
 

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

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_
 

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
 

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

api/controllers/console/workspace/account.py

@@ -3,7 +3,8 @@
 
 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

@@ -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

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
 

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

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
 

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
 

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 @@
     '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 @@ def get(self):
         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 @@ def post(self):
 
 
 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