github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2026-07-16 17:15:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add dbgpt app start cli

Browse Source
This commit is contained in:
alan.cl
2026-03-20 16:50:48 +08:00
parent 03e72271f5
commit aaa32f00d2
28 changed files with 3069 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
packages/dbgpt-app/src/dbgpt_app/_cli.py
View File

@@ -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,

21
packages/dbgpt-app/src/dbgpt_app/base.py
View File

@@ -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

41
packages/dbgpt-app/src/dbgpt_app/initialization/workspace_provisioning.py Normal file
View File

@@ -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)

81
packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic.ini Normal file
View File

@@ -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

1
packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/README Normal file
View File

@@ -0,0 +1 @@
Generic single-database configuration.

74
packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/env.py Normal file
View File

@@ -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()

26
packages/dbgpt-app/src/dbgpt_app/pilot_template/meta_data/alembic/script.py.mako Normal file
View File

@@ -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"}

47
packages/dbgpt-app/src/dbgpt_app/tests/test_workspace_provisioning.py Normal file
View File

@@ -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")

322
packages/dbgpt-core/src/dbgpt/cli/_config.py Normal file
View File

@@ -0,0 +1,322 @@
"""User-level configuration management for DB-GPT CLI.
Manages ``~/.dbgpt/configs/<profile>.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/<profile>.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

95
packages/dbgpt-core/src/dbgpt/cli/_profile_cmd.py Normal file
View File

@@ -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}")

203
packages/dbgpt-core/src/dbgpt/cli/_profiles.py Normal file
View File

@@ -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]

325
packages/dbgpt-core/src/dbgpt/cli/_wizard.py Normal file
View File

@@ -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

122
packages/dbgpt-core/src/dbgpt/cli/cli_scripts.py
View File

@@ -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/<profile>.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)

0
packages/dbgpt-core/src/dbgpt/cli/tests/__init__.py Normal file
View File

24
packages/dbgpt-core/src/dbgpt/cli/tests/conftest.py Normal file
View File

@@ -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)

65
packages/dbgpt-core/src/dbgpt/cli/tests/test_cli_scripts.py Normal file
View File

@@ -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

381
packages/dbgpt-core/src/dbgpt/cli/tests/test_config.py Normal file
View File

@@ -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

170
packages/dbgpt-core/src/dbgpt/cli/tests/test_integration.py Normal file
View File

@@ -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"

126
packages/dbgpt-core/src/dbgpt/cli/tests/test_profile_cmd.py Normal file
View File

@@ -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

175
packages/dbgpt-core/src/dbgpt/cli/tests/test_profiles.py Normal file
View File

@@ -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

43
packages/dbgpt-core/src/dbgpt/cli/tests/test_smoke.py Normal file
View File

@@ -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)

251
packages/dbgpt-core/src/dbgpt/cli/tests/test_wizard.py Normal file
View File

@@ -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()

29
packages/dbgpt-core/src/dbgpt/configs/model_config.py
View File

@@ -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"))

0
packages/dbgpt-core/src/dbgpt/configs/tests/__init__.py Normal file
View File

134
packages/dbgpt-core/src/dbgpt/configs/tests/test_model_config.py Normal file
View File

@@ -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

4
packages/dbgpt-core/src/dbgpt/util/_db_migration_utils.py
View File

@@ -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

158
packages/dbgpt-core/src/dbgpt/util/console/console.py
View File

@@ -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)}.[/]"
)

30
packages/dbgpt-core/src/dbgpt/util/tests/test_db_migration_utils.py Normal file
View File

@@ -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