diff --git a/api/commands.py b/api/commands.py index 86cc9e204f..186b97c3fa 100644 --- a/api/commands.py +++ b/api/commands.py @@ -1,11 +1,13 @@ import base64 import json import secrets +from typing import Optional import click from flask import current_app from werkzeug.exceptions import NotFound +from constants.languages import languages from core.rag.datasource.vdb.vector_factory import Vector from core.rag.models.document import Document from extensions.ext_database import db @@ -17,6 +19,7 @@ from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment from models.dataset import Document as DatasetDocument from models.model import Account, App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation from models.provider import Provider, ProviderModel +from services.account_service import RegisterService, TenantService @click.command('reset-password', help='Reset the account password.') @@ -57,7 +60,7 @@ def reset_password(email, new_password, password_confirm): account.password = base64_password_hashed account.password_salt = base64_salt db.session.commit() - click.echo(click.style('Congratulations!, password has been reset.', fg='green')) + click.echo(click.style('Congratulations! Password has been reset.', fg='green')) @click.command('reset-email', help='Reset the account email.') @@ -501,6 +504,46 @@ def add_qdrant_doc_id_index(field: str): fg='green')) +@click.command('create-tenant', help='Create account and tenant.') +@click.option('--email', prompt=True, help='The email address of the tenant account.') +@click.option('--language', prompt=True, help='Account language, default: en-US.') +def create_tenant(email: str, language: Optional[str] = None): + """ + Create tenant account + """ + if not email: + click.echo(click.style('Sorry, email is required.', fg='red')) + return + + # Create account + email = email.strip() + + if '@' not in email: + click.echo(click.style('Sorry, invalid email address.', fg='red')) + return + + account_name = email.split('@')[0] + + if language not in languages: + language = 'en-US' + + # generate random password + new_password = secrets.token_urlsafe(16) + + # register account + account = RegisterService.register( + email=email, + name=account_name, + password=new_password, + language=language + ) + + TenantService.create_owner_tenant_if_not_exist(account) + + click.echo(click.style('Congratulations! Account and tenant created.\n' + 'Account: {}\nPassword: {}'.format(email, new_password), fg='green')) + + def register_commands(app): app.cli.add_command(reset_password) app.cli.add_command(reset_email) @@ -508,4 +551,5 @@ def register_commands(app): app.cli.add_command(vdb_migrate) app.cli.add_command(convert_to_agent_apps) app.cli.add_command(add_qdrant_doc_id_index) + app.cli.add_command(create_tenant) diff --git a/api/pyproject.toml b/api/pyproject.toml index f65795b4f6..a8920139c6 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -7,7 +7,6 @@ exclude = [ line-length = 120 [tool.ruff.lint] -ignore-init-module-imports = true select = [ "B", # flake8-bugbear rules "F", # pyflakes rules @@ -38,6 +37,7 @@ ignore = [ "B006", # mutable-argument-default "B007", # unused-loop-control-variable "B026", # star-arg-unpacking-after-keyword-arg + "B901", # return-in-generator "B904", # raise-without-from-inside-except "B905", # zip-without-explicit-strict ]