feat: custom webapp logo (#1766)

api/controllers/console/workspace/workspace.py

@@ -10,12 +10,15 @@ 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
+from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
+from controllers.console.datasets.error import NoFileUploadedError, TooManyFilesError, FileTooLargeError, UnsupportedFileTypeError
 from libs.helper import TimestampField
 from extensions.ext_database import db
 from models.account import Tenant
+import services
 from services.account_service import TenantService
 from services.workspace_service import WorkspaceService
+from services.file_service import FileService
 
 provider_fields = {
     'provider_name': fields.String,
@@ -34,6 +37,7 @@ tenant_fields = {
     'providers': fields.List(fields.Nested(provider_fields)),
     'in_trial': fields.Boolean,
     'trial_end_reason': fields.String,
+    'custom_config': fields.Raw(attribute='custom_config'),
 }
 
 tenants_fields = {
@@ -130,6 +134,61 @@ class SwitchWorkspaceApi(Resource):
         new_tenant = db.session.query(Tenant).get(args['tenant_id'])  # Get new tenant
 
         return {'result': 'success', 'new_tenant': marshal(WorkspaceService.get_tenant_info(new_tenant), tenant_fields)}
+    
+
+class CustomConfigWorkspaceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @cloud_edition_billing_resource_check('workspace_custom')
+    def post(self):
+        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()
+
+        custom_config_dict = {
+            'remove_webapp_brand': args['remove_webapp_brand'],
+            'replace_webapp_logo': args['replace_webapp_logo'],
+        }
+
+        tenant = db.session.query(Tenant).filter(Tenant.id == current_user.current_tenant_id).one_or_404()
+
+        tenant.custom_config_dict = custom_config_dict
+        db.session.commit()
+
+        return {'result': 'success', 'tenant': marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)}
+    
+
+class WebappLogoWorkspaceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @cloud_edition_billing_resource_check('workspace_custom')
+    def post(self):
+        # get file from request
+        file = request.files['file']
+
+        # check file
+        if 'file' not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        extension = file.filename.split('.')[-1]
+        if extension.lower() not in ['svg', 'png']:
+            raise UnsupportedFileTypeError()
+
+        try:
+            upload_file = FileService.upload_file(file, current_user, True)
+
+        except services.errors.file.FileTooLargeError as file_too_large_error:
+            raise FileTooLargeError(file_too_large_error.description)
+        except services.errors.file.UnsupportedFileTypeError:
+            raise UnsupportedFileTypeError()
+        
+        return { 'id': upload_file.id }, 201
 
 
 api.add_resource(TenantListApi, '/workspaces')  # GET for getting all tenants
@@ -137,3 +196,5 @@ api.add_resource(WorkspaceListApi, '/all-workspaces')  # GET for getting all ten
 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.add_resource(CustomConfigWorkspaceApi, '/workspaces/custom-config')
+api.add_resource(WebappLogoWorkspaceApi, '/workspaces/custom-config/webapp-logo/upload')

api/controllers/console/wraps.py

@@ -63,6 +63,8 @@ def cloud_edition_billing_resource_check(resource: str,
                     abort(403, error_msg)
                 elif resource == 'vector_space' and 0 < vector_space['limit'] <= vector_space['size']:
                     abort(403, error_msg)
+                elif resource == 'workspace_custom' and not billing_info['can_replace_logo']:
+                    abort(403, error_msg)
                 elif resource == 'annotation' and 0 < annotation_quota_limit['limit'] <= annotation_quota_limit['size']:
                     abort(403, error_msg)
                 else:

api/controllers/files/image_preview.py

@@ -1,10 +1,12 @@
 from flask import request, Response
 from flask_restful import Resource
+from werkzeug.exceptions import NotFound
 
 import services
 from controllers.files import api
 from libs.exception import BaseHTTPException
 from services.file_service import FileService
+from services.account_service import TenantService
 
 
 class ImagePreviewApi(Resource):
@@ -29,9 +31,30 @@ class ImagePreviewApi(Resource):
             raise UnsupportedFileTypeError()
 
         return Response(generator, mimetype=mimetype)
+    
+
+class WorkspaceWebappLogoApi(Resource):
+    def get(self, workspace_id):
+        workspace_id = str(workspace_id)
+
+        custom_config = TenantService.get_custom_config(workspace_id)
+        webapp_logo_file_id = custom_config.get('replace_webapp_logo') if custom_config is not None else None
+
+        if not webapp_logo_file_id:
+            raise NotFound(f'webapp logo is not found')
+
+        try:
+            generator, mimetype = FileService.get_public_image_preview(
+                webapp_logo_file_id,
+            )
+        except services.errors.file.UnsupportedFileTypeError:
+            raise UnsupportedFileTypeError()
+
+        return Response(generator, mimetype=mimetype)
 
 
 api.add_resource(ImagePreviewApi, '/files/<uuid:file_id>/image-preview')
+api.add_resource(WorkspaceWebappLogoApi, '/files/workspaces/<uuid:workspace_id>/webapp-logo')
 
 
 class UnsupportedFileTypeError(BaseHTTPException):

api/controllers/web/site.py

@@ -2,6 +2,7 @@
 import os
 
 from flask_restful import fields, marshal_with
+from flask import current_app
 from werkzeug.exceptions import Forbidden
 
 from controllers.web import api
@@ -43,6 +44,7 @@ class AppSiteApi(WebApiResource):
         'model_config': fields.Nested(model_config_fields, allow_null=True),
         'plan': fields.String,
         'can_replace_logo': fields.Boolean,
+        'custom_config': fields.Raw(attribute='custom_config'),
     }
 
     @marshal_with(app_fields)
@@ -80,6 +82,15 @@ class AppSiteInfo:
         self.plan = tenant.plan
         self.can_replace_logo = can_replace_logo
 
+        if can_replace_logo:
+            base_url = current_app.config.get('FILES_URL')
+            remove_webapp_brand = tenant.custom_config_dict.get('remove_webapp_brand', False)
+            replace_webapp_logo = f'{base_url}/files/workspaces/{tenant.id}/webapp-logo' if tenant.custom_config_dict['replace_webapp_logo'] else None
+            self.custom_config = {
+                'remove_webapp_brand': remove_webapp_brand,
+                'replace_webapp_logo': replace_webapp_logo,
+            }
+
         if app.enable_site and site.prompt_public:
             app_model_config = app.app_model_config
             self.model_config = app_model_config

api/core/file/upload_file_parser.py

@@ -10,7 +10,7 @@ from flask import current_app
 
 from extensions.ext_storage import storage
 
-SUPPORT_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif']
+SUPPORT_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif', 'svg']
 
 
 class UploadFileParser:

api/migrations/versions/88072f0caa04_add_custom_config_in_tenant.py

@@ -0,0 +1,32 @@
+"""add custom config in tenant
+
+Revision ID: 88072f0caa04
+Revises: fca025d3b60f
+Create Date: 2023-12-14 07:36:50.705362
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '88072f0caa04'
+down_revision = '246ba09cbbdb'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tenants', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('custom_config', sa.Text(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tenants', schema=None) as batch_op:
+        batch_op.drop_column('custom_config')
+
+    # ### end Alembic commands ###

api/models/account.py

@@ -1,4 +1,6 @@
+import json
 import enum
+from math import e
 from typing import List
 
 from flask_login import UserMixin
@@ -112,6 +114,7 @@ class Tenant(db.Model):
     encrypt_public_key = db.Column(db.Text)
     plan = db.Column(db.String(255), nullable=False, server_default=db.text("'basic'::character varying"))
     status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying"))
+    custom_config = db.Column(db.Text)
     created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
     updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
 
@@ -121,6 +124,14 @@ class Tenant(db.Model):
             Account.id == TenantAccountJoin.account_id,
             TenantAccountJoin.tenant_id == self.id
         ).all()
+    
+    @property
+    def custom_config_dict(self) -> dict:
+        return json.loads(self.custom_config) if self.custom_config else None
+    
+    @custom_config_dict.setter
+    def custom_config_dict(self, value: dict):
+        self.custom_config = json.dumps(value)
 
 
 class TenantAccountJoinRole(enum.Enum):

api/services/account_service.py

@@ -412,6 +412,12 @@ class TenantService:
         db.session.delete(tenant)
         db.session.commit()
 
+    @staticmethod
+    def get_custom_config(tenant_id: str) -> None:
+        tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).one_or_404()
+
+        return tenant.custom_config_dict
+
 
 class RegisterService:
 

api/services/file_service.py

@@ -17,8 +17,8 @@ from models.model import UploadFile, EndUser
 from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
 
 ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'docx', 'csv',
-                      'jpg', 'jpeg', 'png', 'webp', 'gif']
-IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif']
+                      'jpg', 'jpeg', 'png', 'webp', 'gif', 'svg']
+IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif', 'svg']
 PREVIEW_WORDS_LIMIT = 3000
 
 
@@ -154,3 +154,21 @@ class FileService:
         generator = storage.load(upload_file.key, stream=True)
 
         return generator, upload_file.mime_type
+    
+    @staticmethod
+    def get_public_image_preview(file_id: str) -> str:
+        upload_file = db.session.query(UploadFile) \
+            .filter(UploadFile.id == file_id) \
+            .first()
+
+        if not upload_file:
+            raise NotFound("File not found or signature is invalid")
+
+        # extract text from file
+        extension = upload_file.extension
+        if extension.lower() not in IMAGE_EXTENSIONS:
+            raise UnsupportedFileTypeError()
+
+        generator = storage.load(upload_file.key)
+
+        return generator, upload_file.mime_type

api/services/workspace_service.py

@@ -1,8 +1,11 @@
 from flask_login import current_user
 from extensions.ext_database import db
-from models.account import Tenant, TenantAccountJoin
+from models.account import Tenant, TenantAccountJoin, TenantAccountJoinRole
 from models.provider import Provider
 
+from services.billing_service import BillingService
+from services.account_service import TenantService
+
 
 class WorkspaceService:
     @classmethod
@@ -28,6 +31,11 @@ class WorkspaceService:
         ).first()
         tenant_info['role'] = tenant_account_join.role
 
+        billing_info = BillingService.get_info(tenant_info['id'])
+
+        if billing_info['can_replace_logo'] and TenantService.has_roles(tenant, [TenantAccountJoinRole.OWNER, TenantAccountJoinRole.ADMIN]):
+            tenant_info['custom_config'] = tenant.custom_config_dict
+
         # Get providers
         providers = db.session.query(Provider).filter(
             Provider.tenant_id == tenant.id