From 3c9abb386edd12c8d758464bb9e82f7e86af33a1 Mon Sep 17 00:00:00 2001 From: Lenar Sharipov Date: Mon, 23 Mar 2026 18:01:55 +0100 Subject: [PATCH 1/8] DBA-291 Refactor databao-cli: Phase 1 --- docs/architecture.md | 85 +++-- src/databao_cli/__main__.py | 320 +----------------- src/databao_cli/commands/_utils.py | 17 + src/databao_cli/commands/app.py | 70 +++- src/databao_cli/commands/ask.py | 56 ++- src/databao_cli/commands/build.py | 30 ++ .../commands/datasource/__init__.py | 49 +++ src/databao_cli/commands/index.py | 29 ++ src/databao_cli/commands/init.py | 38 ++- src/databao_cli/commands/mcp.py | 45 +++ src/databao_cli/commands/status.py | 12 +- 11 files changed, 389 insertions(+), 362 deletions(-) create mode 100644 src/databao_cli/commands/_utils.py diff --git a/docs/architecture.md b/docs/architecture.md index 0a7a8686..41ee9c2a 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,66 +1,63 @@ # Architecture ## Overview - Databao CLI is a Python CLI application built with Click that provides a command-line and web interface to the Databao Agent and Context Engine. -Key layers: +Project structure: +``` +src/databao_cli/ +- commands/ # CLI commands +- ui/ # Web UI (Streamlit) +- mcp/ # MCP server +- project/ # Project management +- log/ # Logging +``` -- CLI commands (`src/databao_cli/commands/`) -- Web UI via Streamlit (`src/databao_cli/ui/`) -- MCP server (`src/databao_cli/mcp/`) -- Project management (`src/databao_cli/project/`) -- Logging (`src/databao_cli/log/`) +## Key Dependencies +- `databao-context-engine` — core context indexing and querying +- `databao-agent` — AI agent for natural language interaction +- `click` — CLI framework +- `streamlit` — web UI framework +- `fastmcp` — MCP server framework ## Entry Point - `databao_cli.__main__:cli` — a Click group that registers all subcommands. +Invoked as `uv run databao`. ## CLI Commands +The CLI is implemented using the Click framework. -| Command | Module | Purpose | -|----------------------|---------------------------------|---------------------------------| -| `databao init` | `commands/init.py` | Initialize a new project | -| `databao status` | `commands/status.py` | Show project status | -| `databao datasource` | `commands/datasource/` | Add/check data sources | -| `databao build` | `commands/build.py` | Build context for datasources | -| `databao ask` | `commands/ask.py` | Interactive chat with agent | -| `databao app` | `commands/app.py` | Launch Streamlit web UI | -| `databao mcp` | `commands/mcp.py` | Run MCP server | - -## Web UI (Streamlit) +### Structure +- All CLI commands are located in `src/databao_cli/commands/` +- Single command = single file. +- Each command MUST be implemented in its own module. +- File name SHOULD match the command name. +- Grouped commands MUST be organized in subdirectories. +- Example: `src/databao_cli/commands/datasource/` -`src/databao_cli/ui/` contains the Streamlit application: +### Registration +- All commands MUST be explicitly registered in `src/databao_cli/__main__.py` +- Commands are exposed via the COMMANDS collection. Every new command MUST be added to this collection to be discoverable by the CLI. -- `app.py` — main Streamlit entry -- `pages/` — individual UI pages -- `components/` — reusable UI components -- `services/` — backend service wrappers -- `models/` — data models for UI state +## Web UI (Streamlit) +``` +src/databao_cli/ui/ # contains the Streamlit application + app.py # main Streamlit entry + pages/ # individual UI pages + components/ # reusable UI components + services/ # backend service wrappers + models/ # data models for UI state +``` ## MCP Server - -`src/databao_cli/mcp/` exposes tools via the Model Context Protocol: - -- `server.py` — FastMCP server setup -- `tools/` — individual tool handlers - -## Key Dependencies - -- `databao-context-engine` — core context indexing and querying -- `databao-agent` — AI agent for natural language interaction -- `click` — CLI framework -- `streamlit` — web UI framework -- `fastmcp` — MCP server framework - -## Database Support - -Multi-database via optional dependencies: Snowflake (default), BigQuery, -ClickHouse, Athena, MS SQL, with ADBC drivers. +``` +src/databao_cli/mcp/ # exposes tools via the Model Context Protocol: + server.py # FastMCP server setup + tools/ # individual tool handlers +``` ## Extension Points - - Add CLI command: create module in `commands/`, register with Click group in `__main__.py`. - Add MCP tool: add handler in `mcp/tools/`, register in `mcp/server.py`. diff --git a/src/databao_cli/__main__.py b/src/databao_cli/__main__.py index fe6a92f7..a45d0c5d 100644 --- a/src/databao_cli/__main__.py +++ b/src/databao_cli/__main__.py @@ -1,11 +1,20 @@ -import sys from pathlib import Path import click -from click import Context - +from click import Command, Context + +from databao_cli.commands.app import app +from databao_cli.commands.ask import ask +from databao_cli.commands.build import build +from databao_cli.commands.datasource import datasource +from databao_cli.commands.index import index +from databao_cli.commands.init import init +from databao_cli.commands.mcp import mcp +from databao_cli.commands.status import status from databao_cli.log.logging import configure_logging -from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project +from databao_cli.project.layout import find_project + +COMMANDS: list[Command] = [app, ask, build, datasource, index, init, mcp, status] @click.group() @@ -27,307 +36,8 @@ def cli(ctx: Context, verbose: bool, project_dir: Path | None) -> None: configure_logging(find_project(project_path), verbose=verbose) -@cli.command() -@click.pass_context -def status(ctx: Context) -> None: - """Display project status and system-wide information.""" - from databao_cli.commands.status import status_impl - - status_message = status_impl(ctx.obj["project_dir"]) - click.echo(status_message) - - -@cli.command() -@click.pass_context -def init(ctx: Context) -> None: - """Create a new Databao project.""" - from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl - from databao_cli.commands.init import InitDatabaoProjectError, ProjectDirDoesnotExistError, init_impl - - project_dir = ctx.obj["project_dir"] - project_layout: ProjectLayout - try: - project_layout = init_impl(project_dir) - except ProjectDirDoesnotExistError: - if click.confirm( - f"The directory {project_dir.resolve()} does not exist. Do you want to create it?", - default=True, - ): - project_dir.mkdir(parents=True, exist_ok=False) - project_layout = init_impl(project_dir) - else: - return - except InitDatabaoProjectError as e: - click.echo(e.message, err=True) - sys.exit(1) - - click.echo(f"Project initialized successfully at {project_dir.resolve()}") - - # todo install ollama - # try: - # install_ollama_if_needed() - # except RuntimeError as e: - # click.echo(str(e), err=True) - - if not click.confirm("\nDo you want to configure a domain now?"): - return - - add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN) - - while click.confirm("\nDo you want to add more datasources?"): - add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN) - - -@cli.group() -def datasource() -> None: - """Manage datasource configurations.""" - pass - - -@datasource.command(name="add") -@click.option( - "-d", - "--domain", - type=click.STRING, - default="root", - help="Databao domain name", -) -@click.pass_context -def add_datasource_config(ctx: Context, domain: str) -> None: - """Add a new datasource configuration. - - The command will ask all relevant information for that datasource and save it in a chosen Databao domain - """ - from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl - - project_layout = _get_project_or_exit(ctx.obj["project_dir"]) - add_datasource_config_interactive_impl(project_layout, domain) - - -@datasource.command(name="check") -@click.argument( - "domains", - type=click.STRING, - nargs=-1, -) -@click.pass_context -def check_datasource_config(ctx: Context, domains: list[str] | None) -> None: - """Check whether a datasource configuration is valid. - - The configuration is considered as valid if a connection with the datasource can be established. - - By default, all declared datasources across all domains in the project will be checked. - You can explicitely list which domains to validate by using the [DOMAINS] argument. - """ - from databao_cli.commands.datasource.check_datasource_connection import check_datasource_connection_impl - - project_layout = _get_project_or_exit(ctx.obj["project_dir"]) - check_datasource_connection_impl(project_layout, requested_domains=domains if domains else None) - - -@cli.command() -@click.option( - "-d", - "--domain", - type=click.STRING, - default="root", - help="Databao domain name", -) -@click.option( - "--should-index/--should-not-index", - default=True, - show_default=True, - help="Whether to index the context. If disabled, the context will be built but not indexed.", -) -@click.pass_context -def build(ctx: Context, domain: str, should_index: bool) -> None: - """Build context for all domain's datasources. - - The output of the build command will be saved in the domain's output directory. - - Internally, this indexes the context to be used by the MCP server and the "retrieve" command. - """ - from databao_cli.commands.build import build_impl - - project_layout = _get_project_or_exit(ctx.obj["project_dir"]) - results = build_impl(project_layout, domain, should_index) - click.echo(f"Build complete. Processed {len(results)} datasources.") - - -@cli.command() -@click.option( - "-d", - "--domain", - type=click.STRING, - default="root", - help="Databao domain name", -) -@click.argument( - "datasources-config-files", - nargs=-1, - type=click.STRING, -) -@click.pass_context -def index(ctx: Context, domain: str, datasources_config_files: tuple[str, ...]) -> None: - """Index built contexts into the embeddings database. - - If one or more datasource config file strings are provided, only those datasources will be indexed. - If no values are provided, all built contexts for the selected domain will be indexed. - """ - from databao_cli.commands.index import index_impl - - project_layout = _get_project_or_exit(ctx.obj["project_dir"]) - datasources = list(datasources_config_files) if datasources_config_files else None - results = index_impl(project_layout, domain, datasources) - click.echo(f"Index complete. Processed {len(results)} contexts.") - - -@cli.command() -@click.argument("question", required=False) -@click.option( - "--one-shot", - is_flag=True, - default=False, - help="Run single question and exit (default: interactive mode)", -) -@click.option( - "-m", - "--model", - type=str, - default=None, - help="LLM model in format provider:name (e.g., openai:gpt-4o, anthropic:claude-sonnet-4-6)", -) -@click.option( - "-t", - "--temperature", - type=float, - default=0.0, - help="Temperature 0.0-1.0 (default: 0.0)", -) -@click.option( - "--show-thinking/--no-show-thinking", - default=True, - help="Display reasoning/thinking output (streaming is implicit when enabled)", -) -@click.pass_context -def ask( - ctx: click.Context, - question: str | None, - one_shot: bool, - model: str | None, - temperature: float, - show_thinking: bool, -) -> None: - """Chat with the Databao agent. - - By default, starts an interactive chat session. Use --one-shot with a - QUESTION argument to run a single query and exit. - - \b - Examples: - databao ask # Interactive mode - databao ask --one-shot "What tables exist?" # One-shot mode - databao ask --model anthropic:claude-sonnet-4-6 # With custom model - databao ask --no-show-thinking # Hide reasoning - """ - from databao_cli.commands.ask import ask_impl - - ask_impl(ctx, question, one_shot, model, temperature, show_thinking) - - -@cli.command( - context_settings={"ignore_unknown_options": True, "allow_extra_args": True}, -) -@click.option( - "--read-only-domain", - is_flag=True, - default=False, - help="Disable all domain-editing operations (init, datasources, build) in the UI", -) -@click.option( - "--hide-suggested-questions", - is_flag=True, - default=False, - help="Hide the suggested questions on the empty chat screen", -) -@click.option( - "--hide-build-context-hint", - is_flag=True, - default=False, - help=( - "Hide the 'Context isn't built yet' warning on the empty chat screen and " - "remove the Build Context step from the setup wizard" - ), -) -@click.pass_context -def app(ctx: click.Context, read_only_domain: bool, hide_suggested_questions: bool, hide_build_context_hint: bool) -> None: - """Launch the Databao Streamlit web interface. - - All additional arguments are passed directly to streamlit run. - - \b - Examples: - databao app - databao app --server.port 8502 - databao app --server.headless true - databao app --read-only-domain - """ - from databao_cli.commands.app import app_impl - - ctx.obj["read_only_domain"] = read_only_domain - ctx.obj["hide_suggested_questions"] = hide_suggested_questions - ctx.obj["hide_build_context_hint"] = hide_build_context_hint - app_impl(ctx) - - -@cli.command() -@click.option( - "--transport", - type=click.Choice(["stdio", "sse"]), - default="stdio", - show_default=True, - help="MCP transport type.", -) -@click.option( - "--host", - type=str, - default="localhost", - show_default=True, - help="Host to bind to (SSE transport only).", -) -@click.option( - "--port", - type=int, - default=8765, - show_default=True, - help="Port to listen on (SSE transport only).", -) -@click.pass_context -def mcp(ctx: click.Context, transport: str, host: str, port: int) -> None: - """Run an MCP server exposing Databao tools. - - Starts a Model Context Protocol server that provides tools for - interacting with your Databao project programmatically. - - \b - Examples: - databao mcp # stdio (default) - databao mcp --transport sse # SSE on localhost:8765 - databao mcp --transport sse --port 9000 # SSE on custom port - """ - from databao_cli.commands.mcp import mcp_impl - - if transport == "stdio": - configure_logging(find_project(ctx.obj["project_dir"]), quiet=True) - mcp_impl(ctx.obj["project_dir"], transport, host, port) - - -def _get_project_or_exit(project_dir: Path) -> ProjectLayout: - project_layout = find_project(project_dir) - if not project_layout: - click.echo("No project found.") - exit(1) - return project_layout +for command in COMMANDS: + cli.add_command(command) if __name__ == "__main__": diff --git a/src/databao_cli/commands/_utils.py b/src/databao_cli/commands/_utils.py new file mode 100644 index 00000000..dd76260b --- /dev/null +++ b/src/databao_cli/commands/_utils.py @@ -0,0 +1,17 @@ +"""Shared CLI utilities for command implementations.""" + +import sys +from pathlib import Path + +import click + +from databao_cli.project.layout import ProjectLayout, find_project + + +def get_project_or_exit(project_dir: Path) -> ProjectLayout: + """Return the project layout or exit with an error if no project is found.""" + project_layout = find_project(project_dir) + if not project_layout: + click.echo("No project found.") + sys.exit(1) + return project_layout diff --git a/src/databao_cli/commands/app.py b/src/databao_cli/commands/app.py index 00198fb8..da22f211 100644 --- a/src/databao_cli/commands/app.py +++ b/src/databao_cli/commands/app.py @@ -2,22 +2,80 @@ import subprocess import sys +from pathlib import Path import click from databao_cli.ui.cli import bootstrap_streamlit_app -def app_impl(ctx: click.Context) -> None: +@click.command( + context_settings={"ignore_unknown_options": True, "allow_extra_args": True}, +) +@click.option( + "--read-only-domain", + is_flag=True, + default=False, + help="Disable all domain-editing operations (init, datasources, build) in the UI", +) +@click.option( + "--hide-suggested-questions", + is_flag=True, + default=False, + help="Hide the suggested questions on the empty chat screen", +) +@click.option( + "--hide-build-context-hint", + is_flag=True, + default=False, + help=( + "Hide the 'Context isn't built yet' warning on the empty chat screen and " + "remove the Build Context step from the setup wizard" + ), +) +@click.pass_context +def app( + ctx: click.Context, + read_only_domain: bool, + hide_suggested_questions: bool, + hide_build_context_hint: bool, +) -> None: + """Launch the Databao Streamlit web interface. + + All additional arguments are passed directly to streamlit run. + + \b + Examples: + databao app + databao app --server.port 8502 + databao app --server.headless true + databao app --read-only-domain + """ + app_impl( + ctx.obj["project_dir"], + ctx.args, + read_only_domain=read_only_domain, + hide_suggested_questions=hide_suggested_questions, + hide_build_context_hint=hide_build_context_hint, + ) + + +def app_impl( + project_dir: Path, + extra_args: list[str], + read_only_domain: bool = False, + hide_suggested_questions: bool = False, + hide_build_context_hint: bool = False, +) -> None: click.echo("Starting Databao UI...") try: bootstrap_streamlit_app( - ctx.obj["project_dir"], - ctx.args, - read_only_domain=ctx.obj.get("read_only_domain", False), - hide_suggested_questions=ctx.obj.get("hide_suggested_questions", False), - hide_build_context_hint=ctx.obj.get("hide_build_context_hint", False), + project_dir, + extra_args, + read_only_domain=read_only_domain, + hide_suggested_questions=hide_suggested_questions, + hide_build_context_hint=hide_build_context_hint, ) except subprocess.CalledProcessError as e: click.echo(f"Error running Streamlit: {e}", err=True) diff --git a/src/databao_cli/commands/ask.py b/src/databao_cli/commands/ask.py index e72404c8..80494324 100644 --- a/src/databao_cli/commands/ask.py +++ b/src/databao_cli/commands/ask.py @@ -21,6 +21,57 @@ DEFAULT_MAX_DISPLAY_ROWS = 10 +@click.command() +@click.argument("question", required=False) +@click.option( + "--one-shot", + is_flag=True, + default=False, + help="Run single question and exit (default: interactive mode)", +) +@click.option( + "-m", + "--model", + type=str, + default=None, + help="LLM model in format provider:name (e.g., openai:gpt-4o, anthropic:claude-sonnet-4-6)", +) +@click.option( + "-t", + "--temperature", + type=float, + default=0.0, + help="Temperature 0.0-1.0 (default: 0.0)", +) +@click.option( + "--show-thinking/--no-show-thinking", + default=True, + help="Display reasoning/thinking output (streaming is implicit when enabled)", +) +@click.pass_context +def ask( + ctx: click.Context, + question: str | None, + one_shot: bool, + model: str | None, + temperature: float, + show_thinking: bool, +) -> None: + """Chat with the Databao agent. + + By default, starts an interactive chat session. Use --one-shot with a + QUESTION argument to run a single query and exit. + + \b + Examples: + databao ask # Interactive mode + databao ask --one-shot "What tables exist?" # One-shot mode + databao ask --model anthropic:claude-sonnet-4-6 # With custom model + databao ask --no-show-thinking # Hide reasoning + """ + ask_impl(ctx.obj["project_dir"], question, one_shot, model, temperature, show_thinking) + + def _create_cli_writer() -> StreamingWriter: """Create a StreamingWriter that echoes output to the CLI in real-time.""" return StreamingWriter(on_write=lambda text: click.echo(text, nl=False)) @@ -201,7 +252,7 @@ def run_one_shot_mode(agent: Agent, question: str, show_thinking: bool) -> None: def ask_impl( - ctx: click.Context, + project_path: Path, question: str | None, one_shot: bool, model: str | None, @@ -213,9 +264,6 @@ def ask_impl( click.echo("Error: QUESTION argument is required in one-shot mode.", err=True) sys.exit(1) - # Get project path from CLI context - project_path: Path = ctx.obj["project_dir"] - # Initialize agent (with progress indicator if not showing thinking) if not show_thinking: click.echo("Initializing agent...") diff --git a/src/databao_cli/commands/build.py b/src/databao_cli/commands/build.py index bf4b00df..1a88c3e9 100644 --- a/src/databao_cli/commands/build.py +++ b/src/databao_cli/commands/build.py @@ -1,8 +1,38 @@ +import click from databao_context_engine import BuildDatasourceResult, DatabaoContextDomainManager from databao_cli.project.layout import ProjectLayout +@click.command() +@click.option( + "-d", + "--domain", + type=click.STRING, + default="root", + help="Databao domain name", +) +@click.option( + "--should-index/--should-not-index", + default=True, + show_default=True, + help="Whether to index the context. If disabled, the context will be built but not indexed.", +) +@click.pass_context +def build(ctx: click.Context, domain: str, should_index: bool) -> None: + """Build context for all domain's datasources. + + The output of the build command will be saved in the domain's output directory. + + Internally, this indexes the context to be used by the MCP server and the "retrieve" command. + """ + from databao_cli.commands._utils import get_project_or_exit + + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + results = build_impl(project_layout, domain, should_index) + click.echo(f"Build complete. Processed {len(results)} datasources.") + + def build_impl(project_layout: ProjectLayout, domain: str, should_index: bool) -> list[BuildDatasourceResult]: dce_project_dir = project_layout.domains_dir / domain results: list[BuildDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).build_context( diff --git a/src/databao_cli/commands/datasource/__init__.py b/src/databao_cli/commands/datasource/__init__.py index e69de29b..08f9a1fe 100644 --- a/src/databao_cli/commands/datasource/__init__.py +++ b/src/databao_cli/commands/datasource/__init__.py @@ -0,0 +1,49 @@ +import click + + +@click.group() +def datasource() -> None: + """Manage datasource configurations.""" + + +@datasource.command(name="add") +@click.option( + "-d", + "--domain", + type=click.STRING, + default="root", + help="Databao domain name", +) +@click.pass_context +def add_datasource_config(ctx: click.Context, domain: str) -> None: + """Add a new datasource configuration. + + The command will ask all relevant information for that datasource and save it in a chosen Databao domain + """ + from databao_cli.commands._utils import get_project_or_exit + from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl + + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + add_datasource_config_interactive_impl(project_layout, domain) + + +@datasource.command(name="check") +@click.argument( + "domains", + type=click.STRING, + nargs=-1, +) +@click.pass_context +def check_datasource_config(ctx: click.Context, domains: tuple[str, ...]) -> None: + """Check whether a datasource configuration is valid. + + The configuration is considered as valid if a connection with the datasource can be established. + + By default, all declared datasources across all domains in the project will be checked. + You can explicitely list which domains to validate by using the [DOMAINS] argument. + """ + from databao_cli.commands._utils import get_project_or_exit + from databao_cli.commands.datasource.check_datasource_connection import check_datasource_connection_impl + + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + check_datasource_connection_impl(project_layout, requested_domains=list(domains) if domains else None) diff --git a/src/databao_cli/commands/index.py b/src/databao_cli/commands/index.py index 0e8c4406..445e9135 100644 --- a/src/databao_cli/commands/index.py +++ b/src/databao_cli/commands/index.py @@ -1,8 +1,37 @@ +import click from databao_context_engine import ChunkEmbeddingMode, DatabaoContextDomainManager, DatasourceId, IndexDatasourceResult from databao_cli.project.layout import ProjectLayout +@click.command() +@click.option( + "-d", + "--domain", + type=click.STRING, + default="root", + help="Databao domain name", +) +@click.argument( + "datasources-config-files", + nargs=-1, + type=click.STRING, +) +@click.pass_context +def index(ctx: click.Context, domain: str, datasources_config_files: tuple[str, ...]) -> None: + """Index built contexts into the embeddings database. + + If one or more datasource config file strings are provided, only those datasources will be indexed. + If no values are provided, all built contexts for the selected domain will be indexed. + """ + from databao_cli.commands._utils import get_project_or_exit + + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + datasources = list(datasources_config_files) if datasources_config_files else None + results = index_impl(project_layout, domain, datasources) + click.echo(f"Index complete. Processed {len(results)} contexts.") + + def index_impl( project_layout: ProjectLayout, domain: str, datasources_config_files: list[str] | None ) -> list[IndexDatasourceResult]: diff --git a/src/databao_cli/commands/init.py b/src/databao_cli/commands/init.py index bbcb795c..a85f7ba8 100644 --- a/src/databao_cli/commands/init.py +++ b/src/databao_cli/commands/init.py @@ -1,8 +1,44 @@ +import sys from pathlib import Path +import click from databao_context_engine import InitDomainError, init_dce_domain -from databao_cli.project.layout import ProjectLayout, find_project +from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project + + +@click.command() +@click.pass_context +def init(ctx: click.Context) -> None: + """Create a new Databao project.""" + from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl + + project_dir: Path = ctx.obj["project_dir"] + project_layout: ProjectLayout + try: + project_layout = init_impl(project_dir) + except ProjectDirDoesnotExistError: + if click.confirm( + f"The directory {project_dir.resolve()} does not exist. Do you want to create it?", + default=True, + ): + project_dir.mkdir(parents=True, exist_ok=False) + project_layout = init_impl(project_dir) + else: + return + except InitDatabaoProjectError as e: + click.echo(e.message, err=True) + sys.exit(1) + + click.echo(f"Project initialized successfully at {project_dir.resolve()}") + + if not click.confirm("\nDo you want to configure a domain now?"): + return + + add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN) + + while click.confirm("\nDo you want to add more datasources?"): + add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN) class InitDatabaoProjectError(ValueError): diff --git a/src/databao_cli/commands/mcp.py b/src/databao_cli/commands/mcp.py index 614a7d28..880e644a 100644 --- a/src/databao_cli/commands/mcp.py +++ b/src/databao_cli/commands/mcp.py @@ -2,9 +2,54 @@ from pathlib import Path +import click + from databao_cli.mcp.server import McpContext, run_server +@click.command() +@click.option( + "--transport", + type=click.Choice(["stdio", "sse"]), + default="stdio", + show_default=True, + help="MCP transport type.", +) +@click.option( + "--host", + type=str, + default="localhost", + show_default=True, + help="Host to bind to (SSE transport only).", +) +@click.option( + "--port", + type=int, + default=8765, + show_default=True, + help="Port to listen on (SSE transport only).", +) +@click.pass_context +def mcp(ctx: click.Context, transport: str, host: str, port: int) -> None: + """Run an MCP server exposing Databao tools. + + Starts a Model Context Protocol server that provides tools for + interacting with your Databao project programmatically. + + \b + Examples: + databao mcp # stdio (default) + databao mcp --transport sse # SSE on localhost:8765 + databao mcp --transport sse --port 9000 # SSE on custom port + """ + from databao_cli.log.logging import configure_logging + from databao_cli.project.layout import find_project + + if transport == "stdio": + configure_logging(find_project(ctx.obj["project_dir"]), quiet=True) + mcp_impl(ctx.obj["project_dir"], transport, host, port) + + def mcp_impl(project_dir: Path, transport: str, host: str, port: int) -> None: context = McpContext(project_dir=project_dir) run_server(context, transport=transport, host=host, port=port) diff --git a/src/databao_cli/commands/status.py b/src/databao_cli/commands/status.py index 7737396c..551a4cb7 100644 --- a/src/databao_cli/commands/status.py +++ b/src/databao_cli/commands/status.py @@ -3,6 +3,7 @@ from importlib.metadata import version from pathlib import Path +import click from databao_context_engine import ( DceDomainInfo, DceInfo, @@ -10,11 +11,18 @@ get_databao_context_engine_info, ) -from databao_cli.project.layout import find_project +from databao_cli.commands._utils import get_project_or_exit + + +@click.command() +@click.pass_context +def status(ctx: click.Context) -> None: + """Display project status and system-wide information.""" + click.echo(status_impl(ctx.obj["project_dir"])) def status_impl(project_dir: Path) -> str: - project_layout = find_project(project_dir) + project_layout = get_project_or_exit(project_dir) dce_info = get_databao_context_engine_info() From d8312da08d406c0fde1f856fa6f679b6c72bcc5d Mon Sep 17 00:00:00 2001 From: Lenar Date: Mon, 23 Mar 2026 18:18:02 +0100 Subject: [PATCH 2/8] Update src/databao_cli/commands/datasource/__init__.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/databao_cli/commands/datasource/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databao_cli/commands/datasource/__init__.py b/src/databao_cli/commands/datasource/__init__.py index 08f9a1fe..95b6b6c1 100644 --- a/src/databao_cli/commands/datasource/__init__.py +++ b/src/databao_cli/commands/datasource/__init__.py @@ -40,7 +40,7 @@ def check_datasource_config(ctx: click.Context, domains: tuple[str, ...]) -> Non The configuration is considered as valid if a connection with the datasource can be established. By default, all declared datasources across all domains in the project will be checked. - You can explicitely list which domains to validate by using the [DOMAINS] argument. + You can explicitly list which domains to validate by using the [DOMAINS] argument. """ from databao_cli.commands._utils import get_project_or_exit from databao_cli.commands.datasource.check_datasource_connection import check_datasource_connection_impl From 28bd197d19591e028125992d256a7dc05c27f031 Mon Sep 17 00:00:00 2001 From: Lenar Sharipov Date: Mon, 23 Mar 2026 18:39:13 +0100 Subject: [PATCH 3/8] Refactor datasource commands group --- docs/architecture.md | 7 ++- .../commands/datasource/__init__.py | 46 ++----------------- .../{add_datasource_config.py => add.py} | 24 +++++++++- ...heck_datasource_connection.py => check.py} | 44 +++++++++++++----- src/databao_cli/commands/init.py | 2 +- 5 files changed, 63 insertions(+), 60 deletions(-) rename src/databao_cli/commands/datasource/{add_datasource_config.py => add.py} (76%) rename src/databao_cli/commands/datasource/{check_datasource_connection.py => check.py} (69%) diff --git a/docs/architecture.md b/docs/architecture.md index 41ee9c2a..4af6b467 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -33,12 +33,11 @@ The CLI is implemented using the Click framework. - Single command = single file. - Each command MUST be implemented in its own module. - File name SHOULD match the command name. -- Grouped commands MUST be organized in subdirectories. -- Example: `src/databao_cli/commands/datasource/` +- Grouped commands MUST be organized in subdirectories. Example: `src/databao_cli/commands/datasource/` ### Registration -- All commands MUST be explicitly registered in `src/databao_cli/__main__.py` -- Commands are exposed via the COMMANDS collection. Every new command MUST be added to this collection to be discoverable by the CLI. +- All commands and groups MUST be explicitly registered in `src/databao_cli/__main__.py` +- Commands and groups are exposed via the COMMANDS collection. Every new command or group MUST be added to this collection to be discoverable by the CLI. ## Web UI (Streamlit) ``` diff --git a/src/databao_cli/commands/datasource/__init__.py b/src/databao_cli/commands/datasource/__init__.py index 95b6b6c1..531ccc5f 100644 --- a/src/databao_cli/commands/datasource/__init__.py +++ b/src/databao_cli/commands/datasource/__init__.py @@ -1,49 +1,13 @@ import click +from databao_cli.commands.datasource.add import add +from databao_cli.commands.datasource.check import check + @click.group() def datasource() -> None: """Manage datasource configurations.""" -@datasource.command(name="add") -@click.option( - "-d", - "--domain", - type=click.STRING, - default="root", - help="Databao domain name", -) -@click.pass_context -def add_datasource_config(ctx: click.Context, domain: str) -> None: - """Add a new datasource configuration. - - The command will ask all relevant information for that datasource and save it in a chosen Databao domain - """ - from databao_cli.commands._utils import get_project_or_exit - from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl - - project_layout = get_project_or_exit(ctx.obj["project_dir"]) - add_datasource_config_interactive_impl(project_layout, domain) - - -@datasource.command(name="check") -@click.argument( - "domains", - type=click.STRING, - nargs=-1, -) -@click.pass_context -def check_datasource_config(ctx: click.Context, domains: tuple[str, ...]) -> None: - """Check whether a datasource configuration is valid. - - The configuration is considered as valid if a connection with the datasource can be established. - - By default, all declared datasources across all domains in the project will be checked. - You can explicitly list which domains to validate by using the [DOMAINS] argument. - """ - from databao_cli.commands._utils import get_project_or_exit - from databao_cli.commands.datasource.check_datasource_connection import check_datasource_connection_impl - - project_layout = get_project_or_exit(ctx.obj["project_dir"]) - check_datasource_connection_impl(project_layout, requested_domains=list(domains) if domains else None) +datasource.add_command(add) +datasource.add_command(check) diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add.py similarity index 76% rename from src/databao_cli/commands/datasource/add_datasource_config.py rename to src/databao_cli/commands/datasource/add.py index b1bb83e2..83f9c45d 100644 --- a/src/databao_cli/commands/datasource/add_datasource_config.py +++ b/src/databao_cli/commands/datasource/add.py @@ -7,12 +7,32 @@ DatasourceType, ) +from databao_cli.commands._utils import get_project_or_exit from databao_cli.commands.context_engine_cli import ClickUserInputCallback -from databao_cli.commands.datasource.check_datasource_connection import print_connection_check_results from databao_cli.project.layout import ProjectLayout -def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None: +@click.command(name="add") +@click.option( + "-d", + "--domain", + type=click.STRING, + default="root", + help="Databao domain name", +) +@click.pass_context +def add(ctx: click.Context, domain: str) -> None: + """Add a new datasource configuration. + + The command will ask all relevant information for that datasource and save it in a chosen Databao domain + """ + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + add_impl(project_layout, domain) + + +def add_impl(project_layout: ProjectLayout, domain: str) -> None: + from databao_cli.commands.datasource.check import print_connection_check_results + domain_dir = project_layout.domains_dir / domain domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) plugin_loader = DatabaoContextPluginLoader() diff --git a/src/databao_cli/commands/datasource/check_datasource_connection.py b/src/databao_cli/commands/datasource/check.py similarity index 69% rename from src/databao_cli/commands/datasource/check_datasource_connection.py rename to src/databao_cli/commands/datasource/check.py index 7aad5411..1d2f72f2 100644 --- a/src/databao_cli/commands/datasource/check_datasource_connection.py +++ b/src/databao_cli/commands/datasource/check.py @@ -7,24 +7,30 @@ DatasourceId, ) +from databao_cli.commands._utils import get_project_or_exit from databao_cli.project.layout import ProjectLayout -def print_connection_check_results( - domain: str, datasource_results: dict[DatasourceId, CheckDatasourceConnectionResult] -) -> None: - for result in datasource_results.values(): - fq_datasource_name = domain + os.pathsep + str(result.datasource_id) - status = str(result.connection_status.value) - if result.summary: - status += f" - {result.summary}" - if result.full_message: - status += f": {result.full_message}" +@click.command(name="check") +@click.argument( + "domains", + type=click.STRING, + nargs=-1, +) +@click.pass_context +def check(ctx: click.Context, domains: tuple[str, ...]) -> None: + """Check whether a datasource configuration is valid. - click.echo(f"{fq_datasource_name}: {status}") + The configuration is considered as valid if a connection with the datasource can be established. + + By default, all declared datasources across all domains in the project will be checked. + You can explicitly list which domains to validate by using the [DOMAINS] argument. + """ + project_layout = get_project_or_exit(ctx.obj["project_dir"]) + check_impl(project_layout, requested_domains=list(domains) if domains else None) -def check_datasource_connection_impl(project_layout: ProjectLayout, requested_domains: list[str] | None) -> None: +def check_impl(project_layout: ProjectLayout, requested_domains: list[str] | None) -> None: domains: list[str] = project_layout.get_domain_names() if requested_domains is None else requested_domains results = _check_domains(project_layout, domains) @@ -37,6 +43,20 @@ def check_datasource_connection_impl(project_layout: ProjectLayout, requested_do print_connection_check_results(domain, datasource_results) +def print_connection_check_results( + domain: str, datasource_results: dict[DatasourceId, CheckDatasourceConnectionResult] +) -> None: + for result in datasource_results.values(): + fq_datasource_name = domain + os.pathsep + str(result.datasource_id) + status = str(result.connection_status.value) + if result.summary: + status += f" - {result.summary}" + if result.full_message: + status += f": {result.full_message}" + + click.echo(f"{fq_datasource_name}: {status}") + + def _check_domains( project_layout: ProjectLayout, domains: list[str] ) -> dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]]: diff --git a/src/databao_cli/commands/init.py b/src/databao_cli/commands/init.py index a85f7ba8..d05d1a72 100644 --- a/src/databao_cli/commands/init.py +++ b/src/databao_cli/commands/init.py @@ -11,7 +11,7 @@ @click.pass_context def init(ctx: click.Context) -> None: """Create a new Databao project.""" - from databao_cli.commands.datasource.add_datasource_config import add_datasource_config_interactive_impl + from databao_cli.commands.datasource.add import add_impl as add_datasource_config_interactive_impl project_dir: Path = ctx.obj["project_dir"] project_layout: ProjectLayout From 31c145540f45fb4f3ad860351015ee672eb213b9 Mon Sep 17 00:00:00 2001 From: Lenar Date: Mon, 23 Mar 2026 20:44:08 +0100 Subject: [PATCH 4/8] Update docs/architecture.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- docs/architecture.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 4af6b467..1f5a7fcb 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -22,8 +22,9 @@ src/databao_cli/ - `fastmcp` — MCP server framework ## Entry Point -`databao_cli.__main__:cli` — a Click group that registers all subcommands. -Invoked as `uv run databao`. +`databao_cli.__main__:cli` — a Click group that registers all subcommands and +is exposed as the `databao` console script (run as `databao ...` after installation). +During development, it can also be invoked via `uv run databao`. ## CLI Commands The CLI is implemented using the Click framework. From a86079213e9f8fd78dbba5c14a4cdfe8c2fecb3d Mon Sep 17 00:00:00 2001 From: Lenar Date: Mon, 23 Mar 2026 20:50:04 +0100 Subject: [PATCH 5/8] Update docs/architecture.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- docs/architecture.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 1f5a7fcb..fc2876dd 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -58,8 +58,8 @@ src/databao_cli/mcp/ # exposes tools via the Model Context Protocol: ``` ## Extension Points -- Add CLI command: create module in `commands/`, register with Click group - in `__main__.py`. +- Add CLI command: create module in `commands/`, then add it to the + `COMMANDS` collection in `src/databao_cli/__main__.py`. - Add MCP tool: add handler in `mcp/tools/`, register in `mcp/server.py`. - Add datasource type: extend via `databao-context-engine` optional deps. - Add UI page: add to `ui/pages/`. From 8ada0e38c97000d9afdb8286e256f4df55527cc6 Mon Sep 17 00:00:00 2001 From: Lenar Sharipov Date: Mon, 23 Mar 2026 19:27:00 +0100 Subject: [PATCH 6/8] DBA-292 Refactor databao-cli: Phase 2 --- docs/architecture.md | 108 ++++++--- .../src/databao_snowflake_demo/app.py | 2 +- src/databao_cli/__main__.py | 4 +- src/databao_cli/commands/_utils.py | 17 -- src/databao_cli/commands/app.py | 33 +-- src/databao_cli/commands/ask.py | 219 +----------------- src/databao_cli/commands/build.py | 18 +- src/databao_cli/commands/datasource/add.py | 63 +---- src/databao_cli/commands/datasource/check.py | 61 +---- src/databao_cli/commands/index.py | 23 +- src/databao_cli/commands/init.py | 76 +----- src/databao_cli/commands/mcp.py | 15 +- src/databao_cli/commands/status.py | 49 +--- src/databao_cli/{log => features}/__init__.py | 0 .../{mcp => features/ask}/__init__.py | 0 src/databao_cli/features/ask/agent_factory.py | 41 ++++ src/databao_cli/features/ask/display.py | 13 ++ src/databao_cli/features/ask/service.py | 18 ++ src/databao_cli/features/build.py | 11 + .../tools => features/datasource}/__init__.py | 0 src/databao_cli/features/datasource/add.py | 37 +++ src/databao_cli/features/datasource/check.py | 30 +++ src/databao_cli/features/index.py | 17 ++ .../{project => features/init}/__init__.py | 0 src/databao_cli/features/init/errors.py | 29 +++ src/databao_cli/features/init/service.py | 47 ++++ src/databao_cli/features/mcp/__init__.py | 0 src/databao_cli/{ => features}/mcp/server.py | 7 +- .../features/mcp/tools/__init__.py | 0 .../mcp/tools/databao_ask/__init__.py | 5 + .../mcp/tools/databao_ask/agent_factory.py | 6 +- .../mcp/tools/databao_ask/tool.py | 4 +- src/databao_cli/features/status.py | 46 ++++ src/databao_cli/{ => features}/ui/__init__.py | 0 src/databao_cli/{ => features}/ui/app.py | 46 ++-- .../{ => features}/ui/assets/bao.png | Bin src/databao_cli/{ => features}/ui/cli.py | 25 +- .../{ => features}/ui/components/__init__.py | 0 .../{ => features}/ui/components/chat.py | 12 +- .../ui/components/datasource_form.py | 0 .../ui/components/datasource_manager.py | 6 +- .../{ => features}/ui/components/icons.py | 0 .../{ => features}/ui/components/results.py | 6 +- .../{ => features}/ui/components/sidebar.py | 18 +- .../{ => features}/ui/components/status.py | 0 .../features/ui/models/__init__.py | 11 + .../{ => features}/ui/models/chat_session.py | 2 +- .../{ => features}/ui/models/settings.py | 0 .../{ => features}/ui/pages/__init__.py | 0 .../{ => features}/ui/pages/agent_settings.py | 14 +- .../{ => features}/ui/pages/chat.py | 20 +- .../ui/pages/context_settings.py | 14 +- .../ui/pages/general_settings.py | 8 +- .../{ => features}/ui/pages/welcome.py | 20 +- .../{ => features}/ui/project_utils.py | 2 +- .../{ => features}/ui/services/__init__.py | 10 +- .../ui/services/build_service.py | 2 +- .../ui/services/chat_persistence.py | 6 +- .../{ => features}/ui/services/chat_title.py | 2 +- .../ui/services/dce_operations.py | 6 +- .../{ => features}/ui/services/llm_models.py | 2 +- .../ui/services/query_executor.py | 6 +- .../ui/services/settings_persistence.py | 4 +- .../{ => features}/ui/services/storage.py | 0 .../{ => features}/ui/streaming.py | 0 .../{ => features}/ui/suggestions.py | 0 .../mcp/tools/databao_ask/__init__.py | 5 - src/databao_cli/shared/__init__.py | 0 src/databao_cli/shared/cli_utils.py | 34 +++ .../context_engine_cli.py | 0 src/databao_cli/shared/errors.py | 7 + .../{ => shared}/executor_utils.py | 0 src/databao_cli/shared/log/__init__.py | 0 .../{ => shared}/log/llm_errors.py | 0 src/databao_cli/{ => shared}/log/logging.py | 2 +- src/databao_cli/shared/project/__init__.py | 0 .../{ => shared}/project/layout.py | 0 src/databao_cli/ui/models/__init__.py | 11 - src/databao_cli/workflows/__init__.py | 0 src/databao_cli/workflows/ask.py | 119 ++++++++++ .../workflows/datasource/__init__.py | 0 src/databao_cli/workflows/datasource/add.py | 65 ++++++ src/databao_cli/workflows/datasource/check.py | 20 ++ tests/conftest.py | 4 +- tests/test_add_datasource.py | 2 +- tests/test_app.py | 4 +- tests/test_build.py | 2 +- tests/test_index.py | 2 +- tests/test_query_executor.py | 6 +- tests/test_query_executor_race_conditions.py | 8 +- 90 files changed, 835 insertions(+), 697 deletions(-) delete mode 100644 src/databao_cli/commands/_utils.py rename src/databao_cli/{log => features}/__init__.py (100%) rename src/databao_cli/{mcp => features/ask}/__init__.py (100%) create mode 100644 src/databao_cli/features/ask/agent_factory.py create mode 100644 src/databao_cli/features/ask/display.py create mode 100644 src/databao_cli/features/ask/service.py create mode 100644 src/databao_cli/features/build.py rename src/databao_cli/{mcp/tools => features/datasource}/__init__.py (100%) create mode 100644 src/databao_cli/features/datasource/add.py create mode 100644 src/databao_cli/features/datasource/check.py create mode 100644 src/databao_cli/features/index.py rename src/databao_cli/{project => features/init}/__init__.py (100%) create mode 100644 src/databao_cli/features/init/errors.py create mode 100644 src/databao_cli/features/init/service.py create mode 100644 src/databao_cli/features/mcp/__init__.py rename src/databao_cli/{ => features}/mcp/server.py (79%) create mode 100644 src/databao_cli/features/mcp/tools/__init__.py create mode 100644 src/databao_cli/features/mcp/tools/databao_ask/__init__.py rename src/databao_cli/{ => features}/mcp/tools/databao_ask/agent_factory.py (88%) rename src/databao_cli/{ => features}/mcp/tools/databao_ask/tool.py (97%) create mode 100644 src/databao_cli/features/status.py rename src/databao_cli/{ => features}/ui/__init__.py (100%) rename src/databao_cli/{ => features}/ui/app.py (90%) rename src/databao_cli/{ => features}/ui/assets/bao.png (100%) rename src/databao_cli/{ => features}/ui/cli.py (59%) rename src/databao_cli/{ => features}/ui/components/__init__.py (100%) rename src/databao_cli/{ => features}/ui/components/chat.py (97%) rename src/databao_cli/{ => features}/ui/components/datasource_form.py (100%) rename src/databao_cli/{ => features}/ui/components/datasource_manager.py (97%) rename src/databao_cli/{ => features}/ui/components/icons.py (100%) rename src/databao_cli/{ => features}/ui/components/results.py (98%) rename src/databao_cli/{ => features}/ui/components/sidebar.py (88%) rename src/databao_cli/{ => features}/ui/components/status.py (100%) create mode 100644 src/databao_cli/features/ui/models/__init__.py rename src/databao_cli/{ => features}/ui/models/chat_session.py (98%) rename src/databao_cli/{ => features}/ui/models/settings.py (100%) rename src/databao_cli/{ => features}/ui/pages/__init__.py (100%) rename src/databao_cli/{ => features}/ui/pages/agent_settings.py (95%) rename src/databao_cli/{ => features}/ui/pages/chat.py (90%) rename src/databao_cli/{ => features}/ui/pages/context_settings.py (89%) rename src/databao_cli/{ => features}/ui/pages/general_settings.py (91%) rename src/databao_cli/{ => features}/ui/pages/welcome.py (93%) rename src/databao_cli/{ => features}/ui/project_utils.py (96%) rename src/databao_cli/{ => features}/ui/services/__init__.py (78%) rename src/databao_cli/{ => features}/ui/services/build_service.py (99%) rename src/databao_cli/{ => features}/ui/services/chat_persistence.py (97%) rename src/databao_cli/{ => features}/ui/services/chat_title.py (98%) rename src/databao_cli/{ => features}/ui/services/dce_operations.py (97%) rename src/databao_cli/{ => features}/ui/services/llm_models.py (97%) rename src/databao_cli/{ => features}/ui/services/query_executor.py (97%) rename src/databao_cli/{ => features}/ui/services/settings_persistence.py (93%) rename src/databao_cli/{ => features}/ui/services/storage.py (100%) rename src/databao_cli/{ => features}/ui/streaming.py (100%) rename src/databao_cli/{ => features}/ui/suggestions.py (100%) delete mode 100644 src/databao_cli/mcp/tools/databao_ask/__init__.py create mode 100644 src/databao_cli/shared/__init__.py create mode 100644 src/databao_cli/shared/cli_utils.py rename src/databao_cli/{commands => shared}/context_engine_cli.py (100%) create mode 100644 src/databao_cli/shared/errors.py rename src/databao_cli/{ => shared}/executor_utils.py (100%) create mode 100644 src/databao_cli/shared/log/__init__.py rename src/databao_cli/{ => shared}/log/llm_errors.py (100%) rename src/databao_cli/{ => shared}/log/logging.py (97%) create mode 100644 src/databao_cli/shared/project/__init__.py rename src/databao_cli/{ => shared}/project/layout.py (100%) delete mode 100644 src/databao_cli/ui/models/__init__.py create mode 100644 src/databao_cli/workflows/__init__.py create mode 100644 src/databao_cli/workflows/ask.py create mode 100644 src/databao_cli/workflows/datasource/__init__.py create mode 100644 src/databao_cli/workflows/datasource/add.py create mode 100644 src/databao_cli/workflows/datasource/check.py diff --git a/docs/architecture.md b/docs/architecture.md index fc2876dd..3c98726d 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -7,11 +7,10 @@ command-line and web interface to the Databao Agent and Context Engine. Project structure: ``` src/databao_cli/ -- commands/ # CLI commands -- ui/ # Web UI (Streamlit) -- mcp/ # MCP server -- project/ # Project management -- log/ # Logging +- commands/ # CLI routing layer (Click wrappers only) +- workflows/ # Interactive CLI sequences (REPL, wizards, result display) +- features/ # Business logic, organized by feature — no Click dependency +- shared/ # Cross-feature utilities and models ``` ## Key Dependencies @@ -26,40 +25,91 @@ src/databao_cli/ is exposed as the `databao` console script (run as `databao ...` after installation). During development, it can also be invoked via `uv run databao`. -## CLI Commands -The CLI is implemented using the Click framework. +## Layer Responsibilities -### Structure -- All CLI commands are located in `src/databao_cli/commands/` -- Single command = single file. -- Each command MUST be implemented in its own module. -- File name SHOULD match the command name. -- Grouped commands MUST be organized in subdirectories. Example: `src/databao_cli/commands/datasource/` +### `commands/` — Routing Layer +Contains only Click wiring: decorators, option/argument definitions, and calls +to workflow or feature functions. No business logic lives here. -### Registration -- All commands and groups MUST be explicitly registered in `src/databao_cli/__main__.py` -- Commands and groups are exposed via the COMMANDS collection. Every new command or group MUST be added to this collection to be discoverable by the CLI. +- Single command = single file; file name MUST match the command name. +- Grouped commands MUST be organized in subdirectories (e.g. `commands/datasource/`). +- All commands and groups MUST be registered in `src/databao_cli/__main__.py` + via the `COMMANDS` collection. + +### `workflows/` — Interactive CLI Layer +Interactive CLI sequences that need multi-step user interaction (prompts, +confirmations, REPL loops) or result display. Uses Click freely. Calls into +`features/` for business operations. Complex workflows get a subdirectory; +thin workflows are a single file. -## Web UI (Streamlit) ``` -src/databao_cli/ui/ # contains the Streamlit application - app.py # main Streamlit entry - pages/ # individual UI pages - components/ # reusable UI components - services/ # backend service wrappers - models/ # data models for UI state +workflows/ + ask.py # run_interactive_mode, run_one_shot_mode, display_result + datasource/ + add.py # add_workflow (interactive wizard: prompts, confirms, echoes) + check.py # print_connection_check_results ``` -## MCP Server +Commands that require no interaction (build, index, status) skip this layer and +call `features/` directly. + +### `features/` — Business Logic Layer +All application logic lives here, organized by feature. No Click dependency — +`features/` functions are pure business operations callable from any context +(CLI, MCP, tests). Complex features get a subdirectory; thin features are a +single file. + +``` +features/ + ask/ + agent_factory.py # initialize_agent_from_dce — returns Agent + display.py # dataframe_to_prettytable (pure formatter, returns str) + service.py # ask_impl — validates args, initialises agent + init/ + errors.py # project initialization error types + service.py # init_impl, _ProjectCreator + datasource/ + add.py # datasource_config_exists, create_datasource_config + check.py # check_impl — returns connection results dict + mcp/ + server.py # FastMCP server setup, mcp_impl + tools/ # individual MCP tool handlers + ui/ + app.py # main Streamlit entry point + cli.py # bootstrap_streamlit_app, app_impl + pages/ # individual UI pages + components/ # reusable UI components + services/ # backend service wrappers + models/ # data models for UI state + streaming.py # streaming writer for real-time output + project_utils.py # project status helpers used across features + suggestions.py # suggested questions logic + build.py # build_impl + index.py # index_impl + status.py # status_impl, info string generation +``` + +### `shared/` — Shared Utilities +Cross-feature code with no business logic of its own. + ``` -src/databao_cli/mcp/ # exposes tools via the Model Context Protocol: - server.py # FastMCP server setup - tools/ # individual tool handlers +shared/ + project/ + layout.py # ProjectLayout dataclass, find_project + log/ + logging.py # configure_logging + llm_errors.py # format_llm_error (user-friendly API error messages) + cli_utils.py # get_project_or_raise, handle_feature_errors + context_engine_cli.py # ClickUserInputCallback (Click ↔ DCE adapter) + executor_utils.py # build_llm_config, LLM provider/model constants + errors.py # FeatureError ``` ## Extension Points - Add CLI command: create module in `commands/`, then add it to the `COMMANDS` collection in `src/databao_cli/__main__.py`. -- Add MCP tool: add handler in `mcp/tools/`, register in `mcp/server.py`. +- Add interactive workflow: create module in `workflows/`, call from command. +- Add MCP tool: add handler in `features/mcp/tools/`, register in + `features/mcp/server.py`. - Add datasource type: extend via `databao-context-engine` optional deps. -- Add UI page: add to `ui/pages/`. +- Add UI page: add to `features/ui/pages/`. diff --git a/examples/demo-snowflake-project/src/databao_snowflake_demo/app.py b/examples/demo-snowflake-project/src/databao_snowflake_demo/app.py index 08c8f5c0..4b600cc0 100644 --- a/examples/demo-snowflake-project/src/databao_snowflake_demo/app.py +++ b/examples/demo-snowflake-project/src/databao_snowflake_demo/app.py @@ -6,7 +6,7 @@ import streamlit as st -from databao_cli.ui.app import main +from databao_cli.features.ui.app import main logger = logging.getLogger(__name__) diff --git a/src/databao_cli/__main__.py b/src/databao_cli/__main__.py index a45d0c5d..55b7378e 100644 --- a/src/databao_cli/__main__.py +++ b/src/databao_cli/__main__.py @@ -11,8 +11,8 @@ from databao_cli.commands.init import init from databao_cli.commands.mcp import mcp from databao_cli.commands.status import status -from databao_cli.log.logging import configure_logging -from databao_cli.project.layout import find_project +from databao_cli.shared.log.logging import configure_logging +from databao_cli.shared.project.layout import find_project COMMANDS: list[Command] = [app, ask, build, datasource, index, init, mcp, status] diff --git a/src/databao_cli/commands/_utils.py b/src/databao_cli/commands/_utils.py deleted file mode 100644 index dd76260b..00000000 --- a/src/databao_cli/commands/_utils.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Shared CLI utilities for command implementations.""" - -import sys -from pathlib import Path - -import click - -from databao_cli.project.layout import ProjectLayout, find_project - - -def get_project_or_exit(project_dir: Path) -> ProjectLayout: - """Return the project layout or exit with an error if no project is found.""" - project_layout = find_project(project_dir) - if not project_layout: - click.echo("No project found.") - sys.exit(1) - return project_layout diff --git a/src/databao_cli/commands/app.py b/src/databao_cli/commands/app.py index da22f211..617ee8ba 100644 --- a/src/databao_cli/commands/app.py +++ b/src/databao_cli/commands/app.py @@ -1,12 +1,8 @@ """databao app command - Launch the Databao Streamlit web interface.""" -import subprocess -import sys -from pathlib import Path - import click -from databao_cli.ui.cli import bootstrap_streamlit_app +from databao_cli.shared.cli_utils import handle_feature_errors @click.command( @@ -34,6 +30,7 @@ ), ) @click.pass_context +@handle_feature_errors def app( ctx: click.Context, read_only_domain: bool, @@ -51,34 +48,16 @@ def app( databao app --server.headless true databao app --read-only-domain """ - app_impl( - ctx.obj["project_dir"], - ctx.args, - read_only_domain=read_only_domain, - hide_suggested_questions=hide_suggested_questions, - hide_build_context_hint=hide_build_context_hint, - ) - + from databao_cli.features.ui.cli import app_impl -def app_impl( - project_dir: Path, - extra_args: list[str], - read_only_domain: bool = False, - hide_suggested_questions: bool = False, - hide_build_context_hint: bool = False, -) -> None: click.echo("Starting Databao UI...") - try: - bootstrap_streamlit_app( - project_dir, - extra_args, + app_impl( + ctx.obj["project_dir"], + ctx.args, read_only_domain=read_only_domain, hide_suggested_questions=hide_suggested_questions, hide_build_context_hint=hide_build_context_hint, ) - except subprocess.CalledProcessError as e: - click.echo(f"Error running Streamlit: {e}", err=True) - sys.exit(1) except KeyboardInterrupt: click.echo("\nShutting down Databao...") diff --git a/src/databao_cli/commands/ask.py b/src/databao_cli/commands/ask.py index 80494324..1fb34ca4 100644 --- a/src/databao_cli/commands/ask.py +++ b/src/databao_cli/commands/ask.py @@ -1,24 +1,8 @@ """databao ask command - Interactive CLI chat with the Databao agent.""" -import sys -from pathlib import Path - import click -import pandas as pd -from databao.agent import Agent -from databao.agent import domain as create_domain -from databao.agent.api import agent as create_agent -from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory -from databao.agent.core.thread import Thread -from prettytable import PrettyTable - -from databao_cli.log.llm_errors import format_llm_error -from databao_cli.project.layout import ProjectLayout -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status -from databao_cli.ui.streaming import StreamingWriter -# Default maximum number of rows to display in dataframe output -DEFAULT_MAX_DISPLAY_ROWS = 10 +from databao_cli.shared.cli_utils import handle_feature_errors @click.command() @@ -49,6 +33,7 @@ help="Display reasoning/thinking output (streaming is implicit when enabled)", ) @click.pass_context +@handle_feature_errors def ask( ctx: click.Context, question: str | None, @@ -69,207 +54,13 @@ def ask( databao ask --model anthropic:claude-sonnet-4-6 # With custom model databao ask --no-show-thinking # Hide reasoning """ - ask_impl(ctx.obj["project_dir"], question, one_shot, model, temperature, show_thinking) - - -def _create_cli_writer() -> StreamingWriter: - """Create a StreamingWriter that echoes output to the CLI in real-time.""" - return StreamingWriter(on_write=lambda text: click.echo(text, nl=False)) - - -def dataframe_to_prettytable(df: pd.DataFrame, max_rows: int = DEFAULT_MAX_DISPLAY_ROWS) -> str: - """Convert a pandas DataFrame to a prettytable string.""" - table = PrettyTable() - table.field_names = list(df.columns) - for _, row in df.head(max_rows).iterrows(): - table.add_row([str(v) for v in row]) - return str(table) - - -def initialize_agent_from_dce(project_path: Path, model: str | None, temperature: float) -> Agent: - """Initialize the Databao agent using DCE project at the given path.""" - project = ProjectLayout(project_path) - - status = databao_project_status(project) - if status == DatabaoProjectStatus.NOT_INITIALIZED: - click.echo( - f"No Databao project found at {project.project_dir}. Run 'databao init' first.", - err=True, - ) - sys.exit(1) - - if status == DatabaoProjectStatus.NO_DATASOURCES: - click.echo( - f"No datasources configured in project at {project.project_dir}. Add datasources first.", - err=True, - ) - sys.exit(1) - - click.echo(f"Using DCE project: {project.project_dir}") - - _domain = create_domain(project.root_domain_dir) - - if model: - from databao_cli.executor_utils import build_llm_config - - llm_config = build_llm_config(model, temperature=temperature) - else: - # Use default but with custom temperature if provided - if temperature != 0.0: - llm_config = LLMConfig( - name=LLMConfigDirectory.DEFAULT.name, - temperature=temperature, - ) - else: - llm_config = LLMConfigDirectory.DEFAULT - - agent = create_agent(domain=_domain, llm_config=llm_config) - - num_sources = len(agent.sources.dbs) + len(agent.sources.dfs) - click.echo(f"Connected to {num_sources} data source(s)") - return agent - - -def display_result(thread: Thread) -> None: - """Display the execution result from thread to CLI.""" - # Display text response - text = thread.text() - if text: - click.echo(text) - - # Display SQL code if present - code = thread.code() - if code: - click.echo(f"\n```sql\n{code}\n```") - - # Display dataframe if present - df = thread.df() - if df is not None: - rows_shown = min(DEFAULT_MAX_DISPLAY_ROWS, len(df)) - click.echo(f"\n[DataFrame: {rows_shown} / {len(df)} rows]") - click.echo(dataframe_to_prettytable(df)) - - -def _print_help() -> None: - """Print help message for interactive mode.""" - click.echo("Databao REPL") - click.echo("Ask questions about your data in natural language.\n") - click.echo("Commands:") - click.echo(" \\help - Show this help") - click.echo(" \\clear - Start a new conversation") - click.echo(" \\q - Exit\n") - - -def run_interactive_mode(agent: Agent, show_thinking: bool) -> None: - """Run the interactive REPL mode.""" - click.echo("\nDatabao REPL") - click.echo("\nType \\help for available commands.\n") - - writer = _create_cli_writer() if show_thinking else None - - # Create thread with writer for streaming - thread = agent.thread( - stream_ask=show_thinking, - writer=writer, - ) - - while True: - try: - user_input = click.prompt("You", prompt_suffix="> ") - except (EOFError, KeyboardInterrupt): - click.echo() - break - - user_input = user_input.strip() - if not user_input: - continue - - # Check if it's a command (starts with \) - if user_input.startswith("\\"): - command = user_input[1:].lower() - - if command in ("exit", "quit", "q"): - break - - if command == "clear": - # Clear writer and create new thread - if writer: - writer.clear() - thread = agent.thread( - stream_ask=show_thinking, - writer=writer, - ) - click.echo("Conversation cleared.\n") - continue - - if command == "help": - _print_help() - continue - - click.echo(f"Unknown command: {user_input}. Type \\help for available commands.\n") - continue - - # Process as a question - try: - thread.ask(user_input, stream=show_thinking) - - # Clear writer buffer if present - if writer: - writer.clear() - - # Display result - click.echo("\nAssistant:") - display_result(thread) - click.echo() - - except Exception as e: - if writer: - writer.clear() - click.echo(f"\nError: {format_llm_error(e)}\n", err=True) - - -def run_one_shot_mode(agent: Agent, question: str, show_thinking: bool) -> None: - """Run a single question and exit.""" - - writer = _create_cli_writer() if show_thinking else None - - # Create thread with writer for streaming - thread = agent.thread( - stream_ask=show_thinking, - writer=writer, - ) - - try: - # Execute with streaming enabled when showing thinking - thread.ask(question, stream=show_thinking) - - # Display result - display_result(thread) - - except Exception as e: - click.echo(f"Error: {format_llm_error(e)}", err=True) - sys.exit(1) - - -def ask_impl( - project_path: Path, - question: str | None, - one_shot: bool, - model: str | None, - temperature: float, - show_thinking: bool, -) -> None: - # Validate arguments - if one_shot and not question: - click.echo("Error: QUESTION argument is required in one-shot mode.", err=True) - sys.exit(1) + from databao_cli.features.ask.service import ask_impl + from databao_cli.workflows.ask import run_interactive_mode, run_one_shot_mode - # Initialize agent (with progress indicator if not showing thinking) if not show_thinking: click.echo("Initializing agent...") - agent = initialize_agent_from_dce(project_path, model, temperature) + agent = ask_impl(ctx.obj["project_dir"], question, one_shot, model, temperature) - # Run appropriate mode if one_shot: assert question is not None run_one_shot_mode(agent, question, show_thinking) diff --git a/src/databao_cli/commands/build.py b/src/databao_cli/commands/build.py index 1a88c3e9..28499c9b 100644 --- a/src/databao_cli/commands/build.py +++ b/src/databao_cli/commands/build.py @@ -1,7 +1,6 @@ import click -from databao_context_engine import BuildDatasourceResult, DatabaoContextDomainManager -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import handle_feature_errors @click.command() @@ -19,6 +18,7 @@ help="Whether to index the context. If disabled, the context will be built but not indexed.", ) @click.pass_context +@handle_feature_errors def build(ctx: click.Context, domain: str, should_index: bool) -> None: """Build context for all domain's datasources. @@ -26,17 +26,9 @@ def build(ctx: click.Context, domain: str, should_index: bool) -> None: Internally, this indexes the context to be used by the MCP server and the "retrieve" command. """ - from databao_cli.commands._utils import get_project_or_exit + from databao_cli.features.build import build_impl + from databao_cli.shared.cli_utils import get_project_or_raise - project_layout = get_project_or_exit(ctx.obj["project_dir"]) + project_layout = get_project_or_raise(ctx.obj["project_dir"]) results = build_impl(project_layout, domain, should_index) click.echo(f"Build complete. Processed {len(results)} datasources.") - - -def build_impl(project_layout: ProjectLayout, domain: str, should_index: bool) -> list[BuildDatasourceResult]: - dce_project_dir = project_layout.domains_dir / domain - results: list[BuildDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).build_context( - datasource_ids=None, should_index=should_index - ) - - return results diff --git a/src/databao_cli/commands/datasource/add.py b/src/databao_cli/commands/datasource/add.py index 83f9c45d..ecd73ffe 100644 --- a/src/databao_cli/commands/datasource/add.py +++ b/src/databao_cli/commands/datasource/add.py @@ -1,15 +1,6 @@ -import os - import click -from databao_context_engine import ( - DatabaoContextDomainManager, - DatabaoContextPluginLoader, - DatasourceType, -) -from databao_cli.commands._utils import get_project_or_exit -from databao_cli.commands.context_engine_cli import ClickUserInputCallback -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import get_project_or_raise, handle_feature_errors @click.command(name="add") @@ -21,57 +12,13 @@ help="Databao domain name", ) @click.pass_context +@handle_feature_errors def add(ctx: click.Context, domain: str) -> None: """Add a new datasource configuration. The command will ask all relevant information for that datasource and save it in a chosen Databao domain """ - project_layout = get_project_or_exit(ctx.obj["project_dir"]) - add_impl(project_layout, domain) - - -def add_impl(project_layout: ProjectLayout, domain: str) -> None: - from databao_cli.commands.datasource.check import print_connection_check_results - - domain_dir = project_layout.domains_dir / domain - domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) - plugin_loader = DatabaoContextPluginLoader() - - click.echo(f"We will guide you to add a new datasource into {domain} domain, at {domain_dir.resolve()}") - - datasource_type = _ask_for_datasource_type(plugin_loader.get_all_supported_datasource_types(exclude_file_plugins=True)) - - datasource_name = click.prompt("Datasource name?", type=str) - - datasource_id = domain_manager.datasource_config_exists(datasource_name=datasource_name) - if datasource_id is not None: - click.confirm( - f"A config file already exists for this datasource {datasource_id.relative_path_to_config_file()}. " - f"Do you want to overwrite it?", - abort=True, - default=False, - ) - created_datasource = domain_manager.create_datasource_config_interactively( - datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True - ) - - datasource_id = created_datasource.datasource.id - click.echo( - f"{os.linesep}We've created a new config file for your datasource at: " - f"{domain_manager.get_config_file_path_for_datasource(datasource_id)}" - ) - if click.confirm("\nDo you want to check the connection to this new datasource?"): - results = domain_manager.check_datasource_connection(datasource_ids=[datasource_id]) - print_connection_check_results(domain, results) - - -def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType: - all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types]) - config_type = click.prompt( - "What type of datasource do you want to add?", - type=click.Choice(all_datasource_types), - default=all_datasource_types[0] if len(all_datasource_types) == 1 else None, - ) - click.echo(f"Selected type: {config_type}") + from databao_cli.workflows.datasource.add import add_workflow - return DatasourceType(full_type=config_type) + project_layout = get_project_or_raise(ctx.obj["project_dir"]) + add_workflow(project_layout, domain) diff --git a/src/databao_cli/commands/datasource/check.py b/src/databao_cli/commands/datasource/check.py index 1d2f72f2..c58b85bc 100644 --- a/src/databao_cli/commands/datasource/check.py +++ b/src/databao_cli/commands/datasource/check.py @@ -1,14 +1,6 @@ -import os - import click -from databao_context_engine import ( - CheckDatasourceConnectionResult, - DatabaoContextDomainManager, - DatasourceId, -) -from databao_cli.commands._utils import get_project_or_exit -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import get_project_or_raise, handle_feature_errors @click.command(name="check") @@ -18,6 +10,7 @@ nargs=-1, ) @click.pass_context +@handle_feature_errors def check(ctx: click.Context, domains: tuple[str, ...]) -> None: """Check whether a datasource configuration is valid. @@ -26,48 +19,14 @@ def check(ctx: click.Context, domains: tuple[str, ...]) -> None: By default, all declared datasources across all domains in the project will be checked. You can explicitly list which domains to validate by using the [DOMAINS] argument. """ - project_layout = get_project_or_exit(ctx.obj["project_dir"]) - check_impl(project_layout, requested_domains=list(domains) if domains else None) - + from databao_cli.features.datasource.check import check_impl + from databao_cli.workflows.datasource.check import print_connection_check_results -def check_impl(project_layout: ProjectLayout, requested_domains: list[str] | None) -> None: - domains: list[str] = project_layout.get_domain_names() if requested_domains is None else requested_domains + project_layout = get_project_or_raise(ctx.obj["project_dir"]) + results = check_impl(project_layout, requested_domains=list(domains) if domains else None) - results = _check_domains(project_layout, domains) - - if all([len(domain_results) == 0 for domain_results in results.values()]): + if all(len(v) == 0 for v in results.values()): click.echo("No datasource found") - return - - for domain, datasource_results in results.items(): - print_connection_check_results(domain, datasource_results) - - -def print_connection_check_results( - domain: str, datasource_results: dict[DatasourceId, CheckDatasourceConnectionResult] -) -> None: - for result in datasource_results.values(): - fq_datasource_name = domain + os.pathsep + str(result.datasource_id) - status = str(result.connection_status.value) - if result.summary: - status += f" - {result.summary}" - if result.full_message: - status += f": {result.full_message}" - - click.echo(f"{fq_datasource_name}: {status}") - - -def _check_domains( - project_layout: ProjectLayout, domains: list[str] -) -> dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]]: - results: dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]] = {} - for domain in domains: - domain_dir = project_layout.domains_dir / domain - if not domain_dir.exists(): - raise ValueError( - f"The specified {domain} domain does not exist. " - f"Available domains: {', '.join(project_layout.get_domain_names())}" - ) - domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) - results[domain] = domain_manager.check_datasource_connection() - return results + else: + for domain, datasource_results in results.items(): + print_connection_check_results(domain, datasource_results) diff --git a/src/databao_cli/commands/index.py b/src/databao_cli/commands/index.py index 445e9135..b6f2b5b0 100644 --- a/src/databao_cli/commands/index.py +++ b/src/databao_cli/commands/index.py @@ -1,7 +1,6 @@ import click -from databao_context_engine import ChunkEmbeddingMode, DatabaoContextDomainManager, DatasourceId, IndexDatasourceResult -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import handle_feature_errors @click.command() @@ -18,29 +17,17 @@ type=click.STRING, ) @click.pass_context +@handle_feature_errors def index(ctx: click.Context, domain: str, datasources_config_files: tuple[str, ...]) -> None: """Index built contexts into the embeddings database. If one or more datasource config file strings are provided, only those datasources will be indexed. If no values are provided, all built contexts for the selected domain will be indexed. """ - from databao_cli.commands._utils import get_project_or_exit + from databao_cli.features.index import index_impl + from databao_cli.shared.cli_utils import get_project_or_raise - project_layout = get_project_or_exit(ctx.obj["project_dir"]) + project_layout = get_project_or_raise(ctx.obj["project_dir"]) datasources = list(datasources_config_files) if datasources_config_files else None results = index_impl(project_layout, domain, datasources) click.echo(f"Index complete. Processed {len(results)} contexts.") - - -def index_impl( - project_layout: ProjectLayout, domain: str, datasources_config_files: list[str] | None -) -> list[IndexDatasourceResult]: - dce_project_dir = project_layout.domains_dir / domain - - datasource_ids = [DatasourceId.from_string_repr(p) for p in datasources_config_files] if datasources_config_files else None - - results: list[IndexDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).index_built_contexts( - datasource_ids=datasource_ids, chunk_embedding_mode=ChunkEmbeddingMode.EMBEDDABLE_TEXT_ONLY - ) - - return results diff --git a/src/databao_cli/commands/init.py b/src/databao_cli/commands/init.py index d05d1a72..51f75899 100644 --- a/src/databao_cli/commands/init.py +++ b/src/databao_cli/commands/init.py @@ -2,16 +2,19 @@ from pathlib import Path import click -from databao_context_engine import InitDomainError, init_dce_domain -from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project +from databao_cli.features.init.errors import InitDatabaoProjectError, ProjectDirDoesnotExistError +from databao_cli.features.init.service import init_impl +from databao_cli.shared.cli_utils import handle_feature_errors +from databao_cli.shared.project.layout import ROOT_DOMAIN, ProjectLayout @click.command() @click.pass_context +@handle_feature_errors def init(ctx: click.Context) -> None: """Create a new Databao project.""" - from databao_cli.commands.datasource.add import add_impl as add_datasource_config_interactive_impl + from databao_cli.workflows.datasource.add import add_workflow as add_datasource_config_interactive_impl project_dir: Path = ctx.obj["project_dir"] project_layout: ProjectLayout @@ -39,70 +42,3 @@ def init(ctx: click.Context) -> None: while click.confirm("\nDo you want to add more datasources?"): add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN) - - -class InitDatabaoProjectError(ValueError): - def __init__(self, message: str | None): - super().__init__(message or "") - self.message = message - - -class DatabaoProjectDirAlreadyExistsError(InitDatabaoProjectError): - def __init__(self, message: str) -> None: - super().__init__(message) - - -class ParentDatabaoProjectAlreadyExistsError(InitDatabaoProjectError): - def __init__(self, message: str) -> None: - super().__init__(message) - - -class ProjectDirDoesnotExistError(InitDatabaoProjectError): - def __init__(self, message: str) -> None: - super().__init__(message) - - -class ProjectDirNotDirError(InitDatabaoProjectError): - def __init__(self, message: str) -> None: - super().__init__(message) - - -class DatabaoContextEngineProjectInitError(InitDatabaoProjectError): - def __init__(self, message: str) -> None: - super().__init__(message) - - -def init_impl(project_dir: Path) -> ProjectLayout: - project_creator = _ProjectCreator(project_dir=project_dir) - return project_creator.create() - - -class _ProjectCreator: - def __init__(self, project_dir: Path): - self.project_dir = project_dir - - def create(self) -> ProjectLayout: - self.ensure_can_init_project() - - project_layout = ProjectLayout(self.project_dir) - project_layout.root_domain_dir.mkdir(parents=True, exist_ok=False) - project_layout.agents_dir.mkdir(parents=True, exist_ok=True) - - try: - init_dce_domain(project_layout.root_domain_dir) - except InitDomainError as e: - raise DatabaoContextEngineProjectInitError(str(e)) from e - - return project_layout - - def ensure_can_init_project(self) -> None: - if not self.project_dir.exists(): - raise ProjectDirDoesnotExistError(f"The project directory doesn't exist: {self.project_dir.resolve()}") - if not self.project_dir.is_dir(): - raise ProjectDirNotDirError(f"The project directory is not a directory: {self.project_dir.resolve()}") - - existing_project = find_project(self.project_dir) - if existing_project is not None: - raise DatabaoProjectDirAlreadyExistsError( - f"Can't initialize Databao project. It already exists - {existing_project.databao_dir.resolve()}" - ) diff --git a/src/databao_cli/commands/mcp.py b/src/databao_cli/commands/mcp.py index 880e644a..e06a96a1 100644 --- a/src/databao_cli/commands/mcp.py +++ b/src/databao_cli/commands/mcp.py @@ -1,10 +1,8 @@ """databao mcp command - Run an MCP server exposing Databao tools.""" -from pathlib import Path - import click -from databao_cli.mcp.server import McpContext, run_server +from databao_cli.shared.cli_utils import handle_feature_errors @click.command() @@ -30,6 +28,7 @@ help="Port to listen on (SSE transport only).", ) @click.pass_context +@handle_feature_errors def mcp(ctx: click.Context, transport: str, host: str, port: int) -> None: """Run an MCP server exposing Databao tools. @@ -42,14 +41,10 @@ def mcp(ctx: click.Context, transport: str, host: str, port: int) -> None: databao mcp --transport sse # SSE on localhost:8765 databao mcp --transport sse --port 9000 # SSE on custom port """ - from databao_cli.log.logging import configure_logging - from databao_cli.project.layout import find_project + from databao_cli.features.mcp.server import mcp_impl + from databao_cli.shared.log.logging import configure_logging + from databao_cli.shared.project.layout import find_project if transport == "stdio": configure_logging(find_project(ctx.obj["project_dir"]), quiet=True) mcp_impl(ctx.obj["project_dir"], transport, host, port) - - -def mcp_impl(project_dir: Path, transport: str, host: str, port: int) -> None: - context = McpContext(project_dir=project_dir) - run_server(context, transport=transport, host=host, port=port) diff --git a/src/databao_cli/commands/status.py b/src/databao_cli/commands/status.py index 551a4cb7..492b4b21 100644 --- a/src/databao_cli/commands/status.py +++ b/src/databao_cli/commands/status.py @@ -1,54 +1,13 @@ -import os -import sys -from importlib.metadata import version -from pathlib import Path - import click -from databao_context_engine import ( - DceDomainInfo, - DceInfo, - get_databao_context_engine_domain_info, - get_databao_context_engine_info, -) -from databao_cli.commands._utils import get_project_or_exit +from databao_cli.shared.cli_utils import handle_feature_errors @click.command() @click.pass_context +@handle_feature_errors def status(ctx: click.Context) -> None: """Display project status and system-wide information.""" - click.echo(status_impl(ctx.obj["project_dir"])) - - -def status_impl(project_dir: Path) -> str: - project_layout = get_project_or_exit(project_dir) - - dce_info = get_databao_context_engine_info() + from databao_cli.features.status import status_impl - return _generate_info_string( - dce_info, - [get_databao_context_engine_domain_info(domain) for domain in project_layout.get_domain_dirs()] - if project_layout - else [], - ) - - -def _generate_info_string(command_info: DceInfo, domain_infos: list[DceDomainInfo]) -> str: - info_lines = [ - f"Databao context engine version: {command_info.version}", - f"Databao agent version: {version('databao-agent')}", - f"Databao context engine storage dir: {command_info.dce_path}", - f"Databao context engine plugins: {command_info.plugin_ids}", - "", - f"OS name: {sys.platform}", - f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}", - "", - ] - - for domain_info in domain_infos: - if domain_info.is_initialized: - info_lines.append(f"Databao Domain dir: {domain_info.project_path.resolve()}") - info_lines.append(f"Databao Domain ID: {domain_info.project_id!s}") - - return os.linesep.join(info_lines) + click.echo(status_impl(ctx.obj["project_dir"])) diff --git a/src/databao_cli/log/__init__.py b/src/databao_cli/features/__init__.py similarity index 100% rename from src/databao_cli/log/__init__.py rename to src/databao_cli/features/__init__.py diff --git a/src/databao_cli/mcp/__init__.py b/src/databao_cli/features/ask/__init__.py similarity index 100% rename from src/databao_cli/mcp/__init__.py rename to src/databao_cli/features/ask/__init__.py diff --git a/src/databao_cli/features/ask/agent_factory.py b/src/databao_cli/features/ask/agent_factory.py new file mode 100644 index 00000000..f3bf42b6 --- /dev/null +++ b/src/databao_cli/features/ask/agent_factory.py @@ -0,0 +1,41 @@ +from pathlib import Path + +from databao.agent import Agent +from databao.agent import domain as create_domain +from databao.agent.api import agent as create_agent +from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory + +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.shared.errors import FeatureError +from databao_cli.shared.project.layout import ProjectLayout + + +def initialize_agent_from_dce(project_path: Path, model: str | None, temperature: float) -> Agent: + """Initialize the Databao agent using DCE project at the given path.""" + project = ProjectLayout(project_path) + + status = databao_project_status(project) + if status == DatabaoProjectStatus.NOT_INITIALIZED: + raise FeatureError(f"No Databao project found at {project.project_dir}. Run 'databao init' first.") + + if status == DatabaoProjectStatus.NO_DATASOURCES: + raise FeatureError(f"No datasources configured in project at {project.project_dir}. Add datasources first.") + + _domain = create_domain(project.root_domain_dir) + + if model: + from databao_cli.shared.executor_utils import build_llm_config + + llm_config = build_llm_config(model, temperature=temperature) + else: + if temperature != 0.0: + llm_config = LLMConfig( + name=LLMConfigDirectory.DEFAULT.name, + temperature=temperature, + ) + else: + llm_config = LLMConfigDirectory.DEFAULT + + agent = create_agent(domain=_domain, llm_config=llm_config) + + return agent diff --git a/src/databao_cli/features/ask/display.py b/src/databao_cli/features/ask/display.py new file mode 100644 index 00000000..7a3b0897 --- /dev/null +++ b/src/databao_cli/features/ask/display.py @@ -0,0 +1,13 @@ +import pandas as pd +from prettytable import PrettyTable + +DEFAULT_MAX_DISPLAY_ROWS = 10 + + +def dataframe_to_prettytable(df: pd.DataFrame, max_rows: int = DEFAULT_MAX_DISPLAY_ROWS) -> str: + """Convert a pandas DataFrame to a prettytable string.""" + table = PrettyTable() + table.field_names = list(df.columns) + for _, row in df.head(max_rows).iterrows(): + table.add_row([str(v) for v in row]) + return str(table) diff --git a/src/databao_cli/features/ask/service.py b/src/databao_cli/features/ask/service.py new file mode 100644 index 00000000..13716519 --- /dev/null +++ b/src/databao_cli/features/ask/service.py @@ -0,0 +1,18 @@ +from pathlib import Path + +from databao.agent import Agent + +from databao_cli.features.ask.agent_factory import initialize_agent_from_dce +from databao_cli.shared.errors import FeatureError + + +def ask_impl( + project_path: Path, + question: str | None, + one_shot: bool, + model: str | None, + temperature: float, +) -> Agent: + if one_shot and not question: + raise FeatureError("Error: QUESTION argument is required in one-shot mode.") + return initialize_agent_from_dce(project_path, model, temperature) diff --git a/src/databao_cli/features/build.py b/src/databao_cli/features/build.py new file mode 100644 index 00000000..5149489a --- /dev/null +++ b/src/databao_cli/features/build.py @@ -0,0 +1,11 @@ +from databao_context_engine import BuildDatasourceResult, DatabaoContextDomainManager + +from databao_cli.shared.project.layout import ProjectLayout + + +def build_impl(project_layout: ProjectLayout, domain: str, should_index: bool) -> list[BuildDatasourceResult]: + dce_project_dir = project_layout.domains_dir / domain + results: list[BuildDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).build_context( + datasource_ids=None, should_index=should_index + ) + return results diff --git a/src/databao_cli/mcp/tools/__init__.py b/src/databao_cli/features/datasource/__init__.py similarity index 100% rename from src/databao_cli/mcp/tools/__init__.py rename to src/databao_cli/features/datasource/__init__.py diff --git a/src/databao_cli/features/datasource/add.py b/src/databao_cli/features/datasource/add.py new file mode 100644 index 00000000..ad4366ab --- /dev/null +++ b/src/databao_cli/features/datasource/add.py @@ -0,0 +1,37 @@ +from pathlib import Path + +from databao_context_engine import ( + DatabaoContextDomainManager, + DatasourceId, + DatasourceType, + UserInputCallback, +) + +from databao_cli.shared.project.layout import ProjectLayout + + +def datasource_config_exists(project_layout: ProjectLayout, domain: str, datasource_name: str) -> DatasourceId | None: + """Return the existing DatasourceId if a config already exists for this name, else None.""" + domain_dir = project_layout.domains_dir / domain + domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) + return domain_manager.datasource_config_exists(datasource_name=datasource_name) + + +def create_datasource_config( + project_layout: ProjectLayout, + domain: str, + datasource_type: DatasourceType, + datasource_name: str, + user_input_callback: UserInputCallback, + *, + overwrite_existing: bool = False, +) -> tuple[DatasourceId, Path]: + """Create a datasource config file. Returns (datasource_id, config_file_path).""" + domain_dir = project_layout.domains_dir / domain + domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) + created_datasource = domain_manager.create_datasource_config_interactively( + datasource_type, datasource_name, user_input_callback, overwrite_existing=overwrite_existing + ) + datasource_id = created_datasource.datasource.id + config_path = domain_manager.get_config_file_path_for_datasource(datasource_id) + return datasource_id, config_path diff --git a/src/databao_cli/features/datasource/check.py b/src/databao_cli/features/datasource/check.py new file mode 100644 index 00000000..89e542ec --- /dev/null +++ b/src/databao_cli/features/datasource/check.py @@ -0,0 +1,30 @@ +from databao_context_engine import ( + CheckDatasourceConnectionResult, + DatabaoContextDomainManager, + DatasourceId, +) + +from databao_cli.shared.project.layout import ProjectLayout + + +def check_impl( + project_layout: ProjectLayout, requested_domains: list[str] | None +) -> dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]]: + domains: list[str] = project_layout.get_domain_names() if requested_domains is None else requested_domains + return _check_domains(project_layout, domains) + + +def _check_domains( + project_layout: ProjectLayout, domains: list[str] +) -> dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]]: + results: dict[str, dict[DatasourceId, CheckDatasourceConnectionResult]] = {} + for domain in domains: + domain_dir = project_layout.domains_dir / domain + if not domain_dir.exists(): + raise ValueError( + f"The specified {domain} domain does not exist. " + f"Available domains: {', '.join(project_layout.get_domain_names())}" + ) + domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) + results[domain] = domain_manager.check_datasource_connection() + return results diff --git a/src/databao_cli/features/index.py b/src/databao_cli/features/index.py new file mode 100644 index 00000000..d865c33e --- /dev/null +++ b/src/databao_cli/features/index.py @@ -0,0 +1,17 @@ +from databao_context_engine import ChunkEmbeddingMode, DatabaoContextDomainManager, DatasourceId, IndexDatasourceResult + +from databao_cli.shared.project.layout import ProjectLayout + + +def index_impl( + project_layout: ProjectLayout, domain: str, datasources_config_files: list[str] | None +) -> list[IndexDatasourceResult]: + dce_project_dir = project_layout.domains_dir / domain + + datasource_ids = [DatasourceId.from_string_repr(p) for p in datasources_config_files] if datasources_config_files else None + + results: list[IndexDatasourceResult] = DatabaoContextDomainManager(domain_dir=dce_project_dir).index_built_contexts( + datasource_ids=datasource_ids, chunk_embedding_mode=ChunkEmbeddingMode.EMBEDDABLE_TEXT_ONLY + ) + + return results diff --git a/src/databao_cli/project/__init__.py b/src/databao_cli/features/init/__init__.py similarity index 100% rename from src/databao_cli/project/__init__.py rename to src/databao_cli/features/init/__init__.py diff --git a/src/databao_cli/features/init/errors.py b/src/databao_cli/features/init/errors.py new file mode 100644 index 00000000..0a8ec28b --- /dev/null +++ b/src/databao_cli/features/init/errors.py @@ -0,0 +1,29 @@ +class InitDatabaoProjectError(ValueError): + def __init__(self, message: str | None): + super().__init__(message or "") + self.message = message + + +class DatabaoProjectDirAlreadyExistsError(InitDatabaoProjectError): + def __init__(self, message: str) -> None: + super().__init__(message) + + +class ParentDatabaoProjectAlreadyExistsError(InitDatabaoProjectError): + def __init__(self, message: str) -> None: + super().__init__(message) + + +class ProjectDirDoesnotExistError(InitDatabaoProjectError): + def __init__(self, message: str) -> None: + super().__init__(message) + + +class ProjectDirNotDirError(InitDatabaoProjectError): + def __init__(self, message: str) -> None: + super().__init__(message) + + +class DatabaoContextEngineProjectInitError(InitDatabaoProjectError): + def __init__(self, message: str) -> None: + super().__init__(message) diff --git a/src/databao_cli/features/init/service.py b/src/databao_cli/features/init/service.py new file mode 100644 index 00000000..7531784b --- /dev/null +++ b/src/databao_cli/features/init/service.py @@ -0,0 +1,47 @@ +from pathlib import Path + +from databao_context_engine import InitDomainError, init_dce_domain + +from databao_cli.features.init.errors import ( + DatabaoContextEngineProjectInitError, + DatabaoProjectDirAlreadyExistsError, + ProjectDirDoesnotExistError, + ProjectDirNotDirError, +) +from databao_cli.shared.project.layout import ProjectLayout, find_project + + +def init_impl(project_dir: Path) -> ProjectLayout: + project_creator = _ProjectCreator(project_dir=project_dir) + return project_creator.create() + + +class _ProjectCreator: + def __init__(self, project_dir: Path): + self.project_dir = project_dir + + def create(self) -> ProjectLayout: + self.ensure_can_init_project() + + project_layout = ProjectLayout(self.project_dir) + project_layout.root_domain_dir.mkdir(parents=True, exist_ok=False) + project_layout.agents_dir.mkdir(parents=True, exist_ok=True) + + try: + init_dce_domain(project_layout.root_domain_dir) + except InitDomainError as e: + raise DatabaoContextEngineProjectInitError(str(e)) from e + + return project_layout + + def ensure_can_init_project(self) -> None: + if not self.project_dir.exists(): + raise ProjectDirDoesnotExistError(f"The project directory doesn't exist: {self.project_dir.resolve()}") + if not self.project_dir.is_dir(): + raise ProjectDirNotDirError(f"The project directory is not a directory: {self.project_dir.resolve()}") + + existing_project = find_project(self.project_dir) + if existing_project is not None: + raise DatabaoProjectDirAlreadyExistsError( + f"Can't initialize Databao project. It already exists - {existing_project.databao_dir.resolve()}" + ) diff --git a/src/databao_cli/features/mcp/__init__.py b/src/databao_cli/features/mcp/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/mcp/server.py b/src/databao_cli/features/mcp/server.py similarity index 79% rename from src/databao_cli/mcp/server.py rename to src/databao_cli/features/mcp/server.py index 194c7e35..83668dc1 100644 --- a/src/databao_cli/mcp/server.py +++ b/src/databao_cli/features/mcp/server.py @@ -5,7 +5,7 @@ from fastmcp import FastMCP -from databao_cli.mcp.tools import databao_ask +from databao_cli.features.mcp.tools import databao_ask @dataclass(frozen=True) @@ -38,3 +38,8 @@ def run_server( mcp.run(transport="sse", host=host, port=port) case _: raise ValueError(f"Unknown transport: {transport!r}. Supported: stdio, sse") + + +def mcp_impl(project_dir: Path, transport: str, host: str, port: int) -> None: + context = McpContext(project_dir=project_dir) + run_server(context, transport=transport, host=host, port=port) diff --git a/src/databao_cli/features/mcp/tools/__init__.py b/src/databao_cli/features/mcp/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/features/mcp/tools/databao_ask/__init__.py b/src/databao_cli/features/mcp/tools/databao_ask/__init__.py new file mode 100644 index 00000000..f658a173 --- /dev/null +++ b/src/databao_cli/features/mcp/tools/databao_ask/__init__.py @@ -0,0 +1,5 @@ +"""databao_ask MCP tool package.""" + +from databao_cli.features.mcp.tools.databao_ask.tool import register + +__all__ = ["register"] diff --git a/src/databao_cli/mcp/tools/databao_ask/agent_factory.py b/src/databao_cli/features/mcp/tools/databao_ask/agent_factory.py similarity index 88% rename from src/databao_cli/mcp/tools/databao_ask/agent_factory.py rename to src/databao_cli/features/mcp/tools/databao_ask/agent_factory.py index b4645e7b..8e59ba97 100644 --- a/src/databao_cli/mcp/tools/databao_ask/agent_factory.py +++ b/src/databao_cli/features/mcp/tools/databao_ask/agent_factory.py @@ -8,8 +8,8 @@ from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory from databao.agent.core import Cache -from databao_cli.project.layout import ProjectLayout -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status, has_build_output +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status, has_build_output +from databao_cli.shared.project.layout import ProjectLayout def create_agent_for_tool( @@ -37,7 +37,7 @@ def create_agent_for_tool( llm_config: LLMConfig if model: - from databao_cli.executor_utils import build_llm_config + from databao_cli.shared.executor_utils import build_llm_config llm_config = build_llm_config(model, temperature=temperature) elif temperature != 0.0: diff --git a/src/databao_cli/mcp/tools/databao_ask/tool.py b/src/databao_cli/features/mcp/tools/databao_ask/tool.py similarity index 97% rename from src/databao_cli/mcp/tools/databao_ask/tool.py rename to src/databao_cli/features/mcp/tools/databao_ask/tool.py index 41d1fbf2..8d76c496 100644 --- a/src/databao_cli/mcp/tools/databao_ask/tool.py +++ b/src/databao_cli/features/mcp/tools/databao_ask/tool.py @@ -10,12 +10,12 @@ from pydantic import BaseModel, Field from uuid6 import uuid6 -from databao_cli.mcp.tools.databao_ask.agent_factory import create_agent_for_tool +from databao_cli.features.mcp.tools.databao_ask.agent_factory import create_agent_for_tool if TYPE_CHECKING: from fastmcp import FastMCP - from databao_cli.mcp.server import McpContext + from databao_cli.features.mcp.server import McpContext logger = logging.getLogger(__name__) diff --git a/src/databao_cli/features/status.py b/src/databao_cli/features/status.py new file mode 100644 index 00000000..f2e617d6 --- /dev/null +++ b/src/databao_cli/features/status.py @@ -0,0 +1,46 @@ +import os +import sys +from importlib.metadata import version +from pathlib import Path + +from databao_context_engine import ( + DceDomainInfo, + DceInfo, + get_databao_context_engine_domain_info, + get_databao_context_engine_info, +) + +from databao_cli.shared.project.layout import find_project + + +def status_impl(project_dir: Path) -> str: + project_layout = find_project(project_dir) + + dce_info = get_databao_context_engine_info() + + return _generate_info_string( + dce_info, + [get_databao_context_engine_domain_info(domain) for domain in project_layout.get_domain_dirs()] + if project_layout + else [], + ) + + +def _generate_info_string(command_info: DceInfo, domain_infos: list[DceDomainInfo]) -> str: + info_lines = [ + f"Databao context engine version: {command_info.version}", + f"Databao agent version: {version('databao-agent')}", + f"Databao context engine storage dir: {command_info.dce_path}", + f"Databao context engine plugins: {command_info.plugin_ids}", + "", + f"OS name: {sys.platform}", + f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}", + "", + ] + + for domain_info in domain_infos: + if domain_info.is_initialized: + info_lines.append(f"Databao Domain dir: {domain_info.project_path.resolve()}") + info_lines.append(f"Databao Domain ID: {domain_info.project_id!s}") + + return os.linesep.join(info_lines) diff --git a/src/databao_cli/ui/__init__.py b/src/databao_cli/features/ui/__init__.py similarity index 100% rename from src/databao_cli/ui/__init__.py rename to src/databao_cli/features/ui/__init__.py diff --git a/src/databao_cli/ui/app.py b/src/databao_cli/features/ui/app.py similarity index 90% rename from src/databao_cli/ui/app.py rename to src/databao_cli/features/ui/app.py index 6bcf87d8..d133d17d 100644 --- a/src/databao_cli/ui/app.py +++ b/src/databao_cli/features/ui/app.py @@ -14,20 +14,20 @@ from databao.agent.core.agent import Agent from streamlit.navigation.page import StreamlitPage -from databao_cli.project.layout import ProjectLayout, find_project -from databao_cli.ui.components.status import AppStatus, set_status, status_context -from databao_cli.ui.models.chat_session import ChatSession -from databao_cli.ui.models.settings import LLMSettings -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status, has_build_output -from databao_cli.ui.services.storage import get_cache_dir +from databao_cli.features.ui.components.status import AppStatus, set_status, status_context +from databao_cli.features.ui.models.chat_session import ChatSession +from databao_cli.features.ui.models.settings import LLMSettings +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status, has_build_output +from databao_cli.features.ui.services.storage import get_cache_dir +from databao_cli.shared.project.layout import ProjectLayout, find_project logger = logging.getLogger(__name__) def _load_persisted_state() -> None: """Load settings and chats from disk on startup.""" - from databao_cli.ui.services.chat_persistence import load_all_chats - from databao_cli.ui.services.settings_persistence import get_or_create_settings + from databao_cli.features.ui.services.chat_persistence import load_all_chats + from databao_cli.features.ui.services.settings_persistence import get_or_create_settings if "app_settings" not in st.session_state: settings = get_or_create_settings() @@ -46,8 +46,8 @@ def _load_persisted_state() -> None: def _save_settings_if_changed() -> None: """Save settings to disk if they've changed.""" - from databao_cli.ui.models.settings import Settings - from databao_cli.ui.services.settings_persistence import save_settings + from databao_cli.features.ui.models.settings import Settings + from databao_cli.features.ui.services.settings_persistence import save_settings settings: Settings | None = st.session_state.get("app_settings") if settings is None: @@ -139,7 +139,7 @@ def _initialize_agent(project: ProjectLayout) -> Agent | None: def _build_llm_config() -> LLMConfig | None: """Build an LLMConfig from session-state LLM settings, or None for defaults.""" - from databao_cli.ui.models.settings import _ENV_VAR_MAP + from databao_cli.features.ui.models.settings import _ENV_VAR_MAP llm: LLMSettings = st.session_state.get("llm_settings", LLMSettings()) @@ -157,7 +157,7 @@ def _build_llm_config() -> LLMConfig | None: if provider_type == "ollama" and config.base_url: os.environ["OLLAMA_HOST"] = config.base_url - from databao_cli.executor_utils import build_llm_config + from databao_cli.shared.executor_utils import build_llm_config return build_llm_config( config.model, @@ -221,7 +221,7 @@ def _load_welcome_completed(project: ProjectLayout | None) -> bool: if not settings_path.exists(): return False try: - from databao_cli.ui.models.settings import Settings + from databao_cli.features.ui.models.settings import Settings yaml_content = settings_path.read_text() settings = Settings.from_yaml(yaml_content) @@ -238,8 +238,8 @@ def mark_welcome_completed() -> None: during setup are saved, even if the user didn't change them from defaults (which means auto_apply never triggered). """ - from databao_cli.ui.models.settings import LLMSettings - from databao_cli.ui.services.settings_persistence import get_or_create_settings, save_settings + from databao_cli.features.ui.models.settings import LLMSettings + from databao_cli.features.ui.services.settings_persistence import get_or_create_settings, save_settings settings = get_or_create_settings() settings.welcome_completed = True @@ -325,7 +325,7 @@ def _create_new_chat() -> None: """Create a new chat and navigate to it.""" from uuid6 import uuid6 - from databao_cli.ui.services.chat_persistence import save_chat + from databao_cli.features.ui.services.chat_persistence import save_chat prev_chat_id = st.session_state.get("current_chat_id") chats: dict[str, ChatSession] = st.session_state.get("chats", {}) @@ -347,11 +347,11 @@ def _create_new_chat() -> None: def build_navigation() -> None: """Build the full multipage navigation structure (normal mode).""" - from databao_cli.ui.pages.agent_settings import render_agent_settings_page - from databao_cli.ui.pages.chat import render_chat_page - from databao_cli.ui.pages.context_settings import render_context_settings_page - from databao_cli.ui.pages.general_settings import render_general_settings_page - from databao_cli.ui.pages.welcome import render_welcome_page + from databao_cli.features.ui.pages.agent_settings import render_agent_settings_page + from databao_cli.features.ui.pages.chat import render_chat_page + from databao_cli.features.ui.pages.context_settings import render_context_settings_page + from databao_cli.features.ui.pages.general_settings import render_general_settings_page + from databao_cli.features.ui.pages.welcome import render_welcome_page navigate_to_chat: str | None = st.session_state.get("_navigate_to_chat") if navigate_to_chat: @@ -453,7 +453,7 @@ def page_fn() -> None: def build_setup_navigation() -> None: """Build navigation for setup mode -- only the setup wizard page, no sidebar.""" - from databao_cli.ui.pages.welcome import render_setup_wizard_page + from databao_cli.features.ui.pages.welcome import render_setup_wizard_page setup_page = st.Page( render_setup_wizard_page, @@ -486,7 +486,7 @@ def _render_global_sidebar() -> None: This is purely for UI rendering - initialization is handled by _initialize_app(). """ - from databao_cli.ui.components.sidebar import render_sidebar_header + from databao_cli.features.ui.components.sidebar import render_sidebar_header with st.sidebar: render_sidebar_header() diff --git a/src/databao_cli/ui/assets/bao.png b/src/databao_cli/features/ui/assets/bao.png similarity index 100% rename from src/databao_cli/ui/assets/bao.png rename to src/databao_cli/features/ui/assets/bao.png diff --git a/src/databao_cli/ui/cli.py b/src/databao_cli/features/ui/cli.py similarity index 59% rename from src/databao_cli/ui/cli.py rename to src/databao_cli/features/ui/cli.py index 54a125c7..37a2042a 100644 --- a/src/databao_cli/ui/cli.py +++ b/src/databao_cli/features/ui/cli.py @@ -4,15 +4,17 @@ import sys from pathlib import Path +from databao_cli.shared.errors import FeatureError + def _get_streamlit_app_path() -> str: """Get the path to the Streamlit app without importing it. This avoids triggering module-level Streamlit code during import. """ - spec = importlib.util.find_spec("databao_cli.ui.app") + spec = importlib.util.find_spec("databao_cli.features.ui.app") if spec is None or spec.origin is None: - raise ValueError("Could not find databao_cli.ui.app module. ") + raise ValueError("Could not find databao_cli.features.ui.app module. ") return spec.origin @@ -43,3 +45,22 @@ def bootstrap_streamlit_app( [sys.executable, "-m", "streamlit", "run", app_path, *streamlit_args, "--", *app_args], check=True, ) + + +def app_impl( + project_dir: Path, + extra_args: list[str], + read_only_domain: bool = False, + hide_suggested_questions: bool = False, + hide_build_context_hint: bool = False, +) -> None: + try: + bootstrap_streamlit_app( + project_dir, + extra_args, + read_only_domain=read_only_domain, + hide_suggested_questions=hide_suggested_questions, + hide_build_context_hint=hide_build_context_hint, + ) + except subprocess.CalledProcessError as e: + raise FeatureError(f"Error running Streamlit: {e}") from e diff --git a/src/databao_cli/ui/components/__init__.py b/src/databao_cli/features/ui/components/__init__.py similarity index 100% rename from src/databao_cli/ui/components/__init__.py rename to src/databao_cli/features/ui/components/__init__.py diff --git a/src/databao_cli/ui/components/chat.py b/src/databao_cli/features/ui/components/chat.py similarity index 97% rename from src/databao_cli/ui/components/chat.py rename to src/databao_cli/features/ui/components/chat.py index fe9bf925..6e5aa192 100644 --- a/src/databao_cli/ui/components/chat.py +++ b/src/databao_cli/features/ui/components/chat.py @@ -5,9 +5,9 @@ import streamlit as st from databao.agent.core.agent import Agent -from databao_cli.ui.components.results import render_execution_result -from databao_cli.ui.models.chat_session import ChatMessage -from databao_cli.ui.services import ( +from databao_cli.features.ui.components.results import render_execution_result +from databao_cli.features.ui.models.chat_session import ChatMessage +from databao_cli.features.ui.services import ( check_query_completion, get_query_phase, is_query_running, @@ -15,14 +15,14 @@ start_query_execution, stop_query, ) -from databao_cli.ui.suggestions import ( +from databao_cli.features.ui.suggestions import ( check_suggestions_completion, is_suggestions_loading, start_suggestions_generation, ) if TYPE_CHECKING: - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession def render_user_message(message: ChatMessage) -> None: @@ -557,6 +557,6 @@ def render_chat_interface(chat: "ChatSession") -> None: _render_chat_input_bar(chat, query_running) if "pending_plot_message_index" in st.session_state: - from databao_cli.ui.components.results import execute_pending_plot + from databao_cli.features.ui.components.results import execute_pending_plot execute_pending_plot(chat) diff --git a/src/databao_cli/ui/components/datasource_form.py b/src/databao_cli/features/ui/components/datasource_form.py similarity index 100% rename from src/databao_cli/ui/components/datasource_form.py rename to src/databao_cli/features/ui/components/datasource_form.py diff --git a/src/databao_cli/ui/components/datasource_manager.py b/src/databao_cli/features/ui/components/datasource_manager.py similarity index 97% rename from src/databao_cli/ui/components/datasource_manager.py rename to src/databao_cli/features/ui/components/datasource_manager.py index 5044592c..4da0f250 100644 --- a/src/databao_cli/ui/components/datasource_manager.py +++ b/src/databao_cli/features/ui/components/datasource_manager.py @@ -11,9 +11,9 @@ import streamlit as st from databao_context_engine import ConfiguredDatasource, DatasourceConnectionStatus -from databao_cli.ui.app import invalidate_agent -from databao_cli.ui.components.datasource_form import render_datasource_config_form -from databao_cli.ui.services.dce_operations import ( +from databao_cli.features.ui.app import invalidate_agent +from databao_cli.features.ui.components.datasource_form import render_datasource_config_form +from databao_cli.features.ui.services.dce_operations import ( add_datasource, get_available_datasource_types, get_datasource_config_fields, diff --git a/src/databao_cli/ui/components/icons.py b/src/databao_cli/features/ui/components/icons.py similarity index 100% rename from src/databao_cli/ui/components/icons.py rename to src/databao_cli/features/ui/components/icons.py diff --git a/src/databao_cli/ui/components/results.py b/src/databao_cli/features/ui/components/results.py similarity index 98% rename from src/databao_cli/ui/components/results.py rename to src/databao_cli/features/ui/components/results.py index 3a386690..1ae6483c 100644 --- a/src/databao_cli/ui/components/results.py +++ b/src/databao_cli/features/ui/components/results.py @@ -6,13 +6,13 @@ import nh3 import streamlit as st -from databao_cli.ui.services.chat_persistence import save_current_chat +from databao_cli.features.ui.services.chat_persistence import save_current_chat if TYPE_CHECKING: from databao.agent.core.executor import ExecutionResult from databao.agent.core.thread import Thread - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession logger = logging.getLogger(__name__) @@ -368,7 +368,7 @@ def _render_and_handle_action_buttons( rerun so that the input bar is rendered as disabled before the blocking ``thread.plot()`` call begins. """ - from databao_cli.ui.services import is_query_running + from databao_cli.features.ui.services import is_query_running thread = chat.thread if thread is None: diff --git a/src/databao_cli/ui/components/sidebar.py b/src/databao_cli/features/ui/components/sidebar.py similarity index 88% rename from src/databao_cli/ui/components/sidebar.py rename to src/databao_cli/features/ui/components/sidebar.py index 6d1d417a..f6a40b80 100644 --- a/src/databao_cli/ui/components/sidebar.py +++ b/src/databao_cli/features/ui/components/sidebar.py @@ -2,19 +2,19 @@ import streamlit as st -from databao_cli.executor_utils import EXECUTOR_TYPES -from databao_cli.project.layout import ProjectLayout -from databao_cli.ui.app import _clear_all_chat_threads -from databao_cli.ui.components.icons import get_db_type_and_icon -from databao_cli.ui.components.status import AppStatus, render_status_fragment, set_status -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status -from databao_cli.ui.suggestions import reset_suggestions_state +from databao_cli.features.ui.app import _clear_all_chat_threads +from databao_cli.features.ui.components.icons import get_db_type_and_icon +from databao_cli.features.ui.components.status import AppStatus, render_status_fragment, set_status +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.features.ui.suggestions import reset_suggestions_state +from databao_cli.shared.executor_utils import EXECUTOR_TYPES +from databao_cli.shared.project.layout import ProjectLayout @st.dialog("Delete Chat") def _confirm_delete_chat(chat_id: str, chat_title: str) -> None: """Dialog to confirm deleting the current chat.""" - from databao_cli.ui.services.chat_persistence import delete_chat + from databao_cli.features.ui.services.chat_persistence import delete_chat st.warning(f"⚠️ Delete chat: **{chat_title}**?") st.markdown("This will permanently remove the chat and its history.") @@ -156,7 +156,7 @@ def render_sidebar_chat_content(project: ProjectLayout | None) -> None: This is called only on chat pages to show project info, sources, and executor. Must be called within st.sidebar context. """ - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession render_project_info(project) diff --git a/src/databao_cli/ui/components/status.py b/src/databao_cli/features/ui/components/status.py similarity index 100% rename from src/databao_cli/ui/components/status.py rename to src/databao_cli/features/ui/components/status.py diff --git a/src/databao_cli/features/ui/models/__init__.py b/src/databao_cli/features/ui/models/__init__.py new file mode 100644 index 00000000..01552b77 --- /dev/null +++ b/src/databao_cli/features/ui/models/__init__.py @@ -0,0 +1,11 @@ +"""Data models for the Databao Streamlit app.""" + +from databao_cli.features.ui.models.chat_session import ChatMessage, ChatSession +from databao_cli.features.ui.models.settings import AgentSettings, Settings + +__all__ = [ + "AgentSettings", + "ChatMessage", + "ChatSession", + "Settings", +] diff --git a/src/databao_cli/ui/models/chat_session.py b/src/databao_cli/features/ui/models/chat_session.py similarity index 98% rename from src/databao_cli/ui/models/chat_session.py rename to src/databao_cli/features/ui/models/chat_session.py index 55c487fe..ec9700ca 100644 --- a/src/databao_cli/ui/models/chat_session.py +++ b/src/databao_cli/features/ui/models/chat_session.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from databao.agent.core.thread import Thread - from databao_cli.ui.streaming import StreamingWriter + from databao_cli.features.ui.streaming import StreamingWriter __all__ = ["ChatMessage", "ChatSession"] diff --git a/src/databao_cli/ui/models/settings.py b/src/databao_cli/features/ui/models/settings.py similarity index 100% rename from src/databao_cli/ui/models/settings.py rename to src/databao_cli/features/ui/models/settings.py diff --git a/src/databao_cli/ui/pages/__init__.py b/src/databao_cli/features/ui/pages/__init__.py similarity index 100% rename from src/databao_cli/ui/pages/__init__.py rename to src/databao_cli/features/ui/pages/__init__.py diff --git a/src/databao_cli/ui/pages/agent_settings.py b/src/databao_cli/features/ui/pages/agent_settings.py similarity index 95% rename from src/databao_cli/ui/pages/agent_settings.py rename to src/databao_cli/features/ui/pages/agent_settings.py index c48d4917..05fb0a59 100644 --- a/src/databao_cli/ui/pages/agent_settings.py +++ b/src/databao_cli/features/ui/pages/agent_settings.py @@ -6,10 +6,10 @@ import streamlit as st -from databao_cli.executor_utils import EXECUTOR_TYPES, LLM_PROVIDER_MODELS, LLM_PROVIDERS -from databao_cli.ui.app import _clear_all_chat_threads -from databao_cli.ui.components.status import AppStatus, set_status -from databao_cli.ui.models.settings import _ENV_VAR_MAP, LLMProviderConfig, LLMSettings +from databao_cli.features.ui.app import _clear_all_chat_threads +from databao_cli.features.ui.components.status import AppStatus, set_status +from databao_cli.features.ui.models.settings import _ENV_VAR_MAP, LLMProviderConfig, LLMSettings +from databao_cli.shared.executor_utils import EXECUTOR_TYPES, LLM_PROVIDER_MODELS, LLM_PROVIDERS def render_agent_settings_page(*, auto_apply: bool = False) -> None: @@ -191,7 +191,7 @@ def _persist_current_settings() -> None: from setup wizard to normal app mode (where _load_persisted_state reads from disk). """ - from databao_cli.ui.services.settings_persistence import get_or_create_settings, save_settings + from databao_cli.features.ui.services.settings_persistence import get_or_create_settings, save_settings settings = get_or_create_settings() settings.agent.executor_type = st.session_state.get("executor_type", "claude_code") @@ -205,7 +205,7 @@ def _render_test_connection(provider_type: str, api_key: str, base_url: str) -> cache_key = f"_llm_test_{provider_type}" if st.button("🔗 Test connection", type="tertiary", key=f"test_conn_{provider_type}"): - from databao_cli.ui.services.llm_models import fetch_models + from databao_cli.features.ui.services.llm_models import fetch_models try: models = fetch_models(provider_type, api_key, base_url) @@ -266,7 +266,7 @@ def _model_selectbox(provider_type: str, models: list[str], current_model: str) elif current_model: default_index = _find_closest_model_index(models, current_model, provider_type) else: - from databao_cli.ui.services.llm_models import pick_default_model + from databao_cli.features.ui.services.llm_models import pick_default_model best = pick_default_model(models, provider_type) default_index = models.index(best) if best in models else 0 diff --git a/src/databao_cli/ui/pages/chat.py b/src/databao_cli/features/ui/pages/chat.py similarity index 90% rename from src/databao_cli/ui/pages/chat.py rename to src/databao_cli/features/ui/pages/chat.py index 3fb87344..dbae047b 100644 --- a/src/databao_cli/ui/pages/chat.py +++ b/src/databao_cli/features/ui/pages/chat.py @@ -7,15 +7,15 @@ from databao.agent.core.agent import Agent from databao.agent.core.thread import Thread -from databao_cli.project.layout import ProjectLayout -from databao_cli.ui.app import _clear_all_chat_threads -from databao_cli.ui.components.chat import render_chat_interface -from databao_cli.ui.components.sidebar import render_sidebar_chat_content -from databao_cli.ui.components.status import AppStatus, set_status -from databao_cli.ui.models.chat_session import ChatSession -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status -from databao_cli.ui.services.chat_persistence import save_chat -from databao_cli.ui.services.chat_title import check_title_completion, trigger_title_generation +from databao_cli.features.ui.app import _clear_all_chat_threads +from databao_cli.features.ui.components.chat import render_chat_interface +from databao_cli.features.ui.components.sidebar import render_sidebar_chat_content +from databao_cli.features.ui.components.status import AppStatus, set_status +from databao_cli.features.ui.models.chat_session import ChatSession +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.features.ui.services.chat_persistence import save_chat +from databao_cli.features.ui.services.chat_title import check_title_completion, trigger_title_generation +from databao_cli.shared.project.layout import ProjectLayout logger = logging.getLogger(__name__) @@ -136,7 +136,7 @@ def _get_or_create_thread_for_chat(chat: ChatSession, agent: Agent) -> bool: Returns True if thread is available, False on error. """ - from databao_cli.ui.streaming import StreamingWriter + from databao_cli.features.ui.streaming import StreamingWriter if chat.writer is None: chat.writer = StreamingWriter() diff --git a/src/databao_cli/ui/pages/context_settings.py b/src/databao_cli/features/ui/pages/context_settings.py similarity index 89% rename from src/databao_cli/ui/pages/context_settings.py rename to src/databao_cli/features/ui/pages/context_settings.py index 1a899156..dd3c7fc6 100644 --- a/src/databao_cli/ui/pages/context_settings.py +++ b/src/databao_cli/features/ui/pages/context_settings.py @@ -5,13 +5,13 @@ import streamlit as st from databao.agent import Agent -from databao_cli.project.layout import ProjectLayout -from databao_cli.ui.app import invalidate_agent, is_read_only_domain -from databao_cli.ui.components.datasource_manager import render_datasource_manager -from databao_cli.ui.components.icons import get_db_type_and_icon -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status -from databao_cli.ui.services.build_service import render_build_section -from databao_cli.ui.services.dce_operations import get_status_info +from databao_cli.features.ui.app import invalidate_agent, is_read_only_domain +from databao_cli.features.ui.components.datasource_manager import render_datasource_manager +from databao_cli.features.ui.components.icons import get_db_type_and_icon +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.features.ui.services.build_service import render_build_section +from databao_cli.features.ui.services.dce_operations import get_status_info +from databao_cli.shared.project.layout import ProjectLayout logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/pages/general_settings.py b/src/databao_cli/features/ui/pages/general_settings.py similarity index 91% rename from src/databao_cli/ui/pages/general_settings.py rename to src/databao_cli/features/ui/pages/general_settings.py index 29d22a2c..3edc3ef4 100644 --- a/src/databao_cli/ui/pages/general_settings.py +++ b/src/databao_cli/features/ui/pages/general_settings.py @@ -4,10 +4,10 @@ import streamlit as st -from databao_cli.ui.app import _clear_all_chat_threads -from databao_cli.ui.services.chat_persistence import delete_all_chats -from databao_cli.ui.services.settings_persistence import delete_settings -from databao_cli.ui.services.storage import get_cache_dir, get_chats_dir, get_storage_base_path +from databao_cli.features.ui.app import _clear_all_chat_threads +from databao_cli.features.ui.services.chat_persistence import delete_all_chats +from databao_cli.features.ui.services.settings_persistence import delete_settings +from databao_cli.features.ui.services.storage import get_cache_dir, get_chats_dir, get_storage_base_path logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/pages/welcome.py b/src/databao_cli/features/ui/pages/welcome.py similarity index 93% rename from src/databao_cli/ui/pages/welcome.py rename to src/databao_cli/features/ui/pages/welcome.py index bebc8384..547bd09d 100644 --- a/src/databao_cli/ui/pages/welcome.py +++ b/src/databao_cli/features/ui/pages/welcome.py @@ -6,9 +6,9 @@ import streamlit as st -from databao_cli.project.layout import find_project -from databao_cli.ui.models.settings import LLMSettings -from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.features.ui.models.settings import LLMSettings +from databao_cli.features.ui.project_utils import DatabaoProjectStatus, databao_project_status +from databao_cli.shared.project.layout import find_project logger = logging.getLogger(__name__) @@ -20,7 +20,7 @@ def render_welcome_page() -> None: """Render the welcome/home page for a fully configured project.""" - from databao_cli.ui.app import _create_new_chat + from databao_cli.features.ui.app import _create_new_chat _col1, col2, _col3 = st.columns([1, 2, 1]) @@ -133,13 +133,13 @@ def render_setup_wizard_page() -> None: When read-only-domain mode is active, editing sections are disabled with an explanation banner. """ - from databao_cli.ui.app import _create_new_chat, is_hide_build_context_hint, is_read_only_domain - from databao_cli.ui.components.datasource_manager import render_datasource_manager - from databao_cli.ui.services.build_service import ( + from databao_cli.features.ui.app import _create_new_chat, is_hide_build_context_hint, is_read_only_domain + from databao_cli.features.ui.components.datasource_manager import render_datasource_manager + from databao_cli.features.ui.services.build_service import ( get_build_status, render_build_section, ) - from databao_cli.ui.services.dce_operations import init_project, list_datasources + from databao_cli.features.ui.services.dce_operations import init_project, list_datasources project_dir: Path = st.session_state.get("_project_dir", Path.cwd()) read_only = is_read_only_domain() @@ -249,7 +249,7 @@ def render_setup_wizard_page() -> None: "Configure the execution engine and language model for the AI agent. " "You'll need an API key for your chosen LLM provider." ) - from databao_cli.ui.pages.agent_settings import render_agent_settings_page + from databao_cli.features.ui.pages.agent_settings import render_agent_settings_page render_agent_settings_page(auto_apply=True) @@ -304,7 +304,7 @@ def render_setup_wizard_page() -> None: st.markdown("All configured! You're ready to start using Databao.") if st.button("Start New Chat", key="setup_start_chat", type="primary"): - from databao_cli.ui.app import mark_welcome_completed + from databao_cli.features.ui.app import mark_welcome_completed mark_welcome_completed() _create_new_chat() diff --git a/src/databao_cli/ui/project_utils.py b/src/databao_cli/features/ui/project_utils.py similarity index 96% rename from src/databao_cli/ui/project_utils.py rename to src/databao_cli/features/ui/project_utils.py index 47c3cabe..f16fa067 100644 --- a/src/databao_cli/ui/project_utils.py +++ b/src/databao_cli/features/ui/project_utils.py @@ -3,7 +3,7 @@ from databao.agent.integrations.dce import DatabaoContextApi -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.project.layout import ProjectLayout logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/services/__init__.py b/src/databao_cli/features/ui/services/__init__.py similarity index 78% rename from src/databao_cli/ui/services/__init__.py rename to src/databao_cli/features/ui/services/__init__.py index 85a26fbf..8a7a5c22 100644 --- a/src/databao_cli/ui/services/__init__.py +++ b/src/databao_cli/features/ui/services/__init__.py @@ -1,6 +1,6 @@ """Services for the Databao Streamlit app.""" -from databao_cli.ui.services.chat_persistence import ( +from databao_cli.features.ui.services.chat_persistence import ( delete_all_chats, delete_chat, load_all_chats, @@ -8,11 +8,11 @@ save_chat, save_current_chat, ) -from databao_cli.ui.services.chat_title import ( +from databao_cli.features.ui.services.chat_title import ( check_title_completion, trigger_title_generation, ) -from databao_cli.ui.services.query_executor import ( +from databao_cli.features.ui.services.query_executor import ( QueryResult, check_query_completion, get_query_phase, @@ -20,13 +20,13 @@ start_query_execution, stop_query, ) -from databao_cli.ui.services.settings_persistence import ( +from databao_cli.features.ui.services.settings_persistence import ( delete_settings, get_or_create_settings, load_settings, save_settings, ) -from databao_cli.ui.services.storage import ( +from databao_cli.features.ui.services.storage import ( get_cache_dir, get_chat_dir, get_chats_dir, diff --git a/src/databao_cli/ui/services/build_service.py b/src/databao_cli/features/ui/services/build_service.py similarity index 99% rename from src/databao_cli/ui/services/build_service.py rename to src/databao_cli/features/ui/services/build_service.py index a94d8902..b6c6c345 100644 --- a/src/databao_cli/ui/services/build_service.py +++ b/src/databao_cli/features/ui/services/build_service.py @@ -16,7 +16,7 @@ import streamlit as st from databao_context_engine import BuildDatasourceResult -from databao_cli.ui.services.dce_operations import build_context +from databao_cli.features.ui.services.dce_operations import build_context logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/services/chat_persistence.py b/src/databao_cli/features/ui/services/chat_persistence.py similarity index 97% rename from src/databao_cli/ui/services/chat_persistence.py rename to src/databao_cli/features/ui/services/chat_persistence.py index d410019e..54884c56 100644 --- a/src/databao_cli/ui/services/chat_persistence.py +++ b/src/databao_cli/features/ui/services/chat_persistence.py @@ -11,10 +11,10 @@ import pandas as pd from databao.agent import ExecutionResult -from databao_cli.ui.services.storage import get_cache_dir, get_chat_dir, get_chats_dir, is_valid_chat_id +from databao_cli.features.ui.services.storage import get_cache_dir, get_chat_dir, get_chats_dir, is_valid_chat_id if TYPE_CHECKING: - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession logger = logging.getLogger(__name__) @@ -122,7 +122,7 @@ def load_chat(chat_id: str) -> ChatSession | None: logger.warning(f"Refusing to load chat with invalid id: {chat_id!r}") return None - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession chats_dir = get_chats_dir() chat_dir = chats_dir / chat_id diff --git a/src/databao_cli/ui/services/chat_title.py b/src/databao_cli/features/ui/services/chat_title.py similarity index 98% rename from src/databao_cli/ui/services/chat_title.py rename to src/databao_cli/features/ui/services/chat_title.py index 1c420472..f7461598 100644 --- a/src/databao_cli/ui/services/chat_title.py +++ b/src/databao_cli/features/ui/services/chat_title.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: from databao.agent.core.agent import Agent - from databao_cli.ui.models.chat_session import ChatSession + from databao_cli.features.ui.models.chat_session import ChatSession logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/services/dce_operations.py b/src/databao_cli/features/ui/services/dce_operations.py similarity index 97% rename from src/databao_cli/ui/services/dce_operations.py rename to src/databao_cli/features/ui/services/dce_operations.py index c62588f7..17a436b2 100644 --- a/src/databao_cli/ui/services/dce_operations.py +++ b/src/databao_cli/features/ui/services/dce_operations.py @@ -23,8 +23,8 @@ from databao_context_engine.pluginlib.config import ConfigPropertyDefinition from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource -from databao_cli.commands.init import init_impl -from databao_cli.project.layout import ProjectLayout +from databao_cli.features.init.service import init_impl +from databao_cli.shared.project.layout import ProjectLayout logger = logging.getLogger(__name__) @@ -173,6 +173,6 @@ def get_status_info(project_dir: Path) -> str: Delegates to the CLI's status_impl which handles multi-domain iteration. """ - from databao_cli.commands.status import status_impl + from databao_cli.features.status import status_impl return status_impl(project_dir) diff --git a/src/databao_cli/ui/services/llm_models.py b/src/databao_cli/features/ui/services/llm_models.py similarity index 97% rename from src/databao_cli/ui/services/llm_models.py rename to src/databao_cli/features/ui/services/llm_models.py index ea52fd17..e6aa5d7f 100644 --- a/src/databao_cli/ui/services/llm_models.py +++ b/src/databao_cli/features/ui/services/llm_models.py @@ -70,7 +70,7 @@ def _fetch_anthropic_models(api_key: str) -> list[str]: def pick_default_model(models: list[str], provider_type: str) -> str: """Pick the best default from a fetched model list.""" - from databao_cli.executor_utils import LLM_PROVIDER_MODELS + from databao_cli.shared.executor_utils import LLM_PROVIDER_MODELS preferred = LLM_PROVIDER_MODELS.get(provider_type, []) diff --git a/src/databao_cli/ui/services/query_executor.py b/src/databao_cli/features/ui/services/query_executor.py similarity index 97% rename from src/databao_cli/ui/services/query_executor.py rename to src/databao_cli/features/ui/services/query_executor.py index 867d4d6a..58ad0e2b 100644 --- a/src/databao_cli/ui/services/query_executor.py +++ b/src/databao_cli/features/ui/services/query_executor.py @@ -12,13 +12,13 @@ from streamlit.runtime.scriptrunner import add_script_run_ctx, get_script_run_ctx -from databao_cli.ui.components.results import _extract_visualization_data +from databao_cli.features.ui.components.results import _extract_visualization_data if TYPE_CHECKING: from databao.agent.core.thread import Thread - from databao_cli.ui.models.chat_session import ChatSession - from databao_cli.ui.streaming import StreamingWriter + from databao_cli.features.ui.models.chat_session import ChatSession + from databao_cli.features.ui.streaming import StreamingWriter logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/services/settings_persistence.py b/src/databao_cli/features/ui/services/settings_persistence.py similarity index 93% rename from src/databao_cli/ui/services/settings_persistence.py rename to src/databao_cli/features/ui/services/settings_persistence.py index e24122be..4c028e0f 100644 --- a/src/databao_cli/ui/services/settings_persistence.py +++ b/src/databao_cli/features/ui/services/settings_persistence.py @@ -2,8 +2,8 @@ import logging -from databao_cli.ui.models.settings import Settings -from databao_cli.ui.services.storage import get_settings_path +from databao_cli.features.ui.models.settings import Settings +from databao_cli.features.ui.services.storage import get_settings_path logger = logging.getLogger(__name__) diff --git a/src/databao_cli/ui/services/storage.py b/src/databao_cli/features/ui/services/storage.py similarity index 100% rename from src/databao_cli/ui/services/storage.py rename to src/databao_cli/features/ui/services/storage.py diff --git a/src/databao_cli/ui/streaming.py b/src/databao_cli/features/ui/streaming.py similarity index 100% rename from src/databao_cli/ui/streaming.py rename to src/databao_cli/features/ui/streaming.py diff --git a/src/databao_cli/ui/suggestions.py b/src/databao_cli/features/ui/suggestions.py similarity index 100% rename from src/databao_cli/ui/suggestions.py rename to src/databao_cli/features/ui/suggestions.py diff --git a/src/databao_cli/mcp/tools/databao_ask/__init__.py b/src/databao_cli/mcp/tools/databao_ask/__init__.py deleted file mode 100644 index eef81e8a..00000000 --- a/src/databao_cli/mcp/tools/databao_ask/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""databao_ask MCP tool package.""" - -from databao_cli.mcp.tools.databao_ask.tool import register - -__all__ = ["register"] diff --git a/src/databao_cli/shared/__init__.py b/src/databao_cli/shared/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/shared/cli_utils.py b/src/databao_cli/shared/cli_utils.py new file mode 100644 index 00000000..21a439fe --- /dev/null +++ b/src/databao_cli/shared/cli_utils.py @@ -0,0 +1,34 @@ +"""Shared CLI utilities for command implementations.""" + +import functools +import sys +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import click + +from databao_cli.shared.errors import FeatureError +from databao_cli.shared.project.layout import ProjectLayout, find_project + + +def get_project_or_raise(project_dir: Path) -> ProjectLayout: + """Return the project layout or raise FeatureError if no project is found.""" + project_layout = find_project(project_dir) + if not project_layout: + raise FeatureError("No project found.") + return project_layout + + +def handle_feature_errors(f: Callable[..., Any]) -> Callable[..., Any]: + """Decorator that catches FeatureError and converts it to a CLI error exit.""" + + @functools.wraps(f) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return f(*args, **kwargs) + except FeatureError as e: + click.echo(e.message, err=True) + sys.exit(e.exit_code) + + return wrapper diff --git a/src/databao_cli/commands/context_engine_cli.py b/src/databao_cli/shared/context_engine_cli.py similarity index 100% rename from src/databao_cli/commands/context_engine_cli.py rename to src/databao_cli/shared/context_engine_cli.py diff --git a/src/databao_cli/shared/errors.py b/src/databao_cli/shared/errors.py new file mode 100644 index 00000000..f734fc2f --- /dev/null +++ b/src/databao_cli/shared/errors.py @@ -0,0 +1,7 @@ +class FeatureError(Exception): + """Raised by feature functions to signal a user-facing error and request CLI exit.""" + + def __init__(self, message: str, exit_code: int = 1) -> None: + super().__init__(message) + self.message = message + self.exit_code = exit_code diff --git a/src/databao_cli/executor_utils.py b/src/databao_cli/shared/executor_utils.py similarity index 100% rename from src/databao_cli/executor_utils.py rename to src/databao_cli/shared/executor_utils.py diff --git a/src/databao_cli/shared/log/__init__.py b/src/databao_cli/shared/log/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/log/llm_errors.py b/src/databao_cli/shared/log/llm_errors.py similarity index 100% rename from src/databao_cli/log/llm_errors.py rename to src/databao_cli/shared/log/llm_errors.py diff --git a/src/databao_cli/log/logging.py b/src/databao_cli/shared/log/logging.py similarity index 97% rename from src/databao_cli/log/logging.py rename to src/databao_cli/shared/log/logging.py index 37dcbda7..6911d046 100644 --- a/src/databao_cli/log/logging.py +++ b/src/databao_cli/shared/log/logging.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Any -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.project.layout import ProjectLayout def configure_logging(project_layout: ProjectLayout | None, verbose: bool = False, quiet: bool = False) -> None: diff --git a/src/databao_cli/shared/project/__init__.py b/src/databao_cli/shared/project/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/project/layout.py b/src/databao_cli/shared/project/layout.py similarity index 100% rename from src/databao_cli/project/layout.py rename to src/databao_cli/shared/project/layout.py diff --git a/src/databao_cli/ui/models/__init__.py b/src/databao_cli/ui/models/__init__.py deleted file mode 100644 index 61e86cab..00000000 --- a/src/databao_cli/ui/models/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Data models for the Databao Streamlit app.""" - -from databao_cli.ui.models.chat_session import ChatMessage, ChatSession -from databao_cli.ui.models.settings import AgentSettings, Settings - -__all__ = [ - "AgentSettings", - "ChatMessage", - "ChatSession", - "Settings", -] diff --git a/src/databao_cli/workflows/__init__.py b/src/databao_cli/workflows/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/workflows/ask.py b/src/databao_cli/workflows/ask.py new file mode 100644 index 00000000..50cbc31e --- /dev/null +++ b/src/databao_cli/workflows/ask.py @@ -0,0 +1,119 @@ +"""Interactive CLI workflows for the ask command: REPL, one-shot, and result display.""" + +import click +from databao.agent import Agent +from databao.agent.core.thread import Thread + +from databao_cli.features.ask.display import DEFAULT_MAX_DISPLAY_ROWS, dataframe_to_prettytable +from databao_cli.features.ui.streaming import StreamingWriter +from databao_cli.shared.errors import FeatureError +from databao_cli.shared.log.llm_errors import format_llm_error + + +def _create_cli_writer() -> StreamingWriter: + return StreamingWriter(on_write=lambda text: click.echo(text, nl=False)) + + +def _print_help() -> None: + click.echo("Databao REPL") + click.echo("Ask questions about your data in natural language.\n") + click.echo("Commands:") + click.echo(" \\help - Show this help") + click.echo(" \\clear - Start a new conversation") + click.echo(" \\q - Exit\n") + + +def display_result(thread: Thread) -> None: + """Display the execution result from a thread to the CLI.""" + text = thread.text() + if text: + click.echo(text) + + code = thread.code() + if code: + click.echo(f"\n```sql\n{code}\n```") + + df = thread.df() + if df is not None: + rows_shown = min(DEFAULT_MAX_DISPLAY_ROWS, len(df)) + click.echo(f"\n[DataFrame: {rows_shown} / {len(df)} rows]") + click.echo(dataframe_to_prettytable(df)) + + +def run_interactive_mode(agent: Agent, show_thinking: bool) -> None: + """Run the interactive REPL mode.""" + click.echo("\nDatabao REPL") + click.echo("\nType \\help for available commands.\n") + + writer = _create_cli_writer() if show_thinking else None + + thread = agent.thread( + stream_ask=show_thinking, + writer=writer, + ) + + while True: + try: + user_input = click.prompt("You", prompt_suffix="> ") + except (EOFError, KeyboardInterrupt): + click.echo() + break + + user_input = user_input.strip() + if not user_input: + continue + + if user_input.startswith("\\"): + command = user_input[1:].lower() + + if command in ("exit", "quit", "q"): + break + + if command == "clear": + if writer: + writer.clear() + thread = agent.thread( + stream_ask=show_thinking, + writer=writer, + ) + click.echo("Conversation cleared.\n") + continue + + if command == "help": + _print_help() + continue + + click.echo(f"Unknown command: {user_input}. Type \\help for available commands.\n") + continue + + try: + thread.ask(user_input, stream=show_thinking) + + if writer: + writer.clear() + + click.echo("\nAssistant:") + display_result(thread) + click.echo() + + except Exception as e: + if writer: + writer.clear() + click.echo(f"\nError: {format_llm_error(e)}\n", err=True) + + +def run_one_shot_mode(agent: Agent, question: str, show_thinking: bool) -> None: + """Run a single question and exit.""" + writer = _create_cli_writer() if show_thinking else None + + thread = agent.thread( + stream_ask=show_thinking, + writer=writer, + ) + + try: + thread.ask(question, stream=show_thinking) + display_result(thread) + + except Exception as e: + raise FeatureError(f"Error: {format_llm_error(e)}") from e diff --git a/src/databao_cli/workflows/datasource/__init__.py b/src/databao_cli/workflows/datasource/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/databao_cli/workflows/datasource/add.py b/src/databao_cli/workflows/datasource/add.py new file mode 100644 index 00000000..6517998c --- /dev/null +++ b/src/databao_cli/workflows/datasource/add.py @@ -0,0 +1,65 @@ +"""Interactive CLI workflow for adding a datasource.""" + +import os + +import click +from databao_context_engine import ( + DatabaoContextDomainManager, + DatabaoContextPluginLoader, + DatasourceType, +) + +from databao_cli.shared.context_engine_cli import ClickUserInputCallback +from databao_cli.shared.project.layout import ProjectLayout +from databao_cli.workflows.datasource.check import print_connection_check_results + + +def add_workflow(project_layout: ProjectLayout, domain: str) -> None: + """Interactive wizard: collect inputs, create config, optionally check connection.""" + from databao_cli.features.datasource.add import create_datasource_config, datasource_config_exists + + domain_dir = project_layout.domains_dir / domain + plugin_loader = DatabaoContextPluginLoader() + + click.echo(f"We will guide you to add a new datasource into {domain} domain, at {domain_dir.resolve()}") + + datasource_type = _ask_for_datasource_type(plugin_loader.get_all_supported_datasource_types(exclude_file_plugins=True)) + datasource_name = click.prompt("Datasource name?", type=str) + + overwrite_existing = False + existing_id = datasource_config_exists(project_layout, domain, datasource_name) + if existing_id is not None: + click.confirm( + f"A config file already exists for this datasource {existing_id.relative_path_to_config_file()}. " + f"Do you want to overwrite it?", + abort=True, + default=False, + ) + overwrite_existing = True + + datasource_id, config_path = create_datasource_config( + project_layout, + domain, + datasource_type, + datasource_name, + ClickUserInputCallback(), + overwrite_existing=overwrite_existing, + ) + + click.echo(f"{os.linesep}We've created a new config file for your datasource at: {config_path}") + + if click.confirm("\nDo you want to check the connection to this new datasource?"): + domain_manager = DatabaoContextDomainManager(domain_dir=domain_dir) + results = domain_manager.check_datasource_connection(datasource_ids=[datasource_id]) + print_connection_check_results(domain, results) + + +def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType: + all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types]) + config_type = click.prompt( + "What type of datasource do you want to add?", + type=click.Choice(all_datasource_types), + default=all_datasource_types[0] if len(all_datasource_types) == 1 else None, + ) + click.echo(f"Selected type: {config_type}") + return DatasourceType(full_type=config_type) diff --git a/src/databao_cli/workflows/datasource/check.py b/src/databao_cli/workflows/datasource/check.py new file mode 100644 index 00000000..9dd2e0f2 --- /dev/null +++ b/src/databao_cli/workflows/datasource/check.py @@ -0,0 +1,20 @@ +"""CLI display for datasource connection check results.""" + +import os + +import click +from databao_context_engine import CheckDatasourceConnectionResult, DatasourceId + + +def print_connection_check_results( + domain: str, datasource_results: dict[DatasourceId, CheckDatasourceConnectionResult] +) -> None: + for result in datasource_results.values(): + fq_datasource_name = domain + os.pathsep + str(result.datasource_id) + status = str(result.connection_status.value) + if result.summary: + status += f" - {result.summary}" + if result.full_message: + status += f": {result.full_message}" + + click.echo(f"{fq_datasource_name}: {status}") diff --git a/tests/conftest.py b/tests/conftest.py index a5b58a11..445bbd70 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,8 +2,8 @@ import pytest -from databao_cli.commands.init import init_impl -from databao_cli.project.layout import ProjectLayout +from databao_cli.features.init.service import init_impl +from databao_cli.shared.project.layout import ProjectLayout @pytest.fixture diff --git a/tests/test_add_datasource.py b/tests/test_add_datasource.py index 6265fe3a..49b454b8 100644 --- a/tests/test_add_datasource.py +++ b/tests/test_add_datasource.py @@ -5,7 +5,7 @@ from click.testing import CliRunner from databao_cli.__main__ import cli -from databao_cli.commands.init import init_impl as init_databao_project +from databao_cli.features.init.service import init_impl as init_databao_project @pytest.fixture diff --git a/tests/test_app.py b/tests/test_app.py index 4a6dbe0d..76871071 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -19,7 +19,7 @@ def test_app_hide_flags_forwarded() -> None: runner = CliRunner() mock_bootstrap = MagicMock() - with patch("databao_cli.commands.app.bootstrap_streamlit_app", mock_bootstrap): + with patch("databao_cli.features.ui.cli.bootstrap_streamlit_app", mock_bootstrap): result = runner.invoke( cli, ["app", "--hide-suggested-questions", "--hide-build-context-hint"], @@ -36,7 +36,7 @@ def test_app_hide_flags_default_false() -> None: runner = CliRunner() mock_bootstrap = MagicMock() - with patch("databao_cli.commands.app.bootstrap_streamlit_app", mock_bootstrap): + with patch("databao_cli.features.ui.cli.bootstrap_streamlit_app", mock_bootstrap): result = runner.invoke(cli, ["app"]) assert result.exit_code == 0 diff --git a/tests/test_build.py b/tests/test_build.py index 3067e7b5..e1dec87d 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -2,7 +2,7 @@ from _pytest.tmpdir import TempPathFactory from databao_context_engine import DatabaoContextDomainManager, DatasourceType -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.project.layout import ProjectLayout from tests.utils.project import describe_result, run_build diff --git a/tests/test_index.py b/tests/test_index.py index 2d2a3310..b72995e2 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -2,7 +2,7 @@ from _pytest.tmpdir import TempPathFactory from databao_context_engine import DatabaoContextDomainManager, DatasourceType -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.project.layout import ProjectLayout from tests.utils.project import describe_result, run_build, run_index diff --git a/tests/test_query_executor.py b/tests/test_query_executor.py index 83424a8a..4a017dc9 100644 --- a/tests/test_query_executor.py +++ b/tests/test_query_executor.py @@ -6,8 +6,8 @@ import pytest -from databao_cli.ui.models.chat_session import ChatSession -from databao_cli.ui.services.query_executor import ( +from databao_cli.features.ui.models.chat_session import ChatSession +from databao_cli.features.ui.services.query_executor import ( QueryResult, QueryThread, _raise_in_thread, @@ -118,7 +118,7 @@ def test_already_running(self, chat_session: ChatSession, mock_thread: MagicMock result = start_query_execution(chat_session, mock_thread, "test") assert result is False - @patch("databao_cli.ui.services.query_executor.get_script_run_ctx") + @patch("databao_cli.features.ui.services.query_executor.get_script_run_ctx") def test_starts_execution( self, mock_ctx: MagicMock, diff --git a/tests/test_query_executor_race_conditions.py b/tests/test_query_executor_race_conditions.py index d87c2639..3750126d 100644 --- a/tests/test_query_executor_race_conditions.py +++ b/tests/test_query_executor_race_conditions.py @@ -14,8 +14,8 @@ import pytest -from databao_cli.ui.models.chat_session import ChatMessage, ChatSession -from databao_cli.ui.services.query_executor import ( +from databao_cli.features.ui.models.chat_session import ChatMessage, ChatSession +from databao_cli.features.ui.services.query_executor import ( QueryResult, QueryThread, check_query_completion, @@ -225,7 +225,7 @@ def work_with_phase_change() -> None: class TestRapidSequentialOperations: """Tests for rapid sequential operations.""" - @patch("databao_cli.ui.services.query_executor.get_script_run_ctx") + @patch("databao_cli.features.ui.services.query_executor.get_script_run_ctx") def test_rapid_stop_start_stop(self, mock_ctx: MagicMock, chat_session: ChatSession, mock_writer: MagicMock) -> None: """Quick stop→start→stop sequence should work correctly.""" mock_ctx.return_value = None @@ -268,7 +268,7 @@ def slow_work() -> None: stop_query(chat_session) assert chat_session.query_status == "idle" - @patch("databao_cli.ui.services.query_executor.get_script_run_ctx") + @patch("databao_cli.features.ui.services.query_executor.get_script_run_ctx") def test_start_immediately_after_stop( self, mock_ctx: MagicMock, chat_session: ChatSession, mock_writer: MagicMock ) -> None: From 2010c6b9d6057abd5fb7e51852416b9333a093b4 Mon Sep 17 00:00:00 2001 From: Andrei Gasparian Date: Mon, 23 Mar 2026 19:17:28 +0100 Subject: [PATCH 7/8] [DBA-290] Fix flaky duckdb test by using explicit connection Both test_build and test_index called duckdb.execute() on the shared default in-memory connection instead of the file-scoped connection, causing "table t1 already exists" when both tests ran in the same suite. Co-Authored-By: Claude Opus 4.6 --- tests/test_build.py | 5 +++-- tests/test_index.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_build.py b/tests/test_build.py index e1dec87d..217bf904 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -8,8 +8,9 @@ def test_databao_build_duckdb_datasource(project_layout: ProjectLayout, tmp_path_factory: TempPathFactory) -> None: test_db = tmp_path_factory.mktemp("duckdb") / "test_db.duckdb" - duckdb.connect(str(test_db)) - duckdb.execute("CREATE TABLE t1 AS SELECT 1 AS i, 2 AS j;") + conn = duckdb.connect(str(test_db)) + conn.execute("CREATE TABLE t1 AS SELECT 1 AS i, 2 AS j;") + conn.close() dce_domain_manager = DatabaoContextDomainManager(domain_dir=project_layout.root_domain_dir) dce_domain_manager.create_datasource_config( diff --git a/tests/test_index.py b/tests/test_index.py index b72995e2..c170bb4e 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -8,8 +8,9 @@ def test_databao_index_duckdb_datasource(project_layout: ProjectLayout, tmp_path_factory: TempPathFactory) -> None: test_db = tmp_path_factory.mktemp("duckdb") / "test_db.duckdb" - duckdb.connect(str(test_db)) - duckdb.execute("CREATE TABLE t1 AS SELECT 1 AS i, 2 AS j;") + conn = duckdb.connect(str(test_db)) + conn.execute("CREATE TABLE t1 AS SELECT 1 AS i, 2 AS j;") + conn.close() dce_domain_manager = DatabaoContextDomainManager(domain_dir=project_layout.root_domain_dir) dce_domain_manager.create_datasource_config( From 7c7ea5d9bb1f9375cb8b84575a43f63a4f4be380 Mon Sep 17 00:00:00 2001 From: Lenar Sharipov Date: Thu, 26 Mar 2026 01:17:47 +0100 Subject: [PATCH 8/8] DBA-295 Update SKILL.md to be consistent with architecture.md --- .claude/skills/review-architecture/SKILL.md | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/.claude/skills/review-architecture/SKILL.md b/.claude/skills/review-architecture/SKILL.md index b9320b9e..bea305af 100755 --- a/.claude/skills/review-architecture/SKILL.md +++ b/.claude/skills/review-architecture/SKILL.md @@ -18,13 +18,6 @@ Review in this order: 4. `CLAUDE.md` 5. `README.md` (CLI usage and user-facing workflows) -Then validate actual implementation under: - -- `src/databao_cli/commands/` -- `src/databao_cli/ui/` -- `src/databao_cli/mcp/` -- `src/databao_cli/project/` - ## Review goals - Confirm boundaries are clear and responsibilities are separated. @@ -37,7 +30,10 @@ Then validate actual implementation under: - Are modules aligned with single responsibility? - Are CLI concerns separated from business logic? - Is the Click command structure clean and discoverable? -- Are MCP tools properly isolated in `mcp/tools/`? +- Does `workflows/` stay free of business logic (delegates to `features/`)? +- Are `features/` functions free of Click dependency (pure business operations)? +- Is `shared/` limited to cross-feature utilities with no business logic of its own? +- Are MCP tools properly isolated in their own module? - Are UI components reusable and page-specific logic separated? - Are errors actionable and surfaced at the right layer?