some optimization for admin api key, create tenant and reset-encrypt-…

…key-pair command (langgenius#3013)

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

api/commands.py

@@ -109,19 +109,20 @@ def reset_encrypt_key_pair():
         click.echo(click.style('Sorry, only support SELF_HOSTED mode.', fg='red'))
         return
 
-    tenant = db.session.query(Tenant).first()
-    if not tenant:
-        click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red'))
-        return
+    tenants = db.session.query(Tenant).all()
+    for tenant in tenants:
+        if not tenant:
+            click.echo(click.style('Sorry, no workspace found. Please enter /install to initialize.', fg='red'))
+            return
 
-    tenant.encrypt_public_key = generate_key_pair(tenant.id)
+        tenant.encrypt_public_key = generate_key_pair(tenant.id)
 
-    db.session.query(Provider).filter(Provider.provider_type == 'custom').delete()
-    db.session.query(ProviderModel).delete()
-    db.session.commit()
+        db.session.query(Provider).filter(Provider.provider_type == 'custom', Provider.tenant_id == tenant.id).delete()
+        db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete()
+        db.session.commit()
 
-    click.echo(click.style('Congratulations! '
-                           'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green'))
+        click.echo(click.style('Congratulations! '
+                               'the asymmetric key pair of workspace {} has been reset.'.format(tenant.id), fg='green'))
 
 
 @click.command('vdb-migrate', help='migrate vector db.')

api/libs/login.py

@@ -53,7 +53,7 @@ def post():
     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 admin_api_key_enable.lower() == 'true':
             if auth_header:
                 if ' ' not in auth_header:
                     raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')

api/services/account_service.py

@@ -435,20 +435,20 @@ def register(cls, email, name, password: str = None, open_id: str = None, provid
 
             if open_id is not None or provider is not None:
                 AccountService.link_account_integrate(provider, open_id, account)
+            if current_app.config['EDITION'] != 'SELF_HOSTED':
+                tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
 
-            tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+                TenantService.create_tenant_member(tenant, account, role='owner')
+                account.current_tenant = tenant
 
-            TenantService.create_tenant_member(tenant, account, role='owner')
-            account.current_tenant = tenant
+                tenant_was_created.send(tenant)
 
             db.session.commit()
         except Exception as e:
             db.session.rollback()  # todo: do not work
             logging.error(f'Register failed: {e}')
             raise AccountRegisterError(f'Registration failed: {e}') from e
 
-        tenant_was_created.send(tenant)
-
         return account
 
     @classmethod
@@ -461,7 +461,6 @@ def invite_new_member(cls, tenant: Tenant, email: str, language: str, role: str
             name = email.split('@')[0]
 
             account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
-
             # Create new tenant member for invited tenant
             TenantService.create_tenant_member(tenant, account, role)
             TenantService.switch_tenant(account, tenant.id)