From aaa32f00d2a9aee481d9bf0b810116bdfab0fae7 Mon Sep 17 00:00:00 2001 From: "alan.cl" <1165243776@qq.com> Date: Fri, 20 Mar 2026 16:50:48 +0800 Subject: [PATCH] feat: add dbgpt app start cli --- packages/dbgpt-app/src/dbgpt_app/_cli.py | 140 ++++++- packages/dbgpt-app/src/dbgpt_app/base.py | 21 +- .../initialization/workspace_provisioning.py | 41 ++ .../pilot_template/meta_data/alembic.ini | 81 ++++ .../pilot_template/meta_data/alembic/README | 1 + .../pilot_template/meta_data/alembic/env.py | 74 ++++ .../meta_data/alembic/script.py.mako | 26 ++ .../tests/test_workspace_provisioning.py | 47 +++ packages/dbgpt-core/src/dbgpt/cli/_config.py | 322 +++++++++++++++ .../dbgpt-core/src/dbgpt/cli/_profile_cmd.py | 95 +++++ .../dbgpt-core/src/dbgpt/cli/_profiles.py | 203 ++++++++++ packages/dbgpt-core/src/dbgpt/cli/_wizard.py | 325 +++++++++++++++ .../dbgpt-core/src/dbgpt/cli/cli_scripts.py | 122 +++++- .../src/dbgpt/cli/tests/__init__.py | 0 .../src/dbgpt/cli/tests/conftest.py | 24 ++ .../src/dbgpt/cli/tests/test_cli_scripts.py | 65 +++ .../src/dbgpt/cli/tests/test_config.py | 381 ++++++++++++++++++ .../src/dbgpt/cli/tests/test_integration.py | 170 ++++++++ .../src/dbgpt/cli/tests/test_profile_cmd.py | 126 ++++++ .../src/dbgpt/cli/tests/test_profiles.py | 175 ++++++++ .../src/dbgpt/cli/tests/test_smoke.py | 43 ++ .../src/dbgpt/cli/tests/test_wizard.py | 251 ++++++++++++ .../src/dbgpt/configs/model_config.py | 29 +- .../src/dbgpt/configs/tests/__init__.py | 0 .../dbgpt/configs/tests/test_model_config.py | 134 ++++++ .../src/dbgpt/util/_db_migration_utils.py | 4 +- .../src/dbgpt/util/console/console.py | 158 +++++++- .../util/tests/test_db_migration_utils.py | 30 ++ 28 files changed, 3069 insertions(+), 19 deletions(-) create mode 100644 packages/dbgpt-app/src/dbgpt_app/initialization/workspace_provisioning.py create mode 100644 packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic.ini create mode 100644 packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/README create mode 100644 packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/env.py create mode 100644 packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/script.py.mako create mode 100644 packages/dbgpt-app/src/dbgpt_app/tests/test_workspace_provisioning.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/_config.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/_profile_cmd.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/_profiles.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/_wizard.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/__init__.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/conftest.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_cli_scripts.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_config.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_integration.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_profile_cmd.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_profiles.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_smoke.py create mode 100644 packages/dbgpt-core/src/dbgpt/cli/tests/test_wizard.py create mode 100644 packages/dbgpt-core/src/dbgpt/configs/tests/__init__.py create mode 100644 packages/dbgpt-core/src/dbgpt/configs/tests/test_model_config.py create mode 100644 packages/dbgpt-core/src/dbgpt/util/tests/test_db_migration_utils.py diff --git a/packages/dbgpt-app/src/dbgpt_app/_cli.py b/packages/dbgpt-app/src/dbgpt_app/_cli.py index dd47268e4..e1738f283 100644 --- a/packages/dbgpt-app/src/dbgpt_app/_cli.py +++ b/packages/dbgpt-app/src/dbgpt_app/_cli.py @@ -5,24 +5,142 @@ from typing import Any, Dict, Optional import click from dbgpt.configs.model_config import LOGDIR -from dbgpt.model.cli import add_start_server_options from dbgpt.util.command_utils import _run_current_with_daemon, _stop_service from dbgpt.util.i18n_utils import _ _GLOBAL_CONFIG: str = "" +_BANNER_ART = """\ + ____ ____ ____ ____ _____ + | _ \\| __ ) / ___| _ \\_ _| + | | | | _ \\ ____| | _| |_) || | + | |_| | |_) |____| |_| | __/ | | + |____/|____/ \\____|_| |_|\ +""" + + +def _print_banner() -> None: + """Print the DB-GPT ASCII art banner to the terminal.""" + from dbgpt.util.console.console import CliLogger + + _log = CliLogger() + _log.print(f"[bold bright_blue]{_BANNER_ART}[/bold bright_blue]") + _log.print("") + _log.print(" [dim]πŸš€ DB-GPT Quick Start[/dim]") + _log.print("") + + +def _add_webserver_start_options(func): + """Click options decorator for the webserver start command. + + Unlike the generic ``add_start_server_options``, ``--config`` here is + *optional* so that users can rely on ``--profile`` / wizard flow instead. + """ + + @click.option( + "-c", + "--config", + type=str, + required=False, + default=None, + help=_( + "Path to a TOML config file. If omitted, DB-GPT will use the active " + "profile from ~/.dbgpt/ or run the first-time setup wizard." + ), + ) + @click.option( + "-p", + "--profile", + type=str, + required=False, + default=None, + help=_( + "Name of the provider profile to use (openai / kimi / qwen / minimax / " + "deepseek / ollama). Overrides the active profile in ~/.dbgpt/config.toml." + ), + ) + @click.option( + "-y", + "--yes", + is_flag=True, + default=False, + help=_( + "Non-interactive mode: skip the setup wizard and use defaults / " + "environment variables. Useful for CI/CD and scripted installs." + ), + ) + @click.option( + "--api-key", + type=str, + required=False, + default=None, + envvar="DBGPT_API_KEY", + help=_( + "API key for the chosen provider. Can also be set via the provider's " + "own environment variable (e.g. OPENAI_API_KEY)." + ), + ) + @click.option( + "-d", + "--daemon", + is_flag=True, + help=_( + "Run in daemon mode. It will run in the background. If you want to stop" + " it, use `dbgpt stop` command" + ), + ) + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + @click.command(name="webserver") -@add_start_server_options -def start_webserver(config: str, **kwargs): - """Start webserver(dbgpt_server.py)""" - if kwargs["daemon"]: +@_add_webserver_start_options +def start_webserver( + config: Optional[str], + profile: Optional[str], + yes: bool, + api_key: Optional[str], + **kwargs, +): + """Start webserver (dbgpt_server.py). + + On first run (or when no config is found) DB-GPT will launch an + interactive setup wizard so you can choose your LLM provider and API key. + Use ``--yes`` to skip the wizard in non-interactive environments. + """ + # Print banner first β€” skip in daemon mode (output goes to a log file) + if not kwargs.get("daemon"): + _print_banner() + + if kwargs.get("daemon"): log_file = os.path.join(LOGDIR, "webserver_uvicorn.log") _run_current_with_daemon("WebServer", log_file) - else: - from dbgpt_app.dbgpt_server import run_webserver + return - run_webserver(config) + # Resolve (or create) a config file via the wizard if needed + try: + from dbgpt.cli._wizard import maybe_run_wizard + + resolved_config = maybe_run_wizard( + profile=profile, + config=config, + yes=yes, + api_key=api_key, + ) + except ImportError: + # Graceful fallback: if wizard module is somehow unavailable, require --config + if not config: + raise click.UsageError( + "No config file found. Please pass --config or run `dbgpt setup`." + ) + resolved_config = config + + from dbgpt_app.dbgpt_server import run_webserver + + run_webserver(resolved_config) @click.command(name="webserver") @@ -312,6 +430,12 @@ def _get_migration_config( default_meta_data_path = _initialize_db( db_url, "sqlite", db_name, db_engine_args, try_to_create_db=True ) + from dbgpt_app.initialization.workspace_provisioning import _ensure_pilot_workspace + + # Provision pilot workspace template files for pip-installed users. + # dest_root is the parent of meta_data/ (e.g. ~/.dbgpt/workspace/pilot/) + pilot_root = os.path.dirname(default_meta_data_path) + _ensure_pilot_workspace(pilot_root) alembic_cfg = create_alembic_config( default_meta_data_path, db_manager.engine, diff --git a/packages/dbgpt-app/src/dbgpt_app/base.py b/packages/dbgpt-app/src/dbgpt_app/base.py index b7bc37404..db6e7cfd2 100644 --- a/packages/dbgpt-app/src/dbgpt_app/base.py +++ b/packages/dbgpt-app/src/dbgpt_app/base.py @@ -104,7 +104,26 @@ def _migration_db_storage( from dbgpt_app.initialization.db_model_initialization import _MODELS # noqa: F401 from dbgpt_ext.datasource.rdbms.conn_sqlite import SQLiteConnectorParameters - default_meta_data_path = os.path.join(PILOT_PATH, "meta_data") + # Derive meta_data path from the resolved db path when available. + # For pip-installed users, db_params.path is an absolute path like + # ~/.dbgpt/workspace/pilot/meta_data/dbgpt.db (already resolved by + # _initialize_db_storage), so we use its parent directory. For source-code + # developers with relative paths, we fall back to PILOT_PATH/meta_data. + if ( + isinstance(db_params, SQLiteConnectorParameters) + and hasattr(db_params, "path") + and db_params.path + and os.path.isabs(db_params.path) + ): + default_meta_data_path = os.path.dirname(db_params.path) + else: + default_meta_data_path = os.path.join(PILOT_PATH, "meta_data") + from dbgpt_app.initialization.workspace_provisioning import _ensure_pilot_workspace + + # Provision pilot workspace template files for pip-installed users. + # dest_root is the parent of meta_data/ (e.g. ~/.dbgpt/workspace/pilot/) + pilot_root = os.path.dirname(default_meta_data_path) + _ensure_pilot_workspace(pilot_root) if not disable_alembic_upgrade: from dbgpt.storage.metadata.db_manager import db from dbgpt.util._db_migration_utils import _ddl_init_and_upgrade diff --git a/packages/dbgpt-app/src/dbgpt_app/initialization/workspace_provisioning.py b/packages/dbgpt-app/src/dbgpt_app/initialization/workspace_provisioning.py new file mode 100644 index 000000000..eeaf68f0e --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/initialization/workspace_provisioning.py @@ -0,0 +1,41 @@ +"""Workspace provisioning module for dbgpt-app pip package users. + +Copies pilot template files to the user's workspace directory on first startup. +""" + +import logging +import os +import shutil + +logger = logging.getLogger(__name__) + + +def _ensure_pilot_workspace(dest_root: str) -> None: + """Idempotently copy pilot workspace template files to dest_root. + + This function is safe to call multiple times β€” it will never overwrite + existing files. On first run, it provisions the full pilot/ directory + structure including alembic config and benchmark data. + + Args: + dest_root (str): The destination root directory (parent of meta_data/). + Example: ~/.dbgpt/workspace/pilot/ + """ + template_dir = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "pilot_template", + ) + for src_dir, dirs, files in os.walk(template_dir): + dirs[:] = [d for d in dirs if d not in ("__pycache__", "versions")] + for filename in files: + if filename in (".DS_Store",) or filename.endswith((".pyc",)): + continue + src_file = os.path.join(src_dir, filename) + rel_path = os.path.relpath(src_file, template_dir) + dest_file = os.path.join(dest_root, rel_path) + if not os.path.exists(dest_file): + os.makedirs(os.path.dirname(dest_file), exist_ok=True) + shutil.copy2(src_file, dest_file) + logger.info("Provisioned: %s", dest_file) + else: + logger.debug("Skipped (exists): %s", dest_file) diff --git a/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic.ini b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic.ini new file mode 100644 index 000000000..7315483fe --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic.ini @@ -0,0 +1,81 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = alembic + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python-dateutil library that can be +# installed by adding `alembic[tz]` to the pip requirements +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the +# "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +#sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, execute a binary +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = --fix REVISION_SCRIPT_FILENAME diff --git a/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/README b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/README new file mode 100644 index 000000000..98e4f9c44 --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/env.py b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/env.py new file mode 100644 index 000000000..f10832a27 --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/env.py @@ -0,0 +1,74 @@ +from alembic import context + +from dbgpt.storage.metadata.db_manager import db + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel + + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + target_metadata = db.metadata + url = config.get_main_option("sqlalchemy.url") + assert target_metadata is not None + assert url is not None + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + engine = db.engine + target_metadata = db.metadata + with engine.connect() as connection: + if engine.dialect.name == "sqlite": + context.configure( + connection=engine.connect(), + target_metadata=target_metadata, + render_as_batch=True, + ) + else: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/script.py.mako b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/script.py.mako new file mode 100644 index 000000000..fbc4b07dc --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/packages/dbgpt-app/src/dbgpt_app/tests/test_workspace_provisioning.py b/packages/dbgpt-app/src/dbgpt_app/tests/test_workspace_provisioning.py new file mode 100644 index 000000000..3b724714a --- /dev/null +++ b/packages/dbgpt-app/src/dbgpt_app/tests/test_workspace_provisioning.py @@ -0,0 +1,47 @@ +"""Tests for workspace_provisioning module.""" + +import os + +from dbgpt_app.initialization.workspace_provisioning import _ensure_pilot_workspace + + +def test_ensure_pilot_workspace_copies_alembic_ini(tmp_path): + """Test that alembic.ini is copied to dest_root/meta_data/alembic.ini""" + _ensure_pilot_workspace(str(tmp_path)) + assert os.path.exists(tmp_path / "meta_data" / "alembic.ini") + + +def test_ensure_pilot_workspace_copies_alembic_env_py(tmp_path): + """Test that alembic/env.py is copied""" + _ensure_pilot_workspace(str(tmp_path)) + assert os.path.exists(tmp_path / "meta_data" / "alembic" / "env.py") + + +def test_ensure_pilot_workspace_copies_alembic_script_mako(tmp_path): + """Test that alembic/script.py.mako is copied""" + _ensure_pilot_workspace(str(tmp_path)) + assert os.path.exists(tmp_path / "meta_data" / "alembic" / "script.py.mako") + + +def test_ensure_pilot_workspace_copies_benchmark_xlsx(tmp_path): + """Test that benchmark xlsx is copied""" + _ensure_pilot_workspace(str(tmp_path)) + xlsx_files = list((tmp_path / "benchmark_meta_data").glob("*.xlsx")) + assert len(xlsx_files) == 1 + + +def test_ensure_pilot_workspace_idempotent_no_overwrite(tmp_path): + """Test that calling twice does not overwrite existing files""" + _ensure_pilot_workspace(str(tmp_path)) + ini_path = tmp_path / "meta_data" / "alembic.ini" + # write custom content to simulate user modification + ini_path.write_text("custom content") + _ensure_pilot_workspace(str(tmp_path)) # call again + assert ini_path.read_text() == "custom content" # must not be overwritten + + +def test_ensure_pilot_workspace_creates_missing_directories(tmp_path): + """Test that missing destination directories are created automatically""" + dest = tmp_path / "deep" / "nested" / "pilot" + _ensure_pilot_workspace(str(dest)) + assert os.path.exists(dest / "meta_data" / "alembic.ini") diff --git a/packages/dbgpt-core/src/dbgpt/cli/_config.py b/packages/dbgpt-core/src/dbgpt/cli/_config.py new file mode 100644 index 000000000..97da19a86 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/_config.py @@ -0,0 +1,322 @@ +"""User-level configuration management for DB-GPT CLI. + +Manages ``~/.dbgpt/configs/.toml`` β€” one flat TOML file per +profile β€” and a small ``~/.dbgpt/config.toml`` that records which profile +is active by default. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Optional + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + +_DBGPT_HOME = Path(os.environ.get("DBGPT_HOME", str(Path.home() / ".dbgpt"))) +_CONFIGS_DIR = _DBGPT_HOME / "configs" +_ACTIVE_CONFIG = _DBGPT_HOME / "config.toml" + + +def dbgpt_home() -> Path: + """Return ``~/.dbgpt``, creating it if necessary.""" + _DBGPT_HOME.mkdir(parents=True, exist_ok=True) + return _DBGPT_HOME + + +def configs_dir() -> Path: + """Return ``~/.dbgpt/configs/``, creating it if necessary.""" + _CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + return _CONFIGS_DIR + + +def profile_config_path(profile_name: str) -> Path: + """Return the path for a profile TOML, e.g. ``~/.dbgpt/configs/openai.toml``.""" + return configs_dir() / f"{profile_name}.toml" + + +def active_config_path() -> Path: + """Return ``~/.dbgpt/config.toml``.""" + return _ACTIVE_CONFIG + + +# --------------------------------------------------------------------------- +# Active-profile record +# --------------------------------------------------------------------------- + + +def read_active_profile() -> Optional[str]: + """Return the name of the active profile from ``~/.dbgpt/config.toml``. + + Returns: + Optional[str]: Profile name, or *None* if not yet configured. + """ + path = active_config_path() + if not path.exists(): + return None + try: + import tomlkit # type: ignore[import] + + data = tomlkit.loads(path.read_text(encoding="utf-8")) + return data.get("default", {}).get("profile") or None + except Exception: + return None + + +def write_active_profile(profile_name: str) -> None: + """Persist the active profile name to ``~/.dbgpt/config.toml``. + + Args: + profile_name (str): The profile to activate. + """ + import tomlkit # type: ignore[import] + + path = active_config_path() + dbgpt_home() # ensure directory exists + + if path.exists(): + try: + data = tomlkit.loads(path.read_text(encoding="utf-8")) + except Exception: + data = tomlkit.document() + else: + data = tomlkit.document() + + if "default" not in data: + data["default"] = tomlkit.table() + data["default"]["profile"] = profile_name # type: ignore[index] + path.write_text(tomlkit.dumps(data), encoding="utf-8") + + +# --------------------------------------------------------------------------- +# Profile TOML generation +# --------------------------------------------------------------------------- + + +def _escape_toml_string(value: str) -> str: + """Escape backslashes and double quotes for TOML basic strings.""" + return value.replace("\\", "\\\\").replace('"', '\\"') + + +def _render_profile_toml( + spec: "ProfileSpec", # noqa: F821 + api_key: Optional[str], + llm_model: Optional[str] = None, + embedding_model: Optional[str] = None, + api_base: Optional[str] = None, + embedding_api_key: Optional[str] = None, +) -> str: + """Render a complete TOML config for the given profile. + + The generated file uses a literal API key when *api_key* is provided; + otherwise it uses the ``${env:VAR}`` interpolation syntax so the server + reads the key from the environment at runtime. + + Args: + spec: A :class:`~dbgpt.cli._profiles.ProfileSpec` instance. + api_key (Optional[str]): Literal API key value, or *None*. + llm_model (Optional[str]): Override for the LLM model name. If *None*, + uses ``spec.llm_model``. + embedding_model (Optional[str]): Override for the embedding model name. + If *None*, uses ``spec.embedding_model``. + api_base (Optional[str]): Override for the LLM API base URL. If *None*, + uses ``spec.llm_api_base``. + embedding_api_key (Optional[str]): Literal embedding API key, or + *None*. When *None* and ``spec.embedding_env_var`` is set (and + differs from ``spec.env_var``), the generated TOML will reference + that separate env var instead of the LLM key. + + Returns: + str: TOML content as a string. + """ + from dbgpt.cli._profiles import ProfileSpec # noqa: F401 (type-only) + + if api_key: + api_key_value = _escape_toml_string(api_key) + elif spec.env_var: + api_key_value = f"${{env:{spec.env_var}:-sk-xxx}}" + else: + api_key_value = "" + + emb_env_var = spec.embedding_env_var or spec.env_var + if embedding_api_key: + emb_key_value = _escape_toml_string(embedding_api_key) + elif api_key and not spec.embedding_env_var: + emb_key_value = _escape_toml_string(api_key) + elif emb_env_var: + emb_key_value = f"${{env:{emb_env_var}:-sk-xxx}}" + else: + emb_key_value = "" + + if api_base: + embedding_api_url = f"{api_base.rstrip('/')}/embeddings" + else: + embedding_api_url = spec.embedding_api_url + + data_dir = str(dbgpt_home() / "workspace" / "pilot" / "meta_data" / "dbgpt.db") + vector_dir = str(dbgpt_home() / "workspace" / "pilot" / "data") + + lines = [ + f"# DB-GPT configuration β€” profile: {spec.name}", + "# Generated by `dbgpt setup`", + "", + "[system]", + "# Load language from environment variable(It is set by the hook)", + 'language = "${env:DBGPT_LANG:-en}"', + "api_keys = []", + 'encrypt_key = "your_secret_key"', + "", + "# Server Configurations", + "[service.web]", + 'host = "0.0.0.0"', + "port = 5670", + "", + "[service.web.database]", + 'type = "sqlite"', + f'path = "{data_dir}"', + "", + "[rag.storage]", + "[rag.storage.vector]", + 'type = "chroma"', + f'persist_path = "{vector_dir}"', + "", + "# Model Configurations", + "[models]", + "[[models.llms]]", + ] + + use_env = getattr(spec, "use_env_interpolation", False) and not api_key + + if use_env: + lines.append(f'name = "${{env:LLM_MODEL_NAME:-{spec.llm_model}}}"') + lines.append(f'provider = "${{env:LLM_MODEL_PROVIDER:-{spec.llm_provider}}}"') + effective_api_base = api_base or spec.llm_api_base + if effective_api_base: + lines.append(f'api_base = "${{env:OPENAI_API_BASE:-{effective_api_base}}}"') + lines.append(f'api_key = "{api_key_value}"') + lines.append("") + lines.append("[[models.embeddings]]") + lines.append(f'name = "${{env:EMBEDDING_MODEL_NAME:-{spec.embedding_model}}}"') + lines.append( + f'provider = "${{env:EMBEDDING_MODEL_PROVIDER:-{spec.embedding_provider}}}"' + ) + lines.append( + f'api_url = "${{env:EMBEDDING_MODEL_API_URL:-{embedding_api_url}}}"' + ) + lines.append(f'api_key = "{emb_key_value}"') + else: + lines.append(f'name = "{llm_model or spec.llm_model}"') + lines.append(f'provider = "{spec.llm_provider}"') + + effective_api_base = api_base or spec.llm_api_base + if effective_api_base: + lines.append(f'api_base = "{effective_api_base}"') + + lines.append(f'api_key = "{api_key_value}"') + lines.append("") + lines.append("[[models.embeddings]]") + lines.append(f'name = "{embedding_model or spec.embedding_model}"') + lines.append(f'provider = "{spec.embedding_provider}"') + lines.append(f'api_url = "{embedding_api_url}"') + lines.append(f'api_key = "{emb_key_value}"') + + # Append any provider-specific extras + for extra in spec.extra_toml_lines: + lines.append(extra) + + return "\n".join(lines) + "\n" + + +# --------------------------------------------------------------------------- +# Public write API +# --------------------------------------------------------------------------- + + +def write_profile_config( + profile_name: str, + api_key: Optional[str] = None, + activate: bool = True, + llm_model: Optional[str] = None, + embedding_model: Optional[str] = None, + api_base: Optional[str] = None, + embedding_api_key: Optional[str] = None, +) -> Path: + """Write (or overwrite) the TOML config for *profile_name*. + + Args: + profile_name (str): One of the supported profile names. + api_key (Optional[str]): Literal API key. If *None*, env-var + interpolation is used instead. + activate (bool): Also update ``~/.dbgpt/config.toml`` to make this + profile the active default. + llm_model (Optional[str]): Override LLM model name. Uses spec default + if *None*. + embedding_model (Optional[str]): Override embedding model name. Uses + spec default if *None*. + api_base (Optional[str]): Override LLM API base URL. Uses spec default + if *None*. + embedding_api_key (Optional[str]): Literal embedding API key. When + *None* and the profile has a separate ``embedding_env_var``, the + generated TOML will reference that env var. + + Returns: + Path: Path to the written config file. + """ + from dbgpt.cli._profiles import get_profile + + spec = get_profile(profile_name) + content = _render_profile_toml( + spec, + api_key, + llm_model=llm_model, + embedding_model=embedding_model, + api_base=api_base, + embedding_api_key=embedding_api_key, + ) + path = profile_config_path(profile_name) + path.write_text(content, encoding="utf-8") + + if activate: + write_active_profile(profile_name) + + return path + + +def resolve_config_path( + profile: Optional[str] = None, + config: Optional[str] = None, +) -> Optional[str]: + """Resolve which config file to use, in priority order. + + Priority: + 1. Explicit ``--config`` flag β†’ use as-is. + 2. Explicit ``--profile`` flag β†’ look up ``~/.dbgpt/configs/.toml``. + 3. Active profile from ``~/.dbgpt/config.toml``. + 4. Return *None* (caller should run the setup wizard). + + Args: + profile (Optional[str]): Value of ``--profile`` CLI flag. + config (Optional[str]): Value of ``--config`` CLI flag. + + Returns: + Optional[str]: Absolute path to the config file, or *None*. + """ + if config: + return config + + if profile: + path = profile_config_path(profile) + if path.exists(): + return str(path) + return None # profile specified but not yet configured + + # Fall back to whatever is active + active = read_active_profile() + if active: + path = profile_config_path(active) + if path.exists(): + return str(path) + + return None diff --git a/packages/dbgpt-core/src/dbgpt/cli/_profile_cmd.py b/packages/dbgpt-core/src/dbgpt/cli/_profile_cmd.py new file mode 100644 index 000000000..aa4014183 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/_profile_cmd.py @@ -0,0 +1,95 @@ +"""Profile management subcommands: list, show, create, switch, delete.""" + +import click + + +@click.group() +def profile(): + """Manage DB-GPT configuration profiles.""" + pass + + +@profile.command(name="list") +def profile_list(): + """List all configured profiles.""" + from dbgpt.cli._config import configs_dir, read_active_profile + + configs = configs_dir() + active = read_active_profile() + toml_files = sorted(configs.glob("*.toml")) + + if not toml_files: + click.echo("No profiles configured. Run: dbgpt setup") + return + + for f in toml_files: + name = f.stem + marker = "* " if name == active else " " + click.echo(f"{marker}{name}") + + +@profile.command() +@click.argument("name") +def show(name): + """Show the TOML configuration for a profile.""" + from dbgpt.cli._config import profile_config_path + + path = profile_config_path(name) + if not path.exists(): + raise click.ClickException( + f"Profile '{name}' not found. Run: dbgpt profile create {name}" + ) + click.echo(path.read_text(encoding="utf-8")) + + +@profile.command() +@click.argument("name") +def create(name): + """Create or reconfigure a profile (runs the setup wizard).""" + from dbgpt.cli._wizard import run_setup_wizard + + run_setup_wizard(pre_selected_profile=name) + + +@profile.command() +@click.argument("name") +def switch(name): + """Set a profile as the active default.""" + from dbgpt.cli._config import profile_config_path, write_active_profile + + path = profile_config_path(name) + if not path.exists(): + click.echo( + f"Error: Profile '{name}' not found. Run: dbgpt profile create {name}", + err=False, + ) + raise SystemExit(1) + write_active_profile(name) + click.echo(f"Switched active profile to: {name}") + + +@profile.command() +@click.argument("name") +@click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompt.") +def delete(name, yes): + """Delete a profile configuration file.""" + from dbgpt.cli._config import ( + profile_config_path, + read_active_profile, + write_active_profile, + ) + + path = profile_config_path(name) + if not path.exists(): + raise click.ClickException(f"Profile '{name}' not found.") + + if not yes: + click.confirm(f"Delete profile '{name}'?", abort=True) + + path.unlink() + + # Clear active pointer if this was the active profile + if read_active_profile() == name: + write_active_profile("") # clear by writing empty string + + click.echo(f"Deleted profile: {name}") diff --git a/packages/dbgpt-core/src/dbgpt/cli/_profiles.py b/packages/dbgpt-core/src/dbgpt/cli/_profiles.py new file mode 100644 index 000000000..3646d802b --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/_profiles.py @@ -0,0 +1,203 @@ +"""Profile definitions and API key resolution for DB-GPT CLI. + +Each profile corresponds to a supported LLM provider and contains the +information needed to generate a TOML configuration file and resolve +API credentials from environment variables. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Dict, List, Optional + + +@dataclass +class ProfileSpec: + """Specification for a single LLM provider profile.""" + + name: str + """Internal identifier, e.g. 'openai'.""" + + label: str + """Human-readable display name, e.g. 'OpenAI (GPT-4o)'.""" + + env_var: str + """Primary environment variable that holds the API key.""" + + llm_model: str + """Default LLM model name.""" + + llm_provider: str + """DB-GPT provider string, e.g. 'proxy/openai'.""" + + llm_api_base: str + """Base URL for the LLM API.""" + + embedding_model: str + """Default embedding model name.""" + + embedding_provider: str + """DB-GPT provider string for embeddings.""" + + embedding_api_url: str + """URL for the embedding API endpoint.""" + + needs_api_key: bool = True + """Whether this profile requires an API key.""" + + use_env_interpolation: bool = False + """When True, model fields use ``${env:VAR:-default}`` syntax in generated TOML. + + Set to True only for the *default* profile so that ``default.toml`` mirrors + ``configs/dbgpt-proxy-openai.toml`` exactly β€” allowing runtime override via + environment variables without re-running the setup wizard. + """ + + extra_toml_lines: List[str] = field(default_factory=list) + """Extra TOML lines appended verbatim to the generated config.""" + + embedding_env_var: Optional[str] = None + """Environment variable for the embedding API key. + + When *None* (the default), the same ``env_var`` used for the LLM key is + reused for embeddings. Set this to a different name when the embedding + endpoint requires a separate API key (e.g. Kimi uses MOONSHOT_API_KEY for + LLM but DASHSCOPE_API_KEY for embeddings via DashScope/Tongyi). + """ + + def env_key(self) -> Optional[str]: + """Return the API key from the environment, or None.""" + return os.environ.get(self.env_var) + + def embedding_env_key(self) -> Optional[str]: + """Return the embedding API key from the environment, or None. + + Falls back to the primary ``env_var`` when ``embedding_env_var`` is + not set. + """ + var = self.embedding_env_var or self.env_var + return os.environ.get(var) if var else None + + +# --------------------------------------------------------------------------- +# Supported profiles +# --------------------------------------------------------------------------- + +PROFILES: Dict[str, ProfileSpec] = { + "openai": ProfileSpec( + name="openai", + label="OpenAI (OpenAI or OpenAI API proxy)", + env_var="OPENAI_API_KEY", + llm_model="gpt-4o", + llm_provider="proxy/openai", + llm_api_base="https://api.openai.com/v1", + embedding_model="text-embedding-3-small", + embedding_provider="proxy/openai", + embedding_api_url="https://api.openai.com/v1/embeddings", + ), + "kimi": ProfileSpec( + name="kimi", + label="Kimi (Moonshot AI / kimi-k2)", + env_var="MOONSHOT_API_KEY", + llm_model="kimi-k2", + llm_provider="proxy/moonshot", + llm_api_base="https://api.moonshot.cn/v1", + embedding_model="text-embedding-v3", + embedding_provider="proxy/tongyi", + embedding_api_url="https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings", + embedding_env_var="DASHSCOPE_API_KEY", + ), + "qwen": ProfileSpec( + name="qwen", + label="Qwen (DashScope API)", + env_var="DASHSCOPE_API_KEY", + llm_model="qwen-plus", + llm_provider="proxy/tongyi", + llm_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1", + embedding_model="text-embedding-v3", + embedding_provider="proxy/tongyi", + embedding_api_url="https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings", + ), + "minimax": ProfileSpec( + name="minimax", + label="MiniMax (abab series)", + env_var="MINIMAX_API_KEY", + llm_model="abab6.5s-chat", + llm_provider="proxy/openai", + llm_api_base="https://api.minimax.chat/v1", + embedding_model="embo-01", + embedding_provider="proxy/openai", + embedding_api_url="https://api.minimax.chat/v1/embeddings", + ), + "glm": ProfileSpec( + name="glm", + label="z.ai (zhipu.ai API)", + env_var="ZHIPUAI_API_KEY", + llm_model="glm-4-plus", + llm_provider="proxy/zhipu", + llm_api_base="https://open.bigmodel.cn/api/paas/v4", + embedding_model="embedding-3", + embedding_provider="proxy/zhipu", + embedding_api_url="https://open.bigmodel.cn/api/paas/v4/embeddings", + ), + "custom": ProfileSpec( + name="custom", + label="Custom Provider (Any OpenAI compatible endpoint)", + env_var="OPENAI_API_KEY", + llm_model="gpt-4o", + llm_provider="proxy/openai", + llm_api_base="https://api.openai.com/v1", + embedding_model="text-embedding-3-small", + embedding_provider="proxy/openai", + embedding_api_url="https://api.openai.com/v1/embeddings", + ), + "default": ProfileSpec( + name="default", + label="Skip for now (use OpenAI defaults)", + env_var="OPENAI_API_KEY", + llm_model="gpt-4o", + llm_provider="proxy/openai", + llm_api_base="https://api.openai.com/v1", + embedding_model="text-embedding-3-small", + embedding_provider="proxy/openai", + embedding_api_url="https://api.openai.com/v1/embeddings", + needs_api_key=False, + use_env_interpolation=True, + ), +} + +# Ordered list for display in the wizard +PROFILE_ORDER: List[str] = [ + "openai", + "kimi", + "qwen", + "minimax", + "glm", + "custom", + "default", +] + + +def get_profile(name: str) -> ProfileSpec: + """Return a ProfileSpec by name. + + Args: + name (str): Profile identifier (case-insensitive). + + Returns: + ProfileSpec: The matching profile spec. + + Raises: + ValueError: If the profile name is not recognised. + """ + key = name.lower() + if key not in PROFILES: + valid = ", ".join(PROFILE_ORDER) + raise ValueError(f"Unknown profile '{name}'. Valid profiles: {valid}") + return PROFILES[key] + + +def list_profiles() -> List[ProfileSpec]: + """Return profiles in canonical display order.""" + return [PROFILES[k] for k in PROFILE_ORDER] diff --git a/packages/dbgpt-core/src/dbgpt/cli/_wizard.py b/packages/dbgpt-core/src/dbgpt/cli/_wizard.py new file mode 100644 index 000000000..677e66754 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/_wizard.py @@ -0,0 +1,325 @@ +"""First-run setup wizard for DB-GPT CLI. + +Provides :func:`run_setup_wizard` (interactive) and +:func:`run_setup_noninteractive` (``--yes`` / CI mode). + +Interactive flow:: + + Welcome to DB-GPT! πŸŽ‰ + + Which LLM provider would you like to use? + + ● OpenAI OpenAI or OpenAI API proxy + β—‹ Kimi Moonshot AI + β—‹ Qwen DashScope API + β—‹ MiniMax abab series + β—‹ Z.AI zhipu.ai API + β—‹ Custom Any OpenAI compatible endpoint + β—‹ Skip for now Use OpenAI defaults + + Enter your OPENAI_API_KEY: **** + βœ” Config saved β†’ ~/.dbgpt/configs/openai.toml + +For profiles with a separate embedding key (e.g. Kimi uses MOONSHOT_API_KEY +for LLM but DASHSCOPE_API_KEY for embeddings), the wizard will prompt for +both keys. +""" + +from __future__ import annotations + +import os +from typing import Optional + +from dbgpt.cli._config import ( + resolve_config_path, + write_profile_config, +) +from dbgpt.cli._profiles import ProfileSpec, get_profile, list_profiles +from dbgpt.util.console.console import CliLogger + +_log = CliLogger() + +# --------------------------------------------------------------------------- +# Provider display metadata (description only) +# --------------------------------------------------------------------------- + +_PROVIDER_META = { + "openai": "OpenAI or OpenAI API proxy", + "kimi": "Moonshot AI", + "qwen": "DashScope API", + "minimax": "abab series", + "glm": "zhipu.ai API", + "custom": "Any OpenAI compatible endpoint", + "default": "Use OpenAI defaults", +} + +_DISPLAY_NAMES = { + "openai": "OpenAI", + "kimi": "Kimi", + "qwen": "Qwen", + "minimax": "MiniMax", + "glm": "Z.AI", + "custom": "Custom", + "default": "Skip for now", +} + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def run_setup_wizard( + pre_selected_profile: Optional[str] = None, + pre_set_key: Optional[str] = None, +) -> str: + """Run the interactive first-time setup wizard. + + Args: + pre_selected_profile (Optional[str]): If supplied (e.g. via + ``--profile``), skip the provider selection step. + pre_set_key (Optional[str]): If supplied (e.g. via ``--api-key``), + skip the API-key prompt. + + Returns: + str: Absolute path to the written config file. + """ + _print_welcome() + + # ── step 1: choose profile ────────────────────────────────────────────── + if pre_selected_profile: + try: + spec = get_profile(pre_selected_profile) + except ValueError as exc: + _log.error(str(exc), exit_code=1) + return "" + else: + spec = _ask_profile() + + # ── step 2: API key ───────────────────────────────────────────────────── + api_key: Optional[str] = None + embedding_api_key: Optional[str] = None + if spec.needs_api_key: + api_key = _ask_api_key(spec, pre_set_key) + if spec.embedding_env_var and spec.embedding_env_var != spec.env_var: + embedding_api_key = _ask_embedding_api_key(spec) + + # ── step 2.5: api_base (openai + custom ask; others use spec default) ─── + api_base: Optional[str] = None + if spec.name in ("openai", "custom"): + api_base = _ask_api_base(spec) + + # ── step 3: model names (default profile skips this) ──────────────────── + llm_model: Optional[str] = None + embedding_model: Optional[str] = None + if spec.name != "default": + llm_model, embedding_model = _ask_model_names(spec) + + # ── step 4: write config ───────────────────────────────────────────────── + config_path = write_profile_config( + spec.name, + api_key=api_key, + activate=True, + llm_model=llm_model, + embedding_model=embedding_model, + api_base=api_base, + embedding_api_key=embedding_api_key, + ) + _log.success(f"βœ” Config saved β†’ {config_path}") + return str(config_path) + + +def run_setup_noninteractive( + profile_name: str = "openai", + api_key: Optional[str] = None, +) -> str: + """Create a config without prompting (``--yes`` / CI mode). + + Uses the literal *api_key* if provided; otherwise falls back to the + environment variable defined in the profile spec. + + Args: + profile_name (str): Profile to configure. + api_key (Optional[str]): Explicit API key; *None* means use env var. + + Returns: + str: Absolute path to the written config file. + """ + spec = get_profile(profile_name) + + # If no explicit key, try to read from environment + resolved_key = api_key or (spec.env_key() if spec.needs_api_key else None) + + config_path = write_profile_config(spec.name, api_key=resolved_key, activate=True) + _log.success(f"βœ” Config saved β†’ {config_path}") + return str(config_path) + + +def maybe_run_wizard( + profile: Optional[str], + config: Optional[str], + yes: bool, + api_key: Optional[str], +) -> str: + """Decide whether to run the wizard and return a ready config path. + + This is the single call-site used by ``start webserver`` to obtain a + config path, handling all first-run and re-configuration scenarios. + + Priority: + 1. ``--config`` path supplied β†’ use it directly (skip wizard). + 2. Config already exists (via ``--profile`` or active default) β†’ reuse. + 3. ``--yes`` β†’ non-interactive setup. + 4. Interactive wizard. + + Args: + profile (Optional[str]): ``--profile`` CLI flag value. + config (Optional[str]): ``--config`` CLI flag value. + yes (bool): ``--yes`` / ``-y`` flag. + api_key (Optional[str]): ``--api-key`` flag value. + + Returns: + str: Absolute path to a usable config file. + """ + # 1. Explicit --config + if config: + return config + + # 2. Existing profile config + existing = resolve_config_path(profile=profile, config=None) + if existing: + return existing + + # 3. Non-interactive + if yes: + return run_setup_noninteractive( + profile_name=profile or "openai", + api_key=api_key, + ) + + # 4. Interactive wizard + return run_setup_wizard( + pre_selected_profile=profile, + pre_set_key=api_key, + ) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _print_welcome() -> None: + _log.print("") + _log.print("[bold bright_blue]Welcome to DB-GPT! πŸŽ‰[/bold bright_blue]") + _log.print("") + _log.info( + "Let's set up your configuration. This only takes a moment.\n" + "You can re-run [bold]dbgpt setup[/bold] at any time to change settings." + ) + _log.print("") + + +def _build_provider_option(spec: ProfileSpec) -> tuple[str, str]: + """Return (display_name, description) for a provider.""" + display_name = _DISPLAY_NAMES.get(spec.name, spec.name.capitalize()) + desc = _PROVIDER_META.get(spec.name, spec.label) + return (display_name, desc) + + +def _ask_profile() -> ProfileSpec: + """Prompt the user to choose a provider with an arrow-key selector.""" + profiles = list_profiles() + + options = [_build_provider_option(spec) for spec in profiles] + + idx = _log.select("Which LLM provider would you like to use?", options) + return profiles[idx] + + +def _ask_api_key(spec: ProfileSpec, pre_set_key: Optional[str]) -> Optional[str]: + """Ask for an API key, honouring a pre-set value or env var. + + Returns *None* for providers that don't need a key (Ollama). + Returns an empty string if the user explicitly skips. + """ + if not spec.needs_api_key: + return None + + # If an explicit key was passed in (e.g. via --api-key flag), use it. + if pre_set_key: + _log.success(f"βœ” Using supplied API key for {spec.label}") + return pre_set_key + + # Check environment + env_val = spec.env_key() + if env_val: + _log.success(f"βœ” Found [bold]{spec.env_var}[/bold] in environment β€” using it.") + return env_val + + # Prompt the user + _log.print( + f"\nEnter your [bold]{spec.env_var}[/bold] " + f"(or press Enter to use env-var at runtime):" + ) + + import getpass + + try: + entered = getpass.getpass(prompt=" API key: ").strip() + except (EOFError, KeyboardInterrupt): + _log.warning("\nSkipping API key β€” you can set it via the environment later.") + return None + + if not entered: + _log.warning( + f"No key entered. The config will reference ${spec.env_var} at runtime." + ) + return None + + return entered + + +def _ask_model_names(spec: ProfileSpec) -> tuple[str, str]: + llm = _log.ask("LLM model name", default=spec.llm_model) + emb = _log.ask("Embedding model name", default=spec.embedding_model) + return (llm, emb) + + +def _ask_api_base(spec: ProfileSpec) -> str: + return _log.ask("API base URL", default=spec.llm_api_base) + + +def _ask_embedding_api_key(spec: ProfileSpec) -> Optional[str]: + emb_env_var = spec.embedding_env_var or spec.env_var + + env_val = os.environ.get(emb_env_var) if emb_env_var else None + if env_val: + _log.success( + f"βœ” Found [bold]{emb_env_var}[/bold] in environment" + " β€” using it for embeddings." + ) + return env_val + + _log.print( + f"\nEnter your [bold]{emb_env_var}[/bold] for embeddings " + f"(or press Enter to use env-var at runtime):" + ) + + import getpass + + try: + entered = getpass.getpass(prompt=" Embedding API key: ").strip() + except (EOFError, KeyboardInterrupt): + _log.warning( + "\nSkipping embedding API key β€” you can set it via the environment later." + ) + return None + + if not entered: + _log.warning( + f"No key entered. The config will reference ${emb_env_var} at runtime." + ) + return None + + return entered diff --git a/packages/dbgpt-core/src/dbgpt/cli/cli_scripts.py b/packages/dbgpt-core/src/dbgpt/cli/cli_scripts.py index d66e9f551..750945135 100644 --- a/packages/dbgpt-core/src/dbgpt/cli/cli_scripts.py +++ b/packages/dbgpt-core/src/dbgpt/cli/cli_scripts.py @@ -34,10 +34,17 @@ def add_command_alias(command, name: str, hidden: bool = False, parent_group=Non parent_group.add_command(new_command, name=name) -@click.group() -def start(): +@click.group(invoke_without_command=True) +@click.pass_context +def start(ctx): """Start specific server.""" - pass + if ctx.invoked_subcommand is None: + # Try web first, then webserver as fallback + cmd = start.commands.get("web") or start.commands.get("webserver") + if cmd: + ctx.invoke(cmd) + else: + click.echo(ctx.get_help()) @click.group() @@ -93,6 +100,99 @@ def tool(): """DB-GPT Tools.""" +# --------------------------------------------------------------------------- +# dbgpt setup +# --------------------------------------------------------------------------- + + +@click.command(name="setup") +@click.option( + "-p", + "--profile", + type=str, + required=False, + default=None, + help=( + "Provider profile to configure: openai / kimi / qwen / minimax / " + "deepseek / ollama. If omitted, an interactive menu is shown." + ), +) +@click.option( + "-y", + "--yes", + is_flag=True, + default=False, + help="Non-interactive: skip wizard and use defaults / env variables.", +) +@click.option( + "--api-key", + type=str, + required=False, + default=None, + envvar="DBGPT_API_KEY", + help="API key for the chosen provider.", +) +@click.option( + "--show", + is_flag=True, + default=False, + help="Show the current active profile and config path, then exit.", +) +def setup_command(profile: str, yes: bool, api_key: str, show: bool): + """Configure DB-GPT's LLM provider and write ~/.dbgpt/configs/.toml. + + Run without arguments for an interactive wizard, or use --yes for + non-interactive / CI usage. + + \b + Examples: + dbgpt setup # interactive wizard + dbgpt setup --profile openai --yes # use OPENAI_API_KEY from env + dbgpt setup --profile kimi --api-key sk-xxx + dbgpt setup --show # print current config + """ + try: + from dbgpt.cli._config import ( + profile_config_path, + read_active_profile, + ) + from dbgpt.cli._wizard import run_setup_noninteractive, run_setup_wizard + from dbgpt.util.console.console import CliLogger + + cl = CliLogger() + + if show: + active = read_active_profile() + if active: + path = profile_config_path(active) + cl.info(f"Active profile : [bold]{active}[/bold]") + cl.info(f"Config file : {path}") + if not path.exists(): + cl.warning(" ⚠ Config file does not exist yet. Run `dbgpt setup`.") + else: + cl.warning("No profile configured yet. Run `dbgpt setup`.") + return + + if yes: + run_setup_noninteractive( + profile_name=profile or "openai", + api_key=api_key, + ) + else: + run_setup_wizard( + pre_selected_profile=profile, + pre_set_key=api_key, + ) + + except ImportError as e: + logger.warning(f"Setup wizard unavailable: {e}") + raise click.ClickException(str(e)) + + +# --------------------------------------------------------------------------- +# Stop all +# --------------------------------------------------------------------------- + stop_all_func_list = [] @@ -103,6 +203,17 @@ def stop_all(): stop_func() +@click.command(name="none") +def start_none(): + """Start DB-GPT in API-only mode (no web UI). [Planned]""" + click.echo( + "API-only mode (no web UI) is planned for a future release.\n" + "For now, use: dbgpt start web" + ) + + +start.add_command(start_none) + cli.add_command(start) cli.add_command(stop) # cli.add_command(install) @@ -113,6 +224,10 @@ cli.add_command(repo) cli.add_command(run) cli.add_command(net) cli.add_command(tool) +cli.add_command(setup_command) +from dbgpt.cli._profile_cmd import profile as profile_cmd # noqa: E402 + +cli.add_command(profile_cmd, name="profile") add_command_alias(stop_all, name="all", parent_group=stop) try: @@ -149,6 +264,7 @@ try: ) add_command_alias(start_webserver, name="webserver", parent_group=start) + add_command_alias(start_webserver, name="web", parent_group=start) add_command_alias(stop_webserver, name="webserver", parent_group=stop) # Add migration command add_command_alias(migration, name="migration", parent_group=db) diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/__init__.py b/packages/dbgpt-core/src/dbgpt/cli/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/conftest.py b/packages/dbgpt-core/src/dbgpt/cli/tests/conftest.py new file mode 100644 index 000000000..56164fc62 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/conftest.py @@ -0,0 +1,24 @@ +"""Shared fixtures for CLI tests.""" + +import click.testing +import pytest + + +@pytest.fixture() +def isolated_dbgpt_home(tmp_path, monkeypatch): + """Isolate DBGPT_HOME to a temp directory.""" + home = tmp_path / "dbgpt" + home.mkdir() + monkeypatch.setenv("DBGPT_HOME", str(home)) + import dbgpt.cli._config as _cfg + + monkeypatch.setattr(_cfg, "_DBGPT_HOME", home) + monkeypatch.setattr(_cfg, "_CONFIGS_DIR", home / "configs") + monkeypatch.setattr(_cfg, "_ACTIVE_CONFIG", home / "config.toml") + return home + + +@pytest.fixture() +def cli_runner(): + """Return a Click test runner.""" + return click.testing.CliRunner(mix_stderr=False) diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_cli_scripts.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_cli_scripts.py new file mode 100644 index 000000000..9ebcf832b --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_cli_scripts.py @@ -0,0 +1,65 @@ +"""Tests for cli_scripts.py β€” start group alias and bare-start behavior.""" + +from unittest.mock import patch + +from click.testing import CliRunner + +from dbgpt.cli.cli_scripts import cli + + +def test_start_web_alias_help_shows_options(): + """dbgpt start web --help exits 0 and shows --config and --profile.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["start", "web", "--help"]) + assert result.exit_code == 0, result.output + assert "--config" in result.output or "--profile" in result.output + + +def test_start_webserver_still_works(): + """dbgpt start webserver --help exits 0 (regression).""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["start", "webserver", "--help"]) + assert result.exit_code == 0, result.output + + +def test_start_controller_unaffected(): + """dbgpt start controller --help still works after changes.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["start", "controller", "--help"]) + assert result.exit_code == 0, result.output + + +def test_bare_start_invokes_web_or_shows_help(): + """bare dbgpt start (no subcommand) either invokes web or shows help β€” does NOT crash.""" # noqa: E501 + runner = CliRunner(mix_stderr=False) + with ( + patch("dbgpt.cli._wizard.maybe_run_wizard", return_value="/fake/config.toml"), + patch("dbgpt_app.dbgpt_server.run_webserver"), + ): + result = runner.invoke(cli, ["start"]) + # Must not return an error exit code from a crash + assert result.exit_code in (0, 1, 2), ( + f"Unexpected exit code: {result.exit_code}\n{result.output}" + ) + assert result.exception is None or "No such command" not in str(result.exception) + + +def test_start_none_exits_zero(cli_runner): + from dbgpt.cli.cli_scripts import cli + + result = cli_runner.invoke(cli, ["start", "none"]) + assert result.exit_code == 0 + + +def test_start_none_output_mentions_planned(cli_runner): + from dbgpt.cli.cli_scripts import cli + + result = cli_runner.invoke(cli, ["start", "none"]) + assert "planned" in result.output.lower() + + +def test_start_none_output_mentions_start_web(cli_runner): + from dbgpt.cli.cli_scripts import cli + + result = cli_runner.invoke(cli, ["start", "none"]) + assert "dbgpt start web" in result.output diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_config.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_config.py new file mode 100644 index 000000000..ea191bc6c --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_config.py @@ -0,0 +1,381 @@ +"""Tests for _config.py β€” TOML generation, path handling, etc.""" + +from dbgpt.cli._config import _render_profile_toml +from dbgpt.cli._profiles import get_profile + + +def test_escape_api_key_with_double_quote_generates_valid_toml(isolated_dbgpt_home): + """API key containing double quote should produce valid parseable TOML.""" + import tomlkit + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key='sk-abc"def') + data = tomlkit.loads(content) # should NOT raise + llms = data["models"]["llms"] + assert llms[0]["api_key"] == 'sk-abc"def' + + +def test_escape_api_key_with_backslash_generates_valid_toml(isolated_dbgpt_home): + """API key containing backslash should produce valid parseable TOML.""" + import tomlkit + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="sk-abc\\def") + data = tomlkit.loads(content) # should NOT raise + llms = data["models"]["llms"] + assert llms[0]["api_key"] == "sk-abc\\def" + + +def test_escape_env_var_placeholder_not_escaped(isolated_dbgpt_home): + """Env-var placeholder ${env:VAR:-default} should NOT be double-escaped.""" + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key=None) + assert "${env:OPENAI_API_KEY:-sk-xxx}" in content + assert "\\\\" not in content # no double backslash introduced + + +def test_escape_normal_api_key_unchanged(isolated_dbgpt_home): + """Normal API key (no special chars) should work exactly as before.""" + import tomlkit + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="sk-normal-key-123") + data = tomlkit.loads(content) + llms = data["models"]["llms"] + assert llms[0]["api_key"] == "sk-normal-key-123" + + +class TestDbgptHomeEnvVar: + def test_dbgpt_home_returns_custom_path(self, isolated_dbgpt_home): + from dbgpt.cli._config import dbgpt_home + + result = dbgpt_home() + assert result == isolated_dbgpt_home + + def test_configs_dir_under_custom_home(self, isolated_dbgpt_home): + from dbgpt.cli._config import configs_dir + + result = configs_dir() + assert result == isolated_dbgpt_home / "configs" + + def test_profile_config_path_under_custom_home(self, isolated_dbgpt_home): + from dbgpt.cli._config import profile_config_path + + result = profile_config_path("openai") + assert result == isolated_dbgpt_home / "configs" / "openai.toml" + + def test_active_config_path_under_custom_home(self, isolated_dbgpt_home): + from dbgpt.cli._config import active_config_path + + result = active_config_path() + assert result == isolated_dbgpt_home / "config.toml" + + def test_default_home_is_dotdbgpt(self, monkeypatch): + monkeypatch.delenv("DBGPT_HOME", raising=False) + from pathlib import Path + + import dbgpt.cli._config as _cfg + + expected = Path.home() / ".dbgpt" + monkeypatch.setattr(_cfg, "_DBGPT_HOME", expected) + monkeypatch.setattr(_cfg, "_CONFIGS_DIR", expected / "configs") + monkeypatch.setattr(_cfg, "_ACTIVE_CONFIG", expected / "config.toml") + from dbgpt.cli._config import dbgpt_home + + assert dbgpt_home() == expected + + +class TestWorkspacePaths: + def test_render_toml_data_path_contains_workspace(self, isolated_dbgpt_home): + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="test-key") + assert "workspace/pilot/meta_data/dbgpt.db" in content + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="test-key") + assert "workspace/pilot/data" in content + + def test_render_toml_uses_custom_home_in_data_path(self, isolated_dbgpt_home): + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="test-key") + # θ·―εΎ„δΈ­εŒ…ε« isolated_dbgpt_home ηš„ε­—η¬¦δΈ² + assert str(isolated_dbgpt_home) in content + assert "workspace/pilot/meta_data/dbgpt.db" in content + + +class TestExtendedSignature: + def test_render_toml_with_llm_model_override_uses_override( + self, isolated_dbgpt_home + ): + """llm_model override replaces spec default in generated TOML.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config( + "openai", api_key="test-key", llm_model="gpt-4-turbo" + ) + data = tomlkit.loads(path.read_text()) + assert data["models"]["llms"][0]["name"] == "gpt-4-turbo" + + def test_render_toml_with_embedding_model_override_uses_override( + self, isolated_dbgpt_home + ): + """embedding_model override replaces spec default in generated TOML.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config( + "openai", api_key="test-key", embedding_model="ada-002" + ) + data = tomlkit.loads(path.read_text()) + assert data["models"]["embeddings"][0]["name"] == "ada-002" + + def test_render_toml_without_overrides_uses_spec_defaults( + self, isolated_dbgpt_home + ): + """No overrides β†’ TOML uses spec defaults (regression guard).""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config("openai", api_key="test-key") + data = tomlkit.loads(path.read_text()) + assert data["models"]["llms"][0]["name"] == "gpt-4o" + assert data["models"]["embeddings"][0]["name"] == "text-embedding-3-small" + + def test_render_toml_with_api_base_override_uses_override( + self, isolated_dbgpt_home + ): + """api_base override replaces spec default in generated TOML.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config( + "custom", api_key="test-key", api_base="https://my.api/v1" + ) + data = tomlkit.loads(path.read_text()) + assert data["models"]["llms"][0]["api_base"] == "https://my.api/v1" + + def test_render_toml_custom_api_base_derives_embedding_url( + self, isolated_dbgpt_home + ): + """When api_base is overridden, embedding api_url is derived from it.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config( + "custom", api_key="test-key", api_base="https://my-proxy.com/v1" + ) + data = tomlkit.loads(path.read_text()) + assert ( + data["models"]["embeddings"][0]["api_url"] + == "https://my-proxy.com/v1/embeddings" + ) + + def test_render_toml_custom_api_base_trailing_slash_stripped( + self, isolated_dbgpt_home + ): + """Trailing slash in api_base should not produce double-slash.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config( + "custom", api_key="test-key", api_base="https://my-proxy.com/v1/" + ) + data = tomlkit.loads(path.read_text()) + assert ( + data["models"]["embeddings"][0]["api_url"] + == "https://my-proxy.com/v1/embeddings" + ) + + def test_render_toml_without_api_base_uses_spec_embedding_url( + self, isolated_dbgpt_home + ): + """Without api_base override, embedding api_url uses spec default.""" + import tomlkit + + from dbgpt.cli._config import write_profile_config + + path = write_profile_config("openai", api_key="test-key") + data = tomlkit.loads(path.read_text()) + assert ( + data["models"]["embeddings"][0]["api_url"] + == "https://api.openai.com/v1/embeddings" + ) + + +class TestKimiEmbeddingEnvVar: + def test_kimi_embedding_uses_dashscope_env_var(self, isolated_dbgpt_home): + """Kimi profile should reference DASHSCOPE_API_KEY for embeddings.""" + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("kimi") + content = _render_profile_toml(spec, api_key=None) + assert "${env:DASHSCOPE_API_KEY:-sk-xxx}" in content + + def test_kimi_llm_uses_moonshot_env_var(self, isolated_dbgpt_home): + """Kimi profile LLM section should still reference MOONSHOT_API_KEY.""" + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("kimi") + content = _render_profile_toml(spec, api_key=None) + assert "${env:MOONSHOT_API_KEY:-sk-xxx}" in content + + def test_kimi_literal_embedding_key_overrides_env_var(self, isolated_dbgpt_home): + """When embedding_api_key is supplied, use it literally.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("kimi") + content = _render_profile_toml(spec, api_key=None, embedding_api_key="ds-key") + data = tomlkit.loads(content) + assert data["models"]["embeddings"][0]["api_key"] == "ds-key" + + def test_kimi_embedding_api_url_is_dashscope(self, isolated_dbgpt_home): + """Kimi embedding api_url should point to DashScope.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("kimi") + content = _render_profile_toml(spec, api_key=None) + data = tomlkit.loads(content) + assert data["models"]["embeddings"][0]["api_url"] == ( + "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings" + ) + + def test_openai_same_key_used_for_embeddings(self, isolated_dbgpt_home): + """OpenAI profile uses the same literal key for both LLM and embeddings.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="sk-test") + data = tomlkit.loads(content) + assert data["models"]["llms"][0]["api_key"] == "sk-test" + assert data["models"]["embeddings"][0]["api_key"] == "sk-test" + + +class TestDefaultProfileConfig: + """Tests for config generation with the 'default' (formerly skip) profile.""" + + def test_default_profile_generates_env_var_placeholder(self, isolated_dbgpt_home): + """default profile should reference ${env:OPENAI_API_KEY:-sk-xxx} (no literal key).""" + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("default") + content = _render_profile_toml(spec, api_key=None) + assert "${env:OPENAI_API_KEY:-sk-xxx}" in content + # Also verify env-var interpolation for model fields + assert "${env:LLM_MODEL_NAME:-gpt-4o}" in content + assert "${env:LLM_MODEL_PROVIDER:-proxy/openai}" in content + assert "${env:OPENAI_API_BASE:-https://api.openai.com/v1}" in content + assert "${env:EMBEDDING_MODEL_NAME:-text-embedding-3-small}" in content + assert "${env:EMBEDDING_MODEL_PROVIDER:-proxy/openai}" in content + assert ( + "${env:EMBEDDING_MODEL_API_URL:-https://api.openai.com/v1/embeddings}" + in content + ) + + def test_default_profile_generates_valid_toml(self, isolated_dbgpt_home): + """default profile should produce parseable TOML with env-var interpolation.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("default") + content = _render_profile_toml(spec, api_key=None) + data = tomlkit.loads(content) + # Model fields use env-var interpolation syntax as raw strings + assert data["models"]["llms"][0]["name"] == "${env:LLM_MODEL_NAME:-gpt-4o}" + assert ( + data["models"]["llms"][0]["provider"] + == "${env:LLM_MODEL_PROVIDER:-proxy/openai}" + ) + assert ( + data["models"]["embeddings"][0]["name"] + == "${env:EMBEDDING_MODEL_NAME:-text-embedding-3-small}" + ) + + def test_default_profile_config_file_named_default(self, isolated_dbgpt_home): + """write_profile_config('default') should create default.toml.""" + from dbgpt.cli._config import profile_config_path, write_profile_config + + path = write_profile_config("default", api_key=None) + assert path == profile_config_path("default") + assert path.name == "default.toml" + assert path.exists() + + def test_default_profile_with_literal_api_key_uses_literal_not_env( + self, isolated_dbgpt_home + ): + """When api_key is provided for default profile, use literal values.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("default") + content = _render_profile_toml(spec, api_key="sk-literal") + data = tomlkit.loads(content) + # With a literal api_key, use_env_interpolation should NOT activate + assert data["models"]["llms"][0]["name"] == "gpt-4o" + assert data["models"]["llms"][0]["provider"] == "proxy/openai" + assert data["models"]["llms"][0]["api_key"] == "sk-literal" + # No env-var syntax for model name/provider + assert "${env:LLM_MODEL_NAME" not in content + + def test_openai_profile_no_regression_with_api_key(self, isolated_dbgpt_home): + """openai profile with literal api_key must still use literal model values.""" + import tomlkit + + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key="sk-xxx") + data = tomlkit.loads(content) + assert data["models"]["llms"][0]["name"] == "gpt-4o" + assert data["models"]["llms"][0]["provider"] == "proxy/openai" + assert data["models"]["llms"][0]["api_base"] == "https://api.openai.com/v1" + assert data["models"]["llms"][0]["api_key"] == "sk-xxx" + # No env-var interpolation for openai profile + assert "${env:LLM_MODEL_NAME" not in content + + def test_openai_profile_no_regression_without_api_key(self, isolated_dbgpt_home): + """openai profile with api_key=None uses env-var for key, literal for model.""" + from dbgpt.cli._config import _render_profile_toml + from dbgpt.cli._profiles import get_profile + + spec = get_profile("openai") + content = _render_profile_toml(spec, api_key=None) + # api_key uses env-var syntax with default + assert "${env:OPENAI_API_KEY:-sk-xxx}" in content + # model name/provider/api_base remain literal + assert 'name = "gpt-4o"' in content + assert 'provider = "proxy/openai"' in content + assert 'api_base = "https://api.openai.com/v1"' in content + # No LLM_MODEL_NAME env-var for openai profile + assert "${env:LLM_MODEL_NAME" not in content diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_integration.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_integration.py new file mode 100644 index 000000000..5df755bcc --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_integration.py @@ -0,0 +1,170 @@ +"""Integration tests β€” full wizardβ†’configβ†’start flow.""" + +from __future__ import annotations + +from unittest.mock import patch + +import dbgpt.cli._config as _cfg +from dbgpt.cli._config import ( + read_active_profile, + resolve_config_path, + write_active_profile, + write_profile_config, +) +from dbgpt.cli._wizard import run_setup_noninteractive + +# All tests use `isolated_dbgpt_home` from conftest.py β€” never touches real ~/.dbgpt + + +# --------------------------------------------------------------------------- +# Test 1: First run triggers wizard and creates TOML +# --------------------------------------------------------------------------- + + +def test_first_run_triggers_wizard_creates_toml(isolated_dbgpt_home): + """First run: calling run_setup_wizard creates profile TOML + activates it.""" + from dbgpt.cli._wizard import run_setup_wizard + + # Ensure no config exists before wizard runs + assert not (_cfg._CONFIGS_DIR / "openai.toml").exists() + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key", return_value="sk-test"), + patch( + "dbgpt.cli._wizard._ask_model_names", + return_value=("gpt-4o", "text-embedding-3-small"), + ), + patch("dbgpt.cli._wizard._print_welcome"), + patch( + "dbgpt.cli._wizard._ask_api_base", + return_value="https://api.openai.com/v1", + ), + ): + from dbgpt.cli._profiles import get_profile + + mock_ask_profile.return_value = get_profile("openai") + + run_setup_wizard() + + # openai.toml must exist under the isolated configs dir + assert (_cfg._CONFIGS_DIR / "openai.toml").exists() + # Active profile must be set to "openai" + assert read_active_profile() == "openai" + + +# --------------------------------------------------------------------------- +# Test 2: Profile switch then start resolves correct config +# --------------------------------------------------------------------------- + + +def test_profile_switch_then_start_resolves_config(isolated_dbgpt_home): + """Switching active profile makes resolve_config_path return the new profile path.""" # noqa: E501 + # Create two TOML files + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\n", encoding="utf-8") + (_cfg._CONFIGS_DIR / "qwen.toml").write_text("[models]\n", encoding="utf-8") + + # Switch to qwen + write_active_profile("qwen") + + # resolve_config_path with no args should pick up qwen + result = resolve_config_path() + assert result is not None + assert "qwen.toml" in result + + +# --------------------------------------------------------------------------- +# Test 3: start web help shows config/profile options +# --------------------------------------------------------------------------- + + +def test_start_web_invokes_webserver_with_resolved_config(isolated_dbgpt_home): + """start web --help exits 0 and shows relevant option text.""" + import click.testing + + from dbgpt.cli.cli_scripts import cli + + runner = click.testing.CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["start", "web", "--help"]) + + # --help should always succeed (exit 0) and show options + assert result.exit_code == 0 + # Should show at least one of the expected flags + help_text = result.output + assert "--config" in help_text or "--profile" in help_text or "--yes" in help_text + + +# --------------------------------------------------------------------------- +# Test 4: CLI --profile flag overrides active profile +# --------------------------------------------------------------------------- + + +def test_cli_flag_overrides_active_profile(isolated_dbgpt_home): + """resolve_config_path(profile=...) returns that profile's path, not the active one.""" # noqa: E501 + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\n", encoding="utf-8") + (_cfg._CONFIGS_DIR / "qwen.toml").write_text("[models]\n", encoding="utf-8") + + # Active profile is openai + write_active_profile("openai") + + # But we pass profile="qwen" explicitly + result = resolve_config_path(profile="qwen") + assert result is not None + assert "qwen.toml" in result + assert "openai.toml" not in result + + +# --------------------------------------------------------------------------- +# Test 5: Explicit --config flag overrides everything +# --------------------------------------------------------------------------- + + +def test_explicit_config_flag_overrides_all(isolated_dbgpt_home): + """resolve_config_path(config=...) returns the exact path regardless of active profile.""" # noqa: E501 + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\n", encoding="utf-8") + write_active_profile("openai") + + custom_path = "/custom/path.toml" + result = resolve_config_path(config=custom_path) + assert result == custom_path + + +# --------------------------------------------------------------------------- +# Test 6: Non-interactive mode creates default profile +# --------------------------------------------------------------------------- + + +def test_noninteractive_mode_creates_default_profile(isolated_dbgpt_home): + """run_setup_noninteractive writes openai.toml with the provided API key.""" + # Confirm no config exists first + assert not (_cfg._CONFIGS_DIR / "openai.toml").exists() + + run_setup_noninteractive(profile_name="openai", api_key="sk-test") + + # Profile TOML must be created + toml_path = _cfg._CONFIGS_DIR / "openai.toml" + assert toml_path.exists() + content = toml_path.read_text(encoding="utf-8") + assert "sk-test" in content + + +# --------------------------------------------------------------------------- +# Test 7: Skip provider creates minimal config +# --------------------------------------------------------------------------- + + +def test_skip_provider_creates_minimal_config(isolated_dbgpt_home): + """write_profile_config('default', ...) creates default.toml and activates it.""" + # Confirm no config exists + assert not (_cfg._CONFIGS_DIR / "default.toml").exists() + + write_profile_config("default", api_key=None, activate=True) + + # default.toml must exist + toml_path = _cfg._CONFIGS_DIR / "default.toml" + assert toml_path.exists() + # Active profile must be "default" + assert read_active_profile() == "default" diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_profile_cmd.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_profile_cmd.py new file mode 100644 index 000000000..3d242d553 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_profile_cmd.py @@ -0,0 +1,126 @@ +"""Tests for dbgpt profile subcommands.""" + +from click.testing import CliRunner + +from dbgpt.cli.cli_scripts import cli + + +class TestProfileList: + def test_profile_list_with_no_configs_shows_empty_message( + self, isolated_dbgpt_home + ): + """profile list with no configs dir β†’ 'No profiles configured'.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "list"]) + assert result.exit_code == 0, result.output + assert ( + "no profile" in result.output.lower() + or "no config" in result.output.lower() + ) + + def test_profile_list_shows_all_profiles(self, isolated_dbgpt_home): + """profile list with 2 TOML files β†’ both shown.""" + # Create fake TOML files + import dbgpt.cli._config as _cfg + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\n") + (_cfg._CONFIGS_DIR / "kimi.toml").write_text("[models]\n") + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "list"]) + assert result.exit_code == 0, result.output + assert "openai" in result.output + assert "kimi" in result.output + + def test_profile_list_marks_active_with_asterisk(self, isolated_dbgpt_home): + """profile list marks active profile with *.""" + import dbgpt.cli._config as _cfg + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\n") + (_cfg._CONFIGS_DIR / "kimi.toml").write_text("[models]\n") + _cfg._ACTIVE_CONFIG.write_text('[default]\nprofile = "openai"\n') + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "list"]) + assert result.exit_code == 0, result.output + # openai should have * marker + lines = result.output.splitlines() + openai_line = next((line for line in lines if "openai" in line), "") + assert "*" in openai_line, f"Expected * in openai line, got: {openai_line!r}" + + +class TestProfileShow: + def test_profile_show_prints_toml_content(self, isolated_dbgpt_home): + """profile show openai β†’ prints TOML file content.""" + import dbgpt.cli._config as _cfg + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "openai.toml").write_text("[models]\nllms = []\n") + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "show", "openai"]) + assert result.exit_code == 0, result.output + assert "[models]" in result.output + + def test_profile_show_nonexistent_exits_nonzero(self, isolated_dbgpt_home): + """profile show nonexistent β†’ exit_code != 0.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "show", "nonexistent"]) + assert result.exit_code != 0 + + +class TestProfileSwitch: + def test_profile_switch_updates_active(self, isolated_dbgpt_home): + """profile switch kimi β†’ kimi becomes active.""" + import dbgpt.cli._config as _cfg + from dbgpt.cli._config import read_active_profile + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + (_cfg._CONFIGS_DIR / "kimi.toml").write_text("[models]\n") + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "switch", "kimi"]) + assert result.exit_code == 0, result.output + assert read_active_profile() == "kimi" + + def test_profile_switch_nonexistent_exits_nonzero(self, isolated_dbgpt_home): + """profile switch nonexistent β†’ exit_code != 0, suggests create.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "switch", "nonexistent"]) + assert result.exit_code != 0 + assert ( + "create" in result.output.lower() or "nonexistent" in result.output.lower() + ) + + +class TestProfileDelete: + def test_profile_delete_with_yes_removes_file(self, isolated_dbgpt_home): + """profile delete openai --yes β†’ openai.toml removed.""" + import dbgpt.cli._config as _cfg + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + toml_path = _cfg._CONFIGS_DIR / "openai.toml" + toml_path.write_text("[models]\n") + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "delete", "openai", "--yes"]) + assert result.exit_code == 0, result.output + assert not toml_path.exists() + + def test_profile_delete_active_clears_active_pointer(self, isolated_dbgpt_home): + """profile delete active profile β†’ clears active pointer.""" + import dbgpt.cli._config as _cfg + from dbgpt.cli._config import read_active_profile + + _cfg._CONFIGS_DIR.mkdir(parents=True, exist_ok=True) + toml_path = _cfg._CONFIGS_DIR / "openai.toml" + toml_path.write_text("[models]\n") + _cfg._ACTIVE_CONFIG.write_text('[default]\nprofile = "openai"\n') + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "delete", "openai", "--yes"]) + assert result.exit_code == 0, result.output + assert not toml_path.exists() + assert read_active_profile() is None + + def test_profile_delete_nonexistent_exits_nonzero(self, isolated_dbgpt_home): + """profile delete nonexistent --yes β†’ exit_code != 0.""" + runner = CliRunner(mix_stderr=False) + result = runner.invoke(cli, ["profile", "delete", "nonexistent", "--yes"]) + assert result.exit_code != 0 diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_profiles.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_profiles.py new file mode 100644 index 000000000..f25dbdecd --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_profiles.py @@ -0,0 +1,175 @@ +"""Tests for _profiles.py β€” GLM/Custom/Default profiles and label fixes.""" + +import pytest + +from dbgpt.cli._profiles import ProfileSpec, get_profile, list_profiles + + +class TestGlmProfile: + """Tests for the glm profile.""" + + def test_glm_profile_spec(self): + """get_profile('glm') should return a ProfileSpec with correct fields.""" + p = get_profile("glm") + assert isinstance(p, ProfileSpec) + assert p.name == "glm" + assert p.label == "z.ai (zhipu.ai API)" + assert p.env_var == "ZHIPUAI_API_KEY" + assert p.llm_model == "glm-4-plus" + assert p.llm_provider == "proxy/zhipu" + assert p.llm_api_base == "https://open.bigmodel.cn/api/paas/v4" + assert p.embedding_model == "embedding-3" + assert p.embedding_provider == "proxy/zhipu" + assert p.needs_api_key is True + + def test_glm_case_insensitive(self): + """get_profile('GLM') should work the same as get_profile('glm').""" + p_lower = get_profile("glm") + p_upper = get_profile("GLM") + assert p_lower.name == p_upper.name + assert p_lower.label == p_upper.label + assert p_lower.env_var == p_upper.env_var + + +class TestCustomProfile: + """Tests for the custom profile.""" + + def test_custom_profile_spec(self): + """get_profile('custom') should return a ProfileSpec with correct fields.""" + p = get_profile("custom") + assert isinstance(p, ProfileSpec) + assert p.name == "custom" + assert p.label == "Custom Provider (Any OpenAI compatible endpoint)" + assert p.env_var == "OPENAI_API_KEY" + assert p.llm_model == "gpt-4o" + assert p.llm_provider == "proxy/openai" + assert p.llm_api_base == "https://api.openai.com/v1" + assert p.needs_api_key is True + + +class TestDefaultProfile: + """Tests for the default (formerly skip) profile.""" + + def test_default_profile_spec(self): + """get_profile('default') returns a ProfileSpec with needs_api_key=False.""" + p = get_profile("default") + assert isinstance(p, ProfileSpec) + assert p.name == "default" + assert p.needs_api_key is False + + def test_default_profile_has_real_openai_values(self): + """default profile should have real OpenAI values for config generation.""" + p = get_profile("default") + assert p.env_var == "OPENAI_API_KEY" + assert p.llm_model == "gpt-4o" + assert p.llm_provider == "proxy/openai" + assert p.llm_api_base == "https://api.openai.com/v1" + assert p.embedding_model == "text-embedding-3-small" + assert p.embedding_provider == "proxy/openai" + + def test_default_profile_label(self): + """default profile label should mention 'Skip for now'.""" + p = get_profile("default") + assert "Skip for now" in p.label + + def test_skip_profile_no_longer_exists(self): + """'skip' is no longer a valid profile name.""" + with pytest.raises(ValueError): + get_profile("skip") + + +class TestProfileOrder: + """Tests for profile ordering.""" + + def test_profile_order_ends_with_new_profiles(self): + """list_profiles() should end with glm, custom, default.""" + profiles = list_profiles() + names = [p.name for p in profiles] + assert names[-3:] == ["glm", "custom", "default"] + + +class TestExistingLabels: + """Tests for fixed labels on existing profiles.""" + + def test_openai_label_fixed(self): + """openai profile label should be updated.""" + p = get_profile("openai") + assert p.label == "OpenAI (OpenAI or OpenAI API proxy)" + + def test_qwen_label_fixed(self): + """qwen profile label should be updated.""" + p = get_profile("qwen") + assert p.label == "Qwen (DashScope API)" + + +class TestErrorHandling: + """Tests for error handling.""" + + def test_unknown_profile_raises_value_error(self): + """get_profile with unknown name should raise ValueError.""" + with pytest.raises(ValueError): + get_profile("nonexistent") + + +class TestKimiEmbedding: + """Tests for Kimi embedding configuration using DashScope.""" + + def test_kimi_embedding_uses_tongyi_provider(self): + p = get_profile("kimi") + assert p.embedding_provider == "proxy/tongyi" + + def test_kimi_embedding_model_is_text_embedding_v3(self): + p = get_profile("kimi") + assert p.embedding_model == "text-embedding-v3" + + def test_kimi_embedding_api_url_is_dashscope(self): + p = get_profile("kimi") + assert ( + p.embedding_api_url + == "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings" + ) + + def test_kimi_has_separate_embedding_env_var(self): + p = get_profile("kimi") + assert p.embedding_env_var == "DASHSCOPE_API_KEY" + assert p.embedding_env_var != p.env_var + + def test_kimi_llm_still_uses_moonshot_env_var(self): + p = get_profile("kimi") + assert p.env_var == "MOONSHOT_API_KEY" + + +class TestMinimaxEmbedding: + """Tests for MiniMax embedding configuration.""" + + def test_minimax_embedding_model_is_embo_01(self): + p = get_profile("minimax") + assert p.embedding_model == "embo-01" + + def test_minimax_embedding_provider_is_openai(self): + p = get_profile("minimax") + assert p.embedding_provider == "proxy/openai" + + def test_minimax_embedding_api_url_is_minimax(self): + p = get_profile("minimax") + assert p.embedding_api_url == "https://api.minimax.chat/v1/embeddings" + + def test_minimax_no_separate_embedding_env_var(self): + p = get_profile("minimax") + assert p.embedding_env_var is None + + +class TestEmbeddingEnvVarDefault: + """Tests for the embedding_env_var default behaviour.""" + + def test_openai_has_no_separate_embedding_env_var(self): + p = get_profile("openai") + assert p.embedding_env_var is None + + def test_qwen_has_no_separate_embedding_env_var(self): + p = get_profile("qwen") + assert p.embedding_env_var is None + + def test_glm_has_no_separate_embedding_env_var(self): + p = get_profile("glm") + assert p.embedding_env_var is None diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_smoke.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_smoke.py new file mode 100644 index 000000000..f8ef130d1 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_smoke.py @@ -0,0 +1,43 @@ +"""Smoke tests for CLI test infrastructure β€” verifies fixture isolation.""" + +import os +from pathlib import Path + + +def test_isolated_dbgpt_home_is_not_real_home(isolated_dbgpt_home): + """Verify that isolated_dbgpt_home does NOT point to the real ~/.dbgpt.""" + real_home = Path.home() / ".dbgpt" + assert isolated_dbgpt_home != real_home, ( + f"isolated_dbgpt_home should not be the real home: {real_home}" + ) + + +def test_isolated_dbgpt_home_env_var(isolated_dbgpt_home): + """Verify DBGPT_HOME env var is set to the isolated path.""" + env_home = os.environ.get("DBGPT_HOME") + assert env_home is not None, "DBGPT_HOME env var should be set" + assert env_home == str(isolated_dbgpt_home), ( + f"DBGPT_HOME ({env_home}) should match isolated_dbgpt_home ({isolated_dbgpt_home})" # noqa: E501 + ) + + +def test_isolated_dbgpt_home_patches_config_module(isolated_dbgpt_home): + """Verify _config module constants are patched to isolated temp path.""" + import dbgpt.cli._config as _cfg + + assert _cfg._DBGPT_HOME == isolated_dbgpt_home, ( + f"_config._DBGPT_HOME ({_cfg._DBGPT_HOME}) should equal isolated home" + ) + assert _cfg._CONFIGS_DIR == isolated_dbgpt_home / "configs", ( + "_config._CONFIGS_DIR should be under isolated home" + ) + assert _cfg._ACTIVE_CONFIG == isolated_dbgpt_home / "config.toml", ( + "_config._ACTIVE_CONFIG should be under isolated home" + ) + + +def test_cli_runner_returns_click_runner(cli_runner): + """Verify cli_runner fixture returns a CliRunner instance.""" + import click.testing + + assert isinstance(cli_runner, click.testing.CliRunner) diff --git a/packages/dbgpt-core/src/dbgpt/cli/tests/test_wizard.py b/packages/dbgpt-core/src/dbgpt/cli/tests/test_wizard.py new file mode 100644 index 000000000..94077fbf2 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/cli/tests/test_wizard.py @@ -0,0 +1,251 @@ +"""Tests for _wizard.py β€” model config step and skip/default flow.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from dbgpt.cli._profiles import get_profile + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_log_ask(*return_values): + """Return a mock _log whose .ask() yields values in sequence.""" + mock_log = MagicMock() + mock_log.ask.side_effect = list(return_values) + return mock_log + + +# --------------------------------------------------------------------------- +# _ask_model_names +# --------------------------------------------------------------------------- + + +def test_ask_model_names_returns_defaults_on_enter(): + """User presses Enter β†’ default values from spec returned.""" + from dbgpt.cli._wizard import _ask_model_names + + spec = get_profile("openai") + + mock_log = _make_log_ask("gpt-4o", "text-embedding-3-small") + with patch("dbgpt.cli._wizard._log", mock_log): + llm, emb = _ask_model_names(spec) + + assert llm == "gpt-4o" + assert emb == "text-embedding-3-small" + mock_log.ask.assert_any_call("LLM model name", default="gpt-4o") + mock_log.ask.assert_any_call( + "Embedding model name", default="text-embedding-3-small" + ) + + +def test_ask_model_names_returns_custom_values(): + """User types custom model names β†’ returned as-is.""" + from dbgpt.cli._wizard import _ask_model_names + + spec = get_profile("openai") + + mock_log = _make_log_ask("my-custom-llm", "my-custom-emb") + with patch("dbgpt.cli._wizard._log", mock_log): + llm, emb = _ask_model_names(spec) + + assert llm == "my-custom-llm" + assert emb == "my-custom-emb" + + +# --------------------------------------------------------------------------- +# run_setup_wizard β€” model names integration +# --------------------------------------------------------------------------- + + +def test_run_setup_wizard_calls_ask_model_names(isolated_dbgpt_home): + """run_setup_wizard passes llm_model/embedding_model to write_profile_config.""" + from dbgpt.cli._wizard import run_setup_wizard + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key", return_value="sk-test"), + patch( + "dbgpt.cli._wizard._ask_api_base", return_value="https://api.openai.com/v1" + ), + patch( + "dbgpt.cli._wizard._ask_model_names", + return_value=("gpt-4o", "text-embedding-3-small"), + ) as mock_model, + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "openai.toml", + ) as mock_write, + patch("dbgpt.cli._wizard._print_welcome"), + ): + spec = get_profile("openai") + mock_ask_profile.return_value = spec + + run_setup_wizard() + + mock_model.assert_called_once_with(spec) + mock_write.assert_called_once_with( + "openai", + api_key="sk-test", + activate=True, + llm_model="gpt-4o", + embedding_model="text-embedding-3-small", + api_base="https://api.openai.com/v1", + embedding_api_key=None, + ) + + +# --------------------------------------------------------------------------- +# Default (formerly skip) flow +# --------------------------------------------------------------------------- + + +def test_default_profile_bypasses_all_prompts(isolated_dbgpt_home): + """default spec β†’ _ask_api_key and _ask_model_names never called.""" + from dbgpt.cli._wizard import run_setup_wizard + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key") as mock_api_key, + patch("dbgpt.cli._wizard._ask_model_names") as mock_model, + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "default.toml", + ), + patch("dbgpt.cli._wizard._print_welcome"), + ): + mock_ask_profile.return_value = get_profile("default") + + run_setup_wizard() + + mock_api_key.assert_not_called() + mock_model.assert_not_called() + + +# --------------------------------------------------------------------------- +# run_setup_noninteractive β€” unchanged +# --------------------------------------------------------------------------- + + +def test_run_setup_noninteractive_unchanged(isolated_dbgpt_home): + """run_setup_noninteractive does NOT call _ask_model_names.""" + from dbgpt.cli._wizard import run_setup_noninteractive + + with ( + patch("dbgpt.cli._wizard._ask_model_names") as mock_model, + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "openai.toml", + ), + ): + run_setup_noninteractive("openai", "sk-xxx") + + mock_model.assert_not_called() + + +# --------------------------------------------------------------------------- +# Custom profile β€” api_base +# --------------------------------------------------------------------------- + + +def test_custom_profile_asks_api_base(isolated_dbgpt_home): + """custom spec β†’ _ask_api_base is called and result passed to write_profile_config.""" # noqa: E501 + from dbgpt.cli._wizard import run_setup_wizard + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key", return_value="sk-custom"), + patch( + "dbgpt.cli._wizard._ask_api_base", return_value="https://my.proxy.com/v1" + ) as mock_api_base, + patch( + "dbgpt.cli._wizard._ask_model_names", + return_value=("gpt-4o", "text-embedding-3-small"), + ), + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "custom.toml", + ) as mock_write, + patch("dbgpt.cli._wizard._print_welcome"), + ): + spec = get_profile("custom") + mock_ask_profile.return_value = spec + + run_setup_wizard() + + mock_api_base.assert_called_once_with(spec) + mock_write.assert_called_once_with( + "custom", + api_key="sk-custom", + activate=True, + llm_model="gpt-4o", + embedding_model="text-embedding-3-small", + api_base="https://my.proxy.com/v1", + embedding_api_key=None, + ) + + +def test_openai_profile_asks_api_base(isolated_dbgpt_home): + """openai spec β†’ _ask_api_base IS called (users often use proxies).""" + from dbgpt.cli._wizard import run_setup_wizard + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key", return_value="sk-openai"), + patch( + "dbgpt.cli._wizard._ask_api_base", + return_value="https://api.openai.com/v1", + ) as mock_api_base, + patch( + "dbgpt.cli._wizard._ask_model_names", + return_value=("gpt-4o", "text-embedding-3-small"), + ), + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "openai.toml", + ) as mock_write, + patch("dbgpt.cli._wizard._print_welcome"), + ): + spec = get_profile("openai") + mock_ask_profile.return_value = spec + + run_setup_wizard() + + mock_api_base.assert_called_once_with(spec) + mock_write.assert_called_once_with( + "openai", + api_key="sk-openai", + activate=True, + llm_model="gpt-4o", + embedding_model="text-embedding-3-small", + api_base="https://api.openai.com/v1", + embedding_api_key=None, + ) + + +def test_kimi_profile_does_not_ask_api_base(isolated_dbgpt_home): + """kimi spec β†’ _ask_api_base is NOT called.""" + from dbgpt.cli._wizard import run_setup_wizard + + with ( + patch("dbgpt.cli._wizard._ask_profile") as mock_ask_profile, + patch("dbgpt.cli._wizard._ask_api_key", return_value="sk-test"), + patch("dbgpt.cli._wizard._ask_api_base") as mock_api_base, + patch( + "dbgpt.cli._wizard._ask_model_names", + return_value=("kimi-k2", "text-embedding-v3"), + ), + patch("dbgpt.cli._wizard._ask_embedding_api_key", return_value="ds-key"), + patch( + "dbgpt.cli._wizard.write_profile_config", + return_value=isolated_dbgpt_home / "configs" / "kimi.toml", + ), + patch("dbgpt.cli._wizard._print_welcome"), + ): + mock_ask_profile.return_value = get_profile("kimi") + + run_setup_wizard() + + mock_api_base.assert_not_called() diff --git a/packages/dbgpt-core/src/dbgpt/configs/model_config.py b/packages/dbgpt-core/src/dbgpt/configs/model_config.py index 869469440..367d29c70 100644 --- a/packages/dbgpt-core/src/dbgpt/configs/model_config.py +++ b/packages/dbgpt-core/src/dbgpt/configs/model_config.py @@ -5,13 +5,34 @@ import os from functools import cache from typing import Optional -ROOT_PATH = os.path.dirname( - os.path.dirname( + +def _detect_root_path() -> str: + """Detect the root path of the DB-GPT installation. + + Determines whether running from a source checkout or a pip install, + and returns the appropriate root path. + + Returns: + str: The repo root directory for source installs, or + ``DBGPT_HOME/workspace`` (defaulting to ``~/.dbgpt/workspace``) + for pip installs. + """ + candidate = os.path.dirname( os.path.dirname( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + os.path.dirname( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + ) ) ) -) + if os.path.isfile(os.path.join(candidate, "pyproject.toml")): + return candidate + home = os.environ.get("DBGPT_HOME", os.path.expanduser("~/.dbgpt")) + return os.path.join(home, "workspace") + + +ROOT_PATH = _detect_root_path() MODEL_PATH = os.path.join(ROOT_PATH, "models") PILOT_PATH = os.path.join(ROOT_PATH, "pilot") LOGDIR = os.getenv("DBGPT_LOG_DIR", os.path.join(ROOT_PATH, "logs")) diff --git a/packages/dbgpt-core/src/dbgpt/configs/tests/__init__.py b/packages/dbgpt-core/src/dbgpt/configs/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/dbgpt-core/src/dbgpt/configs/tests/test_model_config.py b/packages/dbgpt-core/src/dbgpt/configs/tests/test_model_config.py new file mode 100644 index 000000000..04db43e87 --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/configs/tests/test_model_config.py @@ -0,0 +1,134 @@ +"""Tests for _detect_root_path() in model_config.py.""" + +import os + +# --------------------------------------------------------------------------- +# Helper: call _detect_root_path() with a controlled filesystem view +# --------------------------------------------------------------------------- + + +def _call_detect(monkeypatch, candidate: str, dbgpt_home: str | None = None) -> str: + """Invoke _detect_root_path() with the given candidate and env. + + Args: + monkeypatch: pytest monkeypatch fixture for isolation. + candidate: The fake "6-level dirname" that the function would compute + from ``__file__``. + dbgpt_home: If not None, set DBGPT_HOME env var to this value. + If None, remove DBGPT_HOME from the environment. + + Returns: + str: The return value of ``_detect_root_path()``. + """ + import dbgpt.configs.model_config as _mc + + # Patch os.path.abspath so that the dirname chain resolves to `candidate` + # The function calls os.path.dirname 6 times on os.path.abspath(__file__). + # We make abspath return a path deep enough that 6 dirname calls yield + # exactly `candidate`. + fake_deep_path = os.path.join(candidate, "a", "b", "c", "d", "e", "model_config.py") + + monkeypatch.setattr( + _mc.os.path, + "abspath", + lambda _path: fake_deep_path, + ) + + if dbgpt_home is None: + monkeypatch.delenv("DBGPT_HOME", raising=False) + else: + monkeypatch.setenv("DBGPT_HOME", dbgpt_home) + + return _mc._detect_root_path() + + +# --------------------------------------------------------------------------- +# Test 1: source install β€” pyproject.toml found at candidate +# --------------------------------------------------------------------------- + + +def test_source_install_returns_candidate(tmp_path, monkeypatch): + """Source install: pyproject.toml at candidate β†’ ROOT_PATH == candidate. + + Creates a fake repo root (tmp_path) with a pyproject.toml sentinel and + verifies that _detect_root_path() returns that directory. + """ + candidate = str(tmp_path) + (tmp_path / "pyproject.toml").write_text("[tool]\n") + + result = _call_detect(monkeypatch, candidate) + + assert result == candidate + + +# --------------------------------------------------------------------------- +# Test 2: pip install default β€” no pyproject.toml, no DBGPT_HOME +# --------------------------------------------------------------------------- + + +def test_pip_install_default_returns_dot_dbgpt_workspace(tmp_path, monkeypatch): + """Pip install (default): no pyproject.toml, no DBGPT_HOME β†’ ~/.dbgpt/workspace. + + Ensures the fallback path is ``~/.dbgpt/workspace`` when the candidate + directory does not contain a pyproject.toml and DBGPT_HOME is not set. + """ + # candidate has no pyproject.toml + candidate = str(tmp_path) + + result = _call_detect(monkeypatch, candidate, dbgpt_home=None) + + expected = os.path.join(os.path.expanduser("~/.dbgpt"), "workspace") + assert result == expected + + +# --------------------------------------------------------------------------- +# Test 3: pip install with custom DBGPT_HOME +# --------------------------------------------------------------------------- + + +def test_pip_install_custom_dbgpt_home_returns_workspace_under_home( + tmp_path, monkeypatch +): + """Pip install + DBGPT_HOME: no pyproject.toml + DBGPT_HOME β†’ DBGPT_HOME/workspace. + + Verifies that when DBGPT_HOME is set to a custom path and pyproject.toml + does not exist at the candidate, the result is ``DBGPT_HOME/workspace``. + """ + candidate = str(tmp_path / "candidate") + os.makedirs(candidate, exist_ok=True) + # no pyproject.toml in candidate + + custom_home = str(tmp_path / "custom_home") + + result = _call_detect(monkeypatch, candidate, dbgpt_home=custom_home) + + assert result == os.path.join(custom_home, "workspace") + + +# --------------------------------------------------------------------------- +# Test 4: derived constants follow ROOT_PATH +# --------------------------------------------------------------------------- + + +def test_derived_constants_follow_root_path(tmp_path, monkeypatch): + """Derived constants PILOT_PATH and STATIC_MESSAGE_IMG_PATH follow ROOT_PATH. + + After _detect_root_path() would return a given path, PILOT_PATH and + STATIC_MESSAGE_IMG_PATH should be derived consistently from it. + Validates the relationship defined in model_config.py: + + PILOT_PATH = ROOT_PATH + "/pilot" + STATIC_MESSAGE_IMG_PATH = PILOT_PATH + "/message/img" + """ + candidate = str(tmp_path) + (tmp_path / "pyproject.toml").write_text("[tool]\n") + + root = _call_detect(monkeypatch, candidate) + + expected_pilot = os.path.join(root, "pilot") + expected_img = os.path.join(expected_pilot, "message/img") + + # Verify the relationships (not the module-level cached constants, which + # were evaluated at import time with the real filesystem). + assert os.path.join(root, "pilot") == expected_pilot + assert os.path.join(expected_pilot, "message/img") == expected_img diff --git a/packages/dbgpt-core/src/dbgpt/util/_db_migration_utils.py b/packages/dbgpt-core/src/dbgpt/util/_db_migration_utils.py index 39428c454..167aecde7 100644 --- a/packages/dbgpt-core/src/dbgpt/util/_db_migration_utils.py +++ b/packages/dbgpt-core/src/dbgpt/util/_db_migration_utils.py @@ -35,14 +35,14 @@ def create_alembic_config( alembic_ini_path = alembic_ini_path or os.path.join( alembic_root_path, "alembic.ini" ) - alembic_cfg = AlembicConfig(alembic_ini_path) - alembic_cfg.set_main_option("sqlalchemy.url", str(engine.url)) script_location = script_location or os.path.join(alembic_root_path, "alembic") versions_dir = os.path.join(script_location, "versions") os.makedirs(script_location, exist_ok=True) os.makedirs(versions_dir, exist_ok=True) + alembic_cfg = AlembicConfig(alembic_ini_path) + alembic_cfg.set_main_option("sqlalchemy.url", str(engine.url)) alembic_cfg.set_main_option("script_location", script_location) alembic_cfg.attributes["target_metadata"] = base.metadata diff --git a/packages/dbgpt-core/src/dbgpt/util/console/console.py b/packages/dbgpt-core/src/dbgpt/util/console/console.py index 423c22d11..49486cc76 100644 --- a/packages/dbgpt-core/src/dbgpt/util/console/console.py +++ b/packages/dbgpt-core/src/dbgpt/util/console/console.py @@ -3,7 +3,7 @@ import dataclasses import sys from functools import lru_cache -from typing import Any +from typing import Any, Callable, List, Optional, Tuple from rich.console import Console from rich.markdown import Markdown @@ -40,6 +40,66 @@ def get_console(output: Output | None = None) -> Console: ) +# --------------------------------------------------------------------------- +# Terminal raw-mode helpers (stdlib only) +# --------------------------------------------------------------------------- + + +def _supports_raw_mode() -> bool: + """Return True if stdin supports raw mode (TTY and termios available).""" + if not sys.stdin.isatty(): + return False + try: + import termios # noqa: F401 + import tty # noqa: F401 + + return True + except ImportError: + return False + + +def _read_key() -> str: + """Read a single keypress from stdin in raw mode. + + Returns: + str: One of ``'up'``, ``'down'``, ``'enter'``, or the raw character. + + Raises: + KeyboardInterrupt: On Ctrl-C. + """ + import termios + import tty + + fd = sys.stdin.fileno() + old = termios.tcgetattr(fd) + try: + tty.setraw(fd) + ch = sys.stdin.read(1) + if ch == "\x1b": + ch2 = sys.stdin.read(1) + ch3 = sys.stdin.read(1) + if ch2 == "[": + if ch3 == "A": + return "up" + if ch3 == "B": + return "down" + return ch + if ch in ("\r", "\n"): + return "enter" + if ch == "\x03": + raise KeyboardInterrupt + return ch + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old) + + +def _move_up(console: Console, lines: int) -> None: + """Move terminal cursor up *lines* rows and clear those lines.""" + # ANSI: move up N lines then erase to end of screen + console.file.write(f"\x1b[{lines}A\x1b[0J") + console.file.flush() + + class CliLogger: def __init__(self, output: Output | None = None): self.console = get_console(output) @@ -70,3 +130,99 @@ class CliLogger: def ask(self, msg: str, **kwargs): return Prompt.ask(msg, **kwargs) + + def select( + self, + prompt: str, + options: List[Tuple[str, str]], + _read_key_fn: Optional[Callable[[], str]] = None, + ) -> int: + """Display an interactive arrow-key selector and return the chosen index. + + Args: + prompt (str): Header text displayed above the option list. + options (List[Tuple[str, str]]): Each tuple is (name, description). + ``name`` is the display name rendered in bold. + ``description`` is dim helper text shown after the name. + _read_key_fn (Optional[Callable[[], str]]): Override the key-reading + function (used for testing). Should return one of: ``'up'``, + ``'down'``, ``'enter'``, or a single character string. + + Returns: + int: Zero-based index of the selected option. + """ + read_key = _read_key_fn or _read_key + + # If raw-mode isn't available, fall back to numbered input. + if not _supports_raw_mode(): + return self._select_fallback(prompt, options) + + current = 0 + n = len(options) + + self.console.print(f"\n {prompt}\n") + + def _render(idx: int, move_up_lines: int = 0) -> None: + buf = "" + if move_up_lines > 0: + buf += f"\x1b[{move_up_lines}A\x1b[0J" + for i, (name, desc) in enumerate(options): + if i == idx: + marker = "\x1b[1;96m●\x1b[0m" + else: + marker = "\x1b[2mβ—‹\x1b[0m" + bold_name = f"\x1b[1m{name}\x1b[0m" + dim_desc = f"\x1b[2m{desc}\x1b[0m" + buf += f" {marker} {bold_name:<20}{dim_desc}\n" + self.console.file.write(buf) + self.console.file.flush() + + _render(current) + + while True: + try: + key = read_key() + except KeyboardInterrupt: + raise + + if key == "up": + current = (current - 1) % n + _render(current, move_up_lines=n) + elif key == "down": + current = (current + 1) % n + _render(current, move_up_lines=n) + elif key == "enter": + _render(current, move_up_lines=n) + self.console.print("") + return current + elif key.isdigit(): + num = int(key) + if 1 <= num <= n: + current = num - 1 + _render(current, move_up_lines=n) + self.console.print("") + return current + + def _select_fallback( + self, + prompt: str, + options: List[Tuple[str, str]], + ) -> int: + """Numbered fallback for non-TTY environments.""" + self.console.print(f"\n {prompt}\n") + for i, (name, desc) in enumerate(options, start=1): + self.console.print( + f" [[bold]{i}[/bold]] [bold]{name}[/bold] [dim]{desc}[/dim]" + ) + self.console.print("") + while True: + raw = Prompt.ask("Enter a number", default="1") + try: + choice = int(raw) + if 1 <= choice <= len(options): + return choice - 1 + except (ValueError, TypeError): + pass + self.console.print( + f"[warning]Please enter a number between 1 and {len(options)}.[/]" + ) diff --git a/packages/dbgpt-core/src/dbgpt/util/tests/test_db_migration_utils.py b/packages/dbgpt-core/src/dbgpt/util/tests/test_db_migration_utils.py new file mode 100644 index 000000000..65ff1c2fb --- /dev/null +++ b/packages/dbgpt-core/src/dbgpt/util/tests/test_db_migration_utils.py @@ -0,0 +1,30 @@ +"""Tests for _db_migration_utils create_alembic_config behaviour.""" + +from unittest.mock import MagicMock, patch + + +def test_create_alembic_config_creates_versions_dir(tmp_path): + """create_alembic_config must create script_location and versions/ directories.""" + from dbgpt.util._db_migration_utils import create_alembic_config + + mock_engine = MagicMock() + mock_engine.url = "sqlite:///test.db" + mock_base = MagicMock() + mock_base.metadata = MagicMock() + mock_session = MagicMock() + + alembic_root = str(tmp_path) + + with patch("dbgpt.util._db_migration_utils.AlembicConfig") as mock_alembic_cfg_cls: + mock_cfg = MagicMock() + mock_alembic_cfg_cls.return_value = mock_cfg + + result = create_alembic_config( + alembic_root, mock_engine, mock_base, mock_session + ) + + alembic_dir = tmp_path / "alembic" + versions_dir = alembic_dir / "versions" + assert alembic_dir.exists(), "script_location directory must be created" + assert versions_dir.exists(), "versions directory must be created" + assert result is mock_cfg