import base64 import json import logging import os import secrets import time from pathlib import Path from typing import Optional import click from flask import current_app from werkzeug.exceptions import NotFound from configs import dify_config from constants.languages import languages from core.rag.datasource.vdb.vector_factory import Vector from core.rag.datasource.vdb.vector_type import VectorType from core.rag.index_processor.constant.built_in_field import BuiltInField from core.rag.models.document import Document from events.app_event import app_was_created from extensions.ext_database import db from extensions.ext_redis import redis_client from extensions.ext_storage import storage from libs.helper import email as email_validate from libs.password import hash_password, password_pattern, valid_password from libs.rsa import generate_key_pair from models import Account, Tenant, TenantAccountJoin from models.account import TenantAccountJoinRole from models.dataset import Dataset, DatasetCollectionBinding, DatasetMetadata, DatasetMetadataBinding, DocumentSegment from models.dataset import Document as DatasetDocument from models.model import App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation, UploadFile from models.provider import Provider, ProviderModel from services.account_service import RegisterService, TenantService from services.clear_free_plan_tenant_expired_logs import ClearFreePlanTenantExpiredLogs from services.plugin.data_migration import PluginDataMigration from services.plugin.plugin_migration import PluginMigration @click.command("reset-password", help="Reset the account password.") @click.option("--email", prompt=True, help="Account email to reset password for") @click.option("--new-password", prompt=True, help="New password") @click.option("--password-confirm", prompt=True, help="Confirm new password") def reset_password(email, new_password, password_confirm): """ Reset password of owner account Only available in SELF_HOSTED mode """ if str(new_password).strip() != str(password_confirm).strip(): click.echo(click.style("Passwords do not match.", fg="red")) return account = db.session.query(Account).filter(Account.email == email).one_or_none() if not account: click.echo(click.style("Account not found for email: {}".format(email), fg="red")) return try: valid_password(new_password) except: click.echo(click.style("Invalid password. Must match {}".format(password_pattern), fg="red")) return # generate password salt salt = secrets.token_bytes(16) base64_salt = base64.b64encode(salt).decode() # encrypt password with salt password_hashed = hash_password(new_password, salt) base64_password_hashed = base64.b64encode(password_hashed).decode() account.password = base64_password_hashed account.password_salt = base64_salt db.session.commit() click.echo(click.style("Password reset successfully.", fg="green")) @click.command("reset-email", help="Reset the account email.") @click.option("--email", prompt=True, help="Current account email") @click.option("--new-email", prompt=True, help="New email") @click.option("--email-confirm", prompt=True, help="Confirm new email") def reset_email(email, new_email, email_confirm): """ Replace account email :return: """ if str(new_email).strip() != str(email_confirm).strip(): click.echo(click.style("New emails do not match.", fg="red")) return account = db.session.query(Account).filter(Account.email == email).one_or_none() if not account: click.echo(click.style("Account not found for email: {}".format(email), fg="red")) return try: email_validate(new_email) except: click.echo(click.style("Invalid email: {}".format(new_email), fg="red")) return account.email = new_email db.session.commit() click.echo(click.style("Email updated successfully.", fg="green")) @click.command( "reset-encrypt-key-pair", help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. " "After the reset, all LLM credentials will become invalid, " "requiring re-entry." "Only support SELF_HOSTED mode.", ) @click.confirmation_option( prompt=click.style( "Are you sure you want to reset encrypt key pair? This operation cannot be rolled back!", fg="red", ) ) def reset_encrypt_key_pair(): """ Reset the encrypted key pair of workspace for encrypt LLM credentials. After the reset, all LLM credentials will become invalid, requiring re-entry. Only support SELF_HOSTED mode. """ if dify_config.EDITION != "SELF_HOSTED": click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red")) return tenants = db.session.query(Tenant).all() for tenant in tenants: if not tenant: click.echo(click.style("No workspaces found. Run /install first.", fg="red")) return tenant.encrypt_public_key = generate_key_pair(tenant.id) 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.command("vdb-migrate", help="Migrate vector db.") @click.option( "--scope", default="all", prompt=False, help="The scope of vector database to migrate, Default is All.", ) def vdb_migrate(scope: str): if scope in {"knowledge", "all"}: migrate_knowledge_vector_database() if scope in {"annotation", "all"}: migrate_annotation_vector_database() def migrate_annotation_vector_database(): """ Migrate annotation datas to target vector database . """ click.echo(click.style("Starting annotation data migration.", fg="green")) create_count = 0 skipped_count = 0 total_count = 0 page = 1 while True: try: # get apps info per_page = 50 apps = ( db.session.query(App) .filter(App.status == "normal") .order_by(App.created_at.desc()) .limit(per_page) .offset((page - 1) * per_page) .all() ) if not apps: break except NotFound: break page += 1 for app in apps: total_count = total_count + 1 click.echo( f"Processing the {total_count} app {app.id}. " + f"{create_count} created, {skipped_count} skipped." ) try: click.echo("Creating app annotation index: {}".format(app.id)) app_annotation_setting = ( db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app.id).first() ) if not app_annotation_setting: skipped_count = skipped_count + 1 click.echo("App annotation setting disabled: {}".format(app.id)) continue # get dataset_collection_binding info dataset_collection_binding = ( db.session.query(DatasetCollectionBinding) .filter(DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id) .first() ) if not dataset_collection_binding: click.echo("App annotation collection binding not found: {}".format(app.id)) continue annotations = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app.id).all() dataset = Dataset( id=app.id, tenant_id=app.tenant_id, indexing_technique="high_quality", embedding_model_provider=dataset_collection_binding.provider_name, embedding_model=dataset_collection_binding.model_name, collection_binding_id=dataset_collection_binding.id, ) documents = [] if annotations: for annotation in annotations: document = Document( page_content=annotation.question, metadata={ "annotation_id": annotation.id, "app_id": app.id, "doc_id": annotation.id, }, ) documents.append(document) vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"]) click.echo(f"Migrating annotations for app: {app.id}.") try: vector.delete() click.echo(click.style(f"Deleted vector index for app {app.id}.", fg="green")) except Exception as e: click.echo(click.style(f"Failed to delete vector index for app {app.id}.", fg="red")) raise e if documents: try: click.echo( click.style( f"Creating vector index with {len(documents)} annotations for app {app.id}.", fg="green", ) ) vector.create(documents) click.echo(click.style(f"Created vector index for app {app.id}.", fg="green")) except Exception as e: click.echo( click.style( f"Failed to created vector index for app {app.id}.", fg="red", ) ) raise e click.echo(f"Successfully migrated app annotation {app.id}.") create_count += 1 except Exception as e: click.echo( click.style( "Error creating app annotation index: {} {}".format(e.__class__.__name__, str(e)), fg="red", ) ) continue click.echo( click.style( f"Migration complete. Created {create_count} app annotation indexes. Skipped {skipped_count} apps.", fg="green", ) ) def migrate_knowledge_vector_database(): """ Migrate vector database datas to target vector database . """ click.echo(click.style("Starting vector database migration.", fg="green")) create_count = 0 skipped_count = 0 total_count = 0 vector_type = dify_config.VECTOR_STORE upper_collection_vector_types = { VectorType.MILVUS, VectorType.PGVECTOR, VectorType.VASTBASE, VectorType.RELYT, VectorType.WEAVIATE, VectorType.ORACLE, VectorType.ELASTICSEARCH, VectorType.OPENGAUSS, VectorType.TABLESTORE, } lower_collection_vector_types = { VectorType.ANALYTICDB, VectorType.CHROMA, VectorType.MYSCALE, VectorType.PGVECTO_RS, VectorType.TIDB_VECTOR, VectorType.OPENSEARCH, VectorType.TENCENT, VectorType.BAIDU, VectorType.VIKINGDB, VectorType.UPSTASH, VectorType.COUCHBASE, VectorType.OCEANBASE, } page = 1 while True: try: datasets = ( Dataset.query.filter(Dataset.indexing_technique == "high_quality") .order_by(Dataset.created_at.desc()) .paginate(page=page, per_page=50) ) except NotFound: break page += 1 for dataset in datasets: total_count = total_count + 1 click.echo( f"Processing the {total_count} dataset {dataset.id}. {create_count} created, {skipped_count} skipped." ) try: click.echo("Creating dataset vector database index: {}".format(dataset.id)) if dataset.index_struct_dict: if dataset.index_struct_dict["type"] == vector_type: skipped_count = skipped_count + 1 continue collection_name = "" dataset_id = dataset.id if vector_type in upper_collection_vector_types: collection_name = Dataset.gen_collection_name_by_id(dataset_id) elif vector_type == VectorType.QDRANT: if dataset.collection_binding_id: dataset_collection_binding = ( db.session.query(DatasetCollectionBinding) .filter(DatasetCollectionBinding.id == dataset.collection_binding_id) .one_or_none() ) if dataset_collection_binding: collection_name = dataset_collection_binding.collection_name else: raise ValueError("Dataset Collection Binding not found") else: collection_name = Dataset.gen_collection_name_by_id(dataset_id) elif vector_type in lower_collection_vector_types: collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower() else: raise ValueError(f"Vector store {vector_type} is not supported.") index_struct_dict = { "type": vector_type, "vector_store": {"class_prefix": collection_name}, } dataset.index_struct = json.dumps(index_struct_dict) vector = Vector(dataset) click.echo(f"Migrating dataset {dataset.id}.") try: vector.delete() click.echo( click.style( f"Deleted vector index {collection_name} for dataset {dataset.id}.", fg="green", ) ) except Exception as e: click.echo( click.style( f"Failed to delete vector index {collection_name} for dataset {dataset.id}.", fg="red", ) ) raise e dataset_documents = ( db.session.query(DatasetDocument) .filter( DatasetDocument.dataset_id == dataset.id, DatasetDocument.indexing_status == "completed", DatasetDocument.enabled == True, DatasetDocument.archived == False, ) .all() ) documents = [] segments_count = 0 for dataset_document in dataset_documents: segments = ( db.session.query(DocumentSegment) .filter( DocumentSegment.document_id == dataset_document.id, DocumentSegment.status == "completed", DocumentSegment.enabled == True, ) .all() ) for segment in segments: document = Document( page_content=segment.content, metadata={ "doc_id": segment.index_node_id, "doc_hash": segment.index_node_hash, "document_id": segment.document_id, "dataset_id": segment.dataset_id, }, ) documents.append(document) segments_count = segments_count + 1 if documents: try: click.echo( click.style( f"Creating vector index with {len(documents)} documents of {segments_count}" f" segments for dataset {dataset.id}.", fg="green", ) ) vector.create(documents) click.echo( click.style( f"Created vector index for dataset {dataset.id}.", fg="green", ) ) except Exception as e: click.echo( click.style( f"Failed to created vector index for dataset {dataset.id}.", fg="red", ) ) raise e db.session.add(dataset) db.session.commit() click.echo(f"Successfully migrated dataset {dataset.id}.") create_count += 1 except Exception as e: db.session.rollback() click.echo( click.style( "Error creating dataset index: {} {}".format(e.__class__.__name__, str(e)), fg="red", ) ) continue click.echo( click.style( f"Migration complete. Created {create_count} dataset indexes. Skipped {skipped_count} datasets.", fg="green", ) ) @click.command("convert-to-agent-apps", help="Convert Agent Assistant to Agent App.") def convert_to_agent_apps(): """ Convert Agent Assistant to Agent App. """ click.echo(click.style("Starting convert to agent apps.", fg="green")) proceeded_app_ids = [] while True: # fetch first 1000 apps sql_query = """SELECT a.id AS id FROM apps a INNER JOIN app_model_configs am ON a.app_model_config_id=am.id WHERE a.mode = 'chat' AND am.agent_mode is not null AND ( am.agent_mode like '%"strategy": "function_call"%' OR am.agent_mode like '%"strategy": "react"%' ) AND ( am.agent_mode like '{"enabled": true%' OR am.agent_mode like '{"max_iteration": %' ) ORDER BY a.created_at DESC LIMIT 1000 """ with db.engine.begin() as conn: rs = conn.execute(db.text(sql_query)) apps = [] for i in rs: app_id = str(i.id) if app_id not in proceeded_app_ids: proceeded_app_ids.append(app_id) app = db.session.query(App).filter(App.id == app_id).first() if app is not None: apps.append(app) if len(apps) == 0: break for app in apps: click.echo("Converting app: {}".format(app.id)) try: app.mode = AppMode.AGENT_CHAT.value db.session.commit() # update conversation mode to agent db.session.query(Conversation).filter(Conversation.app_id == app.id).update( {Conversation.mode: AppMode.AGENT_CHAT.value} ) db.session.commit() click.echo(click.style("Converted app: {}".format(app.id), fg="green")) except Exception as e: click.echo( click.style( "Convert app error: {} {}".format(e.__class__.__name__, str(e)), fg="red", ) ) click.echo( click.style( "Conversion complete. Converted {} agent apps.".format(len(proceeded_app_ids)), fg="green", ) ) @click.command("add-qdrant-index", help="Add Qdrant index.") @click.option( "--field", default="metadata.doc_id", prompt=False, help="Index field , default is metadata.doc_id.", ) def add_qdrant_index(field: str): click.echo(click.style("Starting Qdrant index creation.", fg="green")) create_count = 0 try: bindings = db.session.query(DatasetCollectionBinding).all() if not bindings: click.echo(click.style("No dataset collection bindings found.", fg="red")) return import qdrant_client from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http.models import PayloadSchemaType from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig for binding in bindings: if dify_config.QDRANT_URL is None: raise ValueError("Qdrant URL is required.") qdrant_config = QdrantConfig( endpoint=dify_config.QDRANT_URL, api_key=dify_config.QDRANT_API_KEY, root_path=current_app.root_path, timeout=dify_config.QDRANT_CLIENT_TIMEOUT, grpc_port=dify_config.QDRANT_GRPC_PORT, prefer_grpc=dify_config.QDRANT_GRPC_ENABLED, ) try: client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params()) # type: ignore # create payload index client.create_payload_index( binding.collection_name, field, field_schema=PayloadSchemaType.KEYWORD, ) create_count += 1 except UnexpectedResponse as e: # Collection does not exist, so return if e.status_code == 404: click.echo( click.style( f"Collection not found: {binding.collection_name}.", fg="red", ) ) continue # Some other error occurred, so re-raise the exception else: click.echo( click.style( f"Failed to create Qdrant index for collection: {binding.collection_name}.", fg="red", ) ) except Exception: click.echo(click.style("Failed to create Qdrant client.", fg="red")) click.echo( click.style( f"Index creation complete. Created {create_count} collection indexes.", fg="green", ) ) @click.command("old-metadata-migration", help="Old metadata migration.") def old_metadata_migration(): """ Old metadata migration. """ click.echo(click.style("Starting old metadata migration.", fg="green")) page = 1 while True: try: documents = ( DatasetDocument.query.filter(DatasetDocument.doc_metadata is not None) .order_by(DatasetDocument.created_at.desc()) .paginate(page=page, per_page=50) ) except NotFound: break if not documents: break for document in documents: if document.doc_metadata: doc_metadata = document.doc_metadata for key, value in doc_metadata.items(): for field in BuiltInField: if field.value == key: break else: dataset_metadata = ( db.session.query(DatasetMetadata) .filter( DatasetMetadata.dataset_id == document.dataset_id, DatasetMetadata.name == key, ) .first() ) if not dataset_metadata: dataset_metadata = DatasetMetadata( tenant_id=document.tenant_id, dataset_id=document.dataset_id, name=key, type="string", created_by=document.created_by, ) db.session.add(dataset_metadata) db.session.flush() dataset_metadata_binding = DatasetMetadataBinding( tenant_id=document.tenant_id, dataset_id=document.dataset_id, metadata_id=dataset_metadata.id, document_id=document.id, created_by=document.created_by, ) db.session.add(dataset_metadata_binding) else: dataset_metadata_binding = DatasetMetadataBinding.query.filter( DatasetMetadataBinding.dataset_id == document.dataset_id, DatasetMetadataBinding.document_id == document.id, DatasetMetadataBinding.metadata_id == dataset_metadata.id, ).first() if not dataset_metadata_binding: dataset_metadata_binding = DatasetMetadataBinding( tenant_id=document.tenant_id, dataset_id=document.dataset_id, metadata_id=dataset_metadata.id, document_id=document.id, created_by=document.created_by, ) db.session.add(dataset_metadata_binding) db.session.commit() page += 1 click.echo(click.style("Old metadata migration completed.", fg="green")) @click.command("create-tenant", help="Create account and tenant.") @click.option("--email", prompt=True, help="Tenant account email.") @click.option("--name", prompt=True, help="Workspace name.") @click.option("--language", prompt=True, help="Account language, default: en-US.") def create_tenant(email: str, language: Optional[str] = None, name: Optional[str] = None): """ Create tenant account """ if not email: click.echo(click.style("Email is required.", fg="red")) return # Create account email = email.strip() if "@" not in email: click.echo(click.style("Invalid email address.", fg="red")) return account_name = email.split("@")[0] if language not in languages: language = "en-US" # Validates name encoding for non-Latin characters. name = name.strip().encode("utf-8").decode("utf-8") if name else None # generate random password new_password = secrets.token_urlsafe(16) # register account account = RegisterService.register( email=email, name=account_name, password=new_password, language=language, create_workspace_required=False, ) TenantService.create_owner_tenant_if_not_exist(account, name) click.echo( click.style( "Account and tenant created.\nAccount: {}\nPassword: {}".format(email, new_password), fg="green", ) ) @click.command("upgrade-db", help="Upgrade the database") def upgrade_db(): click.echo("Preparing database migration...") lock = redis_client.lock(name="db_upgrade_lock", timeout=60) if lock.acquire(blocking=False): try: click.echo(click.style("Starting database migration.", fg="green")) # run db migration import flask_migrate flask_migrate.upgrade() click.echo(click.style("Database migration successful!", fg="green")) except Exception: logging.exception("Failed to execute database migration") finally: lock.release() else: click.echo("Database migration skipped") @click.command("fix-app-site-missing", help="Fix app related site missing issue.") def fix_app_site_missing(): """ Fix app related site missing issue. """ click.echo(click.style("Starting fix for missing app-related sites.", fg="green")) failed_app_ids = [] while True: sql = """select apps.id as id from apps left join sites on sites.app_id=apps.id where sites.id is null limit 1000""" with db.engine.begin() as conn: rs = conn.execute(db.text(sql)) processed_count = 0 for i in rs: processed_count += 1 app_id = str(i.id) if app_id in failed_app_ids: continue try: app = db.session.query(App).filter(App.id == app_id).first() if not app: print(f"App {app_id} not found") continue tenant = app.tenant if tenant: accounts = tenant.get_accounts() if not accounts: print("Fix failed for app {}".format(app.id)) continue account = accounts[0] print("Fixing missing site for app {}".format(app.id)) app_was_created.send(app, account=account) except Exception: failed_app_ids.append(app_id) click.echo( click.style( "Failed to fix missing site for app {}".format(app_id), fg="red", ) ) logging.exception(f"Failed to fix app related site missing issue, app_id: {app_id}") continue if not processed_count: break click.echo(click.style("Fix for missing app-related sites completed successfully!", fg="green")) @click.command( "create-admin-account", help="Create or update an admin account for an organization with a phone number or email.", ) @click.option("--name", prompt=True, help="Admin account name") @click.option("--login-id", prompt=True, help="Admin account phone number or email") @click.option( "--login-id-type", prompt=True, type=click.Choice(["phone", "email"]), help="Type of login ID (phone or email)", ) @click.option("--organization-id", required=True, help="Organization ID") def create_admin_account(name: str, login_id: str, login_id_type: str, organization_id: str): """ Create or update an admin account with a phone number or email for a specific organization. This command will create a new account if the login ID doesn't exist, or update an existing account with the specified admin role. """ try: # Check if organization exists from models.organization import Organization, OrganizationMember, OrganizationRole organization = db.session.query(Organization).filter(Organization.id == organization_id).first() if not organization: click.echo(click.style(f"Organization with ID {organization_id} not found.", fg="red")) return # Get tenant from organization tenant = db.session.query(Tenant).filter(Tenant.id == organization.tenant_id).first() if not tenant: click.echo(click.style(f"Tenant for organization {organization_id} not found.", fg="red")) return # Check if account exists with this login ID account = None if login_id_type == "phone": account = db.session.query(Account).filter(Account.phone == login_id).first() else: # email account = db.session.query(Account).filter(Account.email == login_id).first() if account: click.echo(f"Account with {login_id_type} {login_id} already exists. Updating account...") # Update account account.name = name account.current_organization_id = organization_id db.session.commit() else: click.echo(f"Creating new account with {login_id_type} {login_id}...") # Create new account if login_id_type == "phone": account = Account( name=name, email=f"{login_id}@qingsu.chat", # Use phone as part of email phone=login_id, interface_language=languages[0], status="active", current_organization_id=organization_id, ) else: # email account = Account( name=name, email=login_id, interface_language=languages[0], status="active", current_organization_id=organization_id, ) db.session.add(account) db.session.commit() # Check if account is already a member of the tenant ta_join = ( db.session.query(TenantAccountJoin) .filter( TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id, ) .first() ) if not ta_join: # Add account to tenant with end_user role (organization role will control admin access) ta_join = TenantAccountJoin( tenant_id=tenant.id, account_id=account.id, role=TenantAccountJoinRole.END_USER.value, ) db.session.add(ta_join) click.echo(f"Added account to tenant {tenant.name}") # Check if account is already a member of the organization org_member = ( db.session.query(OrganizationMember) .filter( OrganizationMember.organization_id == organization_id, OrganizationMember.account_id == account.id, ) .first() ) if org_member: # Update role to admin org_member.role = OrganizationRole.ADMIN click.echo(f"Updated account role to {OrganizationRole.ADMIN} in organization {organization.name}") else: # Add account to organization with admin role org_member = OrganizationMember( organization_id=organization_id, account_id=account.id, role=OrganizationRole.ADMIN, is_default=True, created_by=account.id, ) db.session.add(org_member) click.echo(f"Added account to organization {organization.name} with role {OrganizationRole.ADMIN}") db.session.commit() click.echo( click.style( f"Successfully {'updated' if account else 'created'} admin account with {login_id_type}.", fg="green", ) ) click.echo(f"Name: {name}") click.echo(f"{login_id_type.capitalize()}: {login_id}") click.echo(f"Organization: {organization.name} (ID: {organization.id})") except Exception as e: db.session.rollback() click.echo(click.style(f"Error: {str(e)}", fg="red")) @click.command("create-organization", help="Create a new organization for multi-school support.") @click.option("--tenant-id", required=True, help="ID of the tenant that owns this organization") @click.option("--name", required=True, help="Name of the organization") @click.option("--code", required=True, help="Unique code for the organization") @click.option( "--type", "org_type", default="school", type=click.Choice(["school", "university", "company", "organization"]), help="Type of organization", ) @click.option("--description", default="", help="Description of the organization") @click.option("--email-domains", default="", help="Comma-separated list of allowed email domains") @click.option("--created-by", required=True, help="Account ID of the creator") def create_organization_cmd(tenant_id, name, code, org_type, description, email_domains, created_by): """Create a new organization under a tenant for multi-school support""" try: # Check if code already exists from models.organization import Organization existing = db.session.query(Organization).filter(Organization.code == code).first() if existing: click.echo(f"Error: Organization with code '{code}' already exists") return # Check if creator account exists creator = db.session.query(Account).filter(Account.id == created_by).first() if not creator: click.echo(f"Error: Creator account with ID '{created_by}' not found") return # Parse email domains allowed_domains = [d.strip() for d in email_domains.split(",") if d.strip()] # Create settings settings = {"allowed_email_domains": allowed_domains} # Create organization organization = Organization( tenant_id=tenant_id, name=name, code=code, type=org_type, description=description, settings=json.dumps(settings), status="active", created_by=created_by, ) db.session.add(organization) db.session.commit() click.echo(f"Organization '{name}' (ID: {organization.id}) created successfully") except Exception as e: db.session.rollback() click.echo(f"Error creating organization: {str(e)}") @click.command("update-organization", help="Update an existing organization.") @click.option("--id", "org_id", required=True, help="ID of the organization to update") @click.option("--name", help="New name for the organization") @click.option("--description", help="New description") @click.option("--email-domains", help="Comma-separated list of allowed email domains") @click.option("--status", type=click.Choice(["active", "inactive"]), help="Organization status") def update_organization_cmd(org_id, name, description, email_domains, status): """Update an existing organization's configuration""" try: from models.organization import Organization organization = db.session.query(Organization).filter(Organization.id == org_id).first() if not organization: click.echo(f"Error: Organization with ID '{org_id}' not found") return if name: organization.name = name if description: organization.description = description if status: organization.status = status if email_domains is not None: settings = organization.settings_dict allowed_domains = [d.strip() for d in email_domains.split(",") if d.strip()] settings["allowed_email_domains"] = allowed_domains organization.settings_dict = settings db.session.commit() click.echo(f"Organization '{organization.name}' updated successfully") except Exception as e: db.session.rollback() click.echo(f"Error updating organization: {str(e)}") @click.command("list-organizations", help="List all organizations.") @click.option("--tenant-id", help="Filter by tenant ID") def list_organizations_cmd(tenant_id): """List all organizations with optional tenant filtering""" try: from models.organization import Organization, OrganizationMember, OrganizationRole query = db.session.query(Organization) if tenant_id: query = query.filter(Organization.tenant_id == tenant_id) organizations = query.all() if not organizations: click.echo("No organizations found") return # Prepare a dictionary to store admin phones for each organization admin_phones_by_org = {} for org in organizations: # Query for admin accounts in this organization admin_members = ( db.session.query(OrganizationMember, Account) .join(Account, OrganizationMember.account_id == Account.id) .filter( OrganizationMember.organization_id == org.id, OrganizationMember.role == OrganizationRole.ADMIN, ) .all() ) # Collect phone numbers phones = [] for member, account in admin_members: if account.phone: phones.append(account.phone) admin_phones_by_org[org.id] = ", ".join(phones) if phones else "None" # Create a header with fixed width that doesn't exceed line length limit click.echo( f"{'ID':<36} | {'Code':<10} | {'Name':<30} | {'Type':<12} | {'Status':<8} | {'Admin Phones':<15} | " f"{'Email Domains'}" ) click.echo("-" * 140) for org in organizations: email_domains = ", ".join(org.allowed_email_domains) admin_phones = admin_phones_by_org[org.id] # Split the long line to avoid exceeding line length limit click.echo( f"{org.id:<36} | {org.code:<10} | {org.name:<30} | {org.type:<12} | {org.status:<8} | " f"{admin_phones:<15} | {email_domains}" ) except Exception as e: click.echo(f"Error listing organizations: {str(e)}") @click.command("show-organization", help="Show details of a specific organization.") @click.option("--id", "org_id", required=True, help="ID of the organization to show") def show_organization_cmd(org_id): """Show detailed information about a specific organization""" try: from models.organization import Organization organization = db.session.query(Organization).filter(Organization.id == org_id).first() if not organization: click.echo(f"Error: Organization with ID '{org_id}' not found") return click.echo(f"ID: {organization.id}") click.echo(f"Tenant ID: {organization.tenant_id}") click.echo(f"Name: {organization.name}") click.echo(f"Code: {organization.code}") click.echo(f"Type: {organization.type}") click.echo(f"Description: {organization.description or ''}") click.echo(f"Status: {organization.status}") click.echo(f"Email Domains: {', '.join(organization.allowed_email_domains)}") click.echo(f"Created At: {organization.created_at}") click.echo(f"Updated At: {organization.updated_at}") except Exception as e: click.echo(f"Error showing organization: {str(e)}") @click.command( "add-account-to-organization", help="Add an account to an organization with a specific role.", ) @click.option("--org-id", required=True, help="ID of the organization") @click.option("--account-id", required=True, help="ID of the account to add") @click.option( "--role", required=True, type=click.Choice(["admin", "teacher", "student", "staff", "manager", "employee", "guest"]), help="Role in the organization", ) @click.option("--department", help="Department within the organization") @click.option("--title", help="Job title or position") @click.option("--is-default", is_flag=True, help="Set as the account's default organization") def add_account_to_organization_cmd(org_id, account_id, role, department, title, is_default): """Add an account to an organization with appropriate role and metadata""" try: from models.organization import Organization, OrganizationMember # Check if organization exists organization = db.session.query(Organization).filter(Organization.id == org_id).first() if not organization: click.echo(f"Error: Organization with ID '{org_id}' not found") return # Check if account exists account = db.session.query(Account).filter(Account.id == account_id).first() if not account: click.echo(f"Error: Account with ID '{account_id}' not found") return # Check if membership already exists existing = ( db.session.query(OrganizationMember) .filter( OrganizationMember.organization_id == org_id, OrganizationMember.account_id == account_id, ) .first() ) if existing: click.echo("Account is already a member of this organization. Updating role and metadata.") existing.role = role existing.department = department existing.title = title existing.is_default = is_default else: # Create new membership with meta_data instead of metadata (reserved word) member = OrganizationMember( organization_id=org_id, account_id=account_id, role=role, department=department, title=title, is_default=is_default, created_by=account_id, # Use meta_data instead of metadata as it's a reserved word in SQLAlchemy meta_data=json.dumps({}), ) db.session.add(member) # If set as default, update the account's current_organization_id if is_default: account.current_organization_id = org_id db.session.commit() click.echo( f"Account successfully {'added to' if not existing else 'updated in'} organization with role '{role}'" ) except Exception as e: db.session.rollback() click.echo(f"Error adding account to organization: {str(e)}") @click.command( "upload-private-key-file-to-cloud-storage", help="upload private key file to cloud storage", ) @click.option("--tenant_id", prompt=False, help="tenant_id") def upload_private_key_file_cloud_storage(tenant_id: Optional[str] = None): """ upload private.pem to cloud storage """ click.echo( click.style( "Start upload private.pem to cloud storage", fg="green", ) ) if not tenant_id: click.echo( click.style( "Warning: did not provide an tenant_id, it will be auto queried in the database", fg="yellow", ) ) tenants_list: list[Tenant] = Tenant.query.all() tenants = [item.id for item in tenants_list] else: tenants = [ tenant_id, ] for tenant_id in tenants: click.echo( click.style( f"Current tenant_id is: {tenant_id}", fg="green", ) ) file_key = f"privkeys/{tenant_id}/private.pem" file_content = Path(f"{os.environ.get('STORAGE_LOCAL_PATH', 'storage')}/{file_key}").read_bytes() storage.save(filename=file_key, data=file_content) click.echo( click.style( f"Congratulations! file uploaded. file.key: {file_key}", fg="green", ) ) @click.command("upload-local-files-to-cloud-storage", help="upload local files to cloud storage") def upload_local_files_to_cloud_storage(): """ upload local files to cloud storage """ click.echo( click.style( "Start upload local files to cloud storage", fg="green", ) ) total_count = UploadFile.query.filter_by(storage_type="local").count() click.echo(click.style(f"Total files to process: {total_count}", fg="green")) batch_size = 100 processed_count = 0 while processed_count < total_count: files: list[UploadFile] = UploadFile.query.filter_by(storage_type="local").limit(batch_size).all() for file in files: target_filepath = f"{os.environ.get('STORAGE_LOCAL_PATH', 'storage')}/{file.key}" # if the file exists if not os.path.exists(target_filepath): click.echo( click.style( f"Warning! file not exist. filepath: {target_filepath}, ignore this, continue", fg="yellow", ) ) processed_count += 1 if processed_count % 10 == 0 or processed_count == total_count: click.echo( click.style( f"Processed {processed_count}/{total_count} files\n", fg="blue", ) ) continue # Upload to cloud storage file_content = Path(target_filepath).read_bytes() storage.save(filename=file.key, data=file_content) click.echo( click.style( f"File uploaded. file.key: {file.key}", fg="green", ) ) # Update database record try: file.storage_type = os.environ["STORAGE_TYPE"] db.session.commit() click.echo( click.style( f"file.storage_type updated to database. file.key: {file.key}", fg="green", ) ) except Exception as e: click.echo(click.style(f"An error occurred: {str(e)}", fg="red")) db.session.rollback() processed_count += 1 if processed_count % 10 == 0 or processed_count == total_count: click.echo(click.style(f"Processed {processed_count}/{total_count} files\n", fg="blue")) time.sleep(3) click.echo( click.style( "Congratulations! finish files uploaded.", fg="green", ) ) @click.command("migrate-data-for-plugin", help="Migrate data for plugin.") def migrate_data_for_plugin(): """ Migrate data for plugin. """ click.echo(click.style("Starting migrate data for plugin.", fg="white")) PluginDataMigration.migrate() click.echo(click.style("Migrate data for plugin completed.", fg="green")) @click.command("extract-plugins", help="Extract plugins.") @click.option( "--output_file", prompt=True, help="The file to store the extracted plugins.", default="plugins.jsonl", ) @click.option( "--workers", prompt=True, help="The number of workers to extract plugins.", default=10, ) def extract_plugins(output_file: str, workers: int): """ Extract plugins. """ click.echo(click.style("Starting extract plugins.", fg="white")) PluginMigration.extract_plugins(output_file, workers) click.echo(click.style("Extract plugins completed.", fg="green")) @click.command("extract-unique-identifiers", help="Extract unique identifiers.") @click.option( "--output_file", prompt=True, help="The file to store the extracted unique identifiers.", default="unique_identifiers.json", ) @click.option( "--input_file", prompt=True, help="The file to store the extracted unique identifiers.", default="plugins.jsonl", ) def extract_unique_plugins(output_file: str, input_file: str): """ Extract unique plugins. """ click.echo(click.style("Starting extract unique plugins.", fg="white")) PluginMigration.extract_unique_plugins_to_file(input_file, output_file) click.echo(click.style("Extract unique plugins completed.", fg="green")) @click.command("install-plugins", help="Install plugins.") @click.option( "--input_file", prompt=True, help="The file to store the extracted unique identifiers.", default="plugins.jsonl", ) @click.option( "--output_file", prompt=True, help="The file to store the installed plugins.", default="installed_plugins.jsonl", ) @click.option( "--workers", prompt=True, help="The number of workers to install plugins.", default=100, ) def install_plugins(input_file: str, output_file: str, workers: int): """ Install plugins. """ click.echo(click.style("Starting install plugins.", fg="white")) PluginMigration.install_plugins(input_file, output_file, workers) click.echo(click.style("Install plugins completed.", fg="green")) @click.command("clear-free-plan-tenant-expired-logs", help="Clear free plan tenant expired logs.") @click.option( "--days", prompt=True, help="The days to clear free plan tenant expired logs.", default=30, ) @click.option( "--batch", prompt=True, help="The batch size to clear free plan tenant expired logs.", default=100, ) @click.option( "--tenant_ids", prompt=True, multiple=True, help="The tenant ids to clear free plan tenant expired logs.", ) def clear_free_plan_tenant_expired_logs(days: int, batch: int, tenant_ids: list[str]): """ Clear free plan tenant expired logs. """ click.echo(click.style("Starting clear free plan tenant expired logs.", fg="white")) ClearFreePlanTenantExpiredLogs.process(days, batch, tenant_ids) click.echo(click.style("Clear free plan tenant expired logs completed.", fg="green")) @click.option( "-f", "--force", is_flag=True, help="Skip user confirmation and force the command to execute.", ) @click.command("clear-orphaned-file-records", help="Clear orphaned file records.") def clear_orphaned_file_records(force: bool): """ Clear orphaned file records in the database. """ # define tables and columns to process files_tables = [ {"table": "upload_files", "id_column": "id", "key_column": "key"}, {"table": "tool_files", "id_column": "id", "key_column": "file_key"}, ] ids_tables = [ {"type": "uuid", "table": "message_files", "column": "upload_file_id"}, {"type": "text", "table": "documents", "column": "data_source_info"}, {"type": "text", "table": "document_segments", "column": "content"}, {"type": "text", "table": "messages", "column": "answer"}, {"type": "text", "table": "workflow_node_executions", "column": "inputs"}, {"type": "text", "table": "workflow_node_executions", "column": "process_data"}, {"type": "text", "table": "workflow_node_executions", "column": "outputs"}, {"type": "text", "table": "conversations", "column": "introduction"}, {"type": "text", "table": "conversations", "column": "system_instruction"}, {"type": "json", "table": "messages", "column": "inputs"}, {"type": "json", "table": "messages", "column": "message"}, ] # notify user and ask for confirmation click.echo( click.style( "This command will first find and delete orphaned file records from the message_files table,", fg="yellow", ) ) click.echo( click.style( "and then it will find and delete orphaned file records in the following tables:", fg="yellow", ) ) for files_table in files_tables: click.echo(click.style(f"- {files_table['table']}", fg="yellow")) click.echo( click.style( "The following tables and columns will be scanned to find orphaned file records:", fg="yellow", ) ) for ids_table in ids_tables: click.echo(click.style(f"- {ids_table['table']} ({ids_table['column']})", fg="yellow")) click.echo("") click.echo(click.style("!!! USE WITH CAUTION !!!", fg="red")) click.echo( click.style( ( "Since not all patterns have been fully tested, " "please note that this command may delete unintended file records." ), fg="yellow", ) ) click.echo( click.style( "This cannot be undone. Please make sure to back up your database before proceeding.", fg="yellow", ) ) click.echo( click.style( ( "It is also recommended to run this during the maintenance window, " "as this may cause high load on your instance." ), fg="yellow", ) ) if not force: click.confirm("Do you want to proceed?", abort=True) # start the cleanup process click.echo(click.style("Starting orphaned file records cleanup.", fg="white")) # clean up the orphaned records in the message_files table where message_id doesn't exist in messages table try: click.echo( click.style( "- Listing message_files records where message_id doesn't exist in messages table", fg="white", ) ) query = ( "SELECT mf.id, mf.message_id " "FROM message_files mf LEFT JOIN messages m ON mf.message_id = m.id " "WHERE m.id IS NULL" ) orphaned_message_files = [] with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: orphaned_message_files.append({"id": str(i[0]), "message_id": str(i[1])}) if orphaned_message_files: click.echo( click.style( f"Found {len(orphaned_message_files)} orphaned message_files records:", fg="white", ) ) for record in orphaned_message_files: click.echo( click.style( f" - id: {record['id']}, message_id: {record['message_id']}", fg="black", ) ) if not force: click.confirm( ( f"Do you want to proceed " f"to delete all {len(orphaned_message_files)} orphaned message_files records?" ), abort=True, ) click.echo(click.style("- Deleting orphaned message_files records", fg="white")) query = "DELETE FROM message_files WHERE id IN :ids" with db.engine.begin() as conn: conn.execute( db.text(query), {"ids": tuple([record["id"] for record in orphaned_message_files])}, ) click.echo( click.style( f"Removed {len(orphaned_message_files)} orphaned message_files records.", fg="green", ) ) else: click.echo( click.style( "No orphaned message_files records found. There is nothing to delete.", fg="green", ) ) except Exception as e: click.echo(click.style(f"Error deleting orphaned message_files records: {str(e)}", fg="red")) # clean up the orphaned records in the rest of the *_files tables try: # fetch file id and keys from each table all_files_in_tables = [] for files_table in files_tables: click.echo( click.style( f"- Listing file records in table {files_table['table']}", fg="white", ) ) query = f"SELECT {files_table['id_column']}, {files_table['key_column']} FROM {files_table['table']}" with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: all_files_in_tables.append({"table": files_table["table"], "id": str(i[0]), "key": i[1]}) click.echo(click.style(f"Found {len(all_files_in_tables)} files in tables.", fg="white")) # fetch referred table and columns guid_regexp = "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" all_ids_in_tables = [] for ids_table in ids_tables: query = "" if ids_table["type"] == "uuid": click.echo( click.style( f"- Listing file ids in column {ids_table['column']} in table {ids_table['table']}", fg="white", ) ) query = ( f"SELECT {ids_table['column']} FROM {ids_table['table']} WHERE {ids_table['column']} IS NOT NULL" ) with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: all_ids_in_tables.append({"table": ids_table["table"], "id": str(i[0])}) elif ids_table["type"] == "text": click.echo( click.style( f"- Listing file-id-like strings in column {ids_table['column']} in table {ids_table['table']}", fg="white", ) ) query = ( f"SELECT regexp_matches({ids_table['column']}, '{guid_regexp}', 'g') AS extracted_id " f"FROM {ids_table['table']}" ) with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: for j in i[0]: all_ids_in_tables.append({"table": ids_table["table"], "id": j}) elif ids_table["type"] == "json": click.echo( click.style( ( f"- Listing file-id-like JSON string in column {ids_table['column']} " f"in table {ids_table['table']}" ), fg="white", ) ) query = ( f"SELECT regexp_matches({ids_table['column']}::text, '{guid_regexp}', 'g') AS extracted_id " f"FROM {ids_table['table']}" ) with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: for j in i[0]: all_ids_in_tables.append({"table": ids_table["table"], "id": j}) click.echo(click.style(f"Found {len(all_ids_in_tables)} file ids in tables.", fg="white")) except Exception as e: click.echo(click.style(f"Error fetching keys: {str(e)}", fg="red")) return # find orphaned files all_files = [file["id"] for file in all_files_in_tables] all_ids = [file["id"] for file in all_ids_in_tables] orphaned_files = list(set(all_files) - set(all_ids)) if not orphaned_files: click.echo( click.style( "No orphaned file records found. There is nothing to delete.", fg="green", ) ) return click.echo(click.style(f"Found {len(orphaned_files)} orphaned file records.", fg="white")) for file in orphaned_files: click.echo(click.style(f"- orphaned file id: {file}", fg="black")) if not force: click.confirm( f"Do you want to proceed to delete all {len(orphaned_files)} orphaned file records?", abort=True, ) # delete orphaned records for each file try: for files_table in files_tables: click.echo( click.style( f"- Deleting orphaned file records in table {files_table['table']}", fg="white", ) ) query = f"DELETE FROM {files_table['table']} WHERE {files_table['id_column']} IN :ids" with db.engine.begin() as conn: conn.execute(db.text(query), {"ids": tuple(orphaned_files)}) except Exception as e: click.echo(click.style(f"Error deleting orphaned file records: {str(e)}", fg="red")) return click.echo(click.style(f"Removed {len(orphaned_files)} orphaned file records.", fg="green")) @click.option( "-f", "--force", is_flag=True, help="Skip user confirmation and force the command to execute.", ) @click.command("remove-orphaned-files-on-storage", help="Remove orphaned files on the storage.") def remove_orphaned_files_on_storage(force: bool): """ Remove orphaned files on the storage. """ # define tables and columns to process files_tables = [ {"table": "upload_files", "key_column": "key"}, {"table": "tool_files", "key_column": "file_key"}, ] storage_paths = ["image_files", "tools", "upload_files"] # notify user and ask for confirmation click.echo( click.style( "This command will find and remove orphaned files on the storage,", fg="yellow", ) ) click.echo( click.style( "by comparing the files on the storage with the records in the following tables:", fg="yellow", ) ) for files_table in files_tables: click.echo(click.style(f"- {files_table['table']}", fg="yellow")) click.echo( click.style( "The following paths on the storage will be scanned to find orphaned files:", fg="yellow", ) ) for storage_path in storage_paths: click.echo(click.style(f"- {storage_path}", fg="yellow")) click.echo("") click.echo(click.style("!!! USE WITH CAUTION !!!", fg="red")) click.echo( click.style( "Currently, this command will work only for opendal based storage (STORAGE_TYPE=opendal).", fg="yellow", ) ) click.echo( click.style( "Since not all patterns have been fully tested, please note that this command may delete unintended files.", fg="yellow", ) ) click.echo( click.style( "This cannot be undone. Please make sure to back up your storage before proceeding.", fg="yellow", ) ) click.echo( click.style( ( "It is also recommended to run this during the maintenance window, " "as this may cause high load on your instance." ), fg="yellow", ) ) if not force: click.confirm("Do you want to proceed?", abort=True) # start the cleanup process click.echo(click.style("Starting orphaned files cleanup.", fg="white")) # fetch file id and keys from each table all_files_in_tables = [] try: for files_table in files_tables: click.echo(click.style(f"- Listing files from table {files_table['table']}", fg="white")) query = f"SELECT {files_table['key_column']} FROM {files_table['table']}" with db.engine.begin() as conn: rs = conn.execute(db.text(query)) for i in rs: all_files_in_tables.append(str(i[0])) click.echo(click.style(f"Found {len(all_files_in_tables)} files in tables.", fg="white")) except Exception as e: click.echo(click.style(f"Error fetching keys: {str(e)}", fg="red")) all_files_on_storage = [] for storage_path in storage_paths: try: click.echo(click.style(f"- Scanning files on storage path {storage_path}", fg="white")) files = storage.scan(path=storage_path, files=True, directories=False) all_files_on_storage.extend(files) except FileNotFoundError as e: click.echo( click.style( f" -> Skipping path {storage_path} as it does not exist.", fg="yellow", ) ) continue except Exception as e: click.echo( click.style( f" -> Error scanning files on storage path {storage_path}: {str(e)}", fg="red", ) ) continue click.echo(click.style(f"Found {len(all_files_on_storage)} files on storage.", fg="white")) # find orphaned files orphaned_files = list(set(all_files_on_storage) - set(all_files_in_tables)) if not orphaned_files: click.echo(click.style("No orphaned files found. There is nothing to remove.", fg="green")) return click.echo(click.style(f"Found {len(orphaned_files)} orphaned files.", fg="white")) for file in orphaned_files: click.echo(click.style(f"- orphaned file: {file}", fg="black")) if not force: click.confirm( f"Do you want to proceed to remove all {len(orphaned_files)} orphaned files?", abort=True, ) # delete orphaned files removed_files = 0 error_files = 0 for file in orphaned_files: try: storage.delete(file) removed_files += 1 click.echo(click.style(f"- Removing orphaned file: {file}", fg="white")) except Exception as e: error_files += 1 click.echo(click.style(f"- Error deleting orphaned file {file}: {str(e)}", fg="red")) continue if error_files == 0: click.echo(click.style(f"Removed {removed_files} orphaned files without errors.", fg="green")) else: click.echo( click.style( f"Removed {removed_files} orphaned files, with {error_files} errors.", fg="yellow", ) ) # Keep the original function for backward compatibility @click.command( "create-admin-with-phone", help="Create or update an admin account for an organization with a phone number.", ) @click.option("--name", prompt=True, help="Admin account name") @click.option("--phone", prompt=True, help="Admin account phone number") @click.option("--organization-id", required=True, help="Organization ID") def create_admin_with_phone(name: str, phone: str, organization_id: str): """ Create or update an admin account with a phone number for a specific organization. This command will create a new account if the phone doesn't exist, or update an existing account with the specified admin role. """ return create_admin_account(name, phone, "phone", organization_id)