diff --git a/docs/architecture.md b/docs/architecture.md index 0a7a8686..3c98726d 100755 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,68 +1,115 @@ # 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: - -- 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/`) - -## Entry Point - -`databao_cli.__main__:cli` — a Click group that registers all subcommands. - -## CLI Commands - -| 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) - -`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 +Project structure: +``` +src/databao_cli/ +- 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 - - `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. +## Entry Point +`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`. + +## Layer Responsibilities + +### `commands/` — Routing Layer +Contains only Click wiring: decorators, option/argument definitions, and calls +to workflow or feature functions. No business logic lives here. + +- 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. + +``` +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 +``` + +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. + +``` +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/`, register with Click group - in `__main__.py`. -- Add MCP tool: add handler in `mcp/tools/`, register in `mcp/server.py`. +- Add CLI command: create module in `commands/`, then add it to the + `COMMANDS` collection in `src/databao_cli/__main__.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 fe6a92f7..55b7378e 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.log.logging import configure_logging -from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project +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.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] @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/app.py b/src/databao_cli/commands/app.py index 00198fb8..617ee8ba 100644 --- a/src/databao_cli/commands/app.py +++ b/src/databao_cli/commands/app.py @@ -1,26 +1,63 @@ """databao app command - Launch the Databao Streamlit web interface.""" -import subprocess -import sys - import click -from databao_cli.ui.cli import bootstrap_streamlit_app +from databao_cli.shared.cli_utils import handle_feature_errors -def app_impl(ctx: click.Context) -> None: - click.echo("Starting Databao UI...") +@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 +@handle_feature_errors +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.features.ui.cli import app_impl + + click.echo("Starting Databao UI...") try: - bootstrap_streamlit_app( + app_impl( 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), + 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 e72404c8..1fb34ca4 100644 --- a/src/databao_cli/commands/ask.py +++ b/src/databao_cli/commands/ask.py @@ -1,206 +1,40 @@ """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 - - -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( +from databao_cli.shared.cli_utils import handle_feature_errors + + +@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 +@handle_feature_errors +def ask( ctx: click.Context, question: str | None, one_shot: bool, @@ -208,20 +42,25 @@ def ask_impl( 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) + """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. - # Get project path from CLI context - project_path: Path = ctx.obj["project_dir"] + \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.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 bf4b00df..28499c9b 100644 --- a/src/databao_cli/commands/build.py +++ b/src/databao_cli/commands/build.py @@ -1,12 +1,34 @@ -from databao_context_engine import BuildDatasourceResult, DatabaoContextDomainManager +import click -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import handle_feature_errors -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 - ) +@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 +@handle_feature_errors +def build(ctx: click.Context, domain: str, should_index: bool) -> None: + """Build context for all domain's datasources. - return results + 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.features.build import build_impl + from databao_cli.shared.cli_utils import get_project_or_raise + + 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.") diff --git a/src/databao_cli/commands/datasource/__init__.py b/src/databao_cli/commands/datasource/__init__.py index e69de29b..531ccc5f 100644 --- a/src/databao_cli/commands/datasource/__init__.py +++ b/src/databao_cli/commands/datasource/__init__.py @@ -0,0 +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.add_command(add) +datasource.add_command(check) diff --git a/src/databao_cli/commands/datasource/add.py b/src/databao_cli/commands/datasource/add.py new file mode 100644 index 00000000..ecd73ffe --- /dev/null +++ b/src/databao_cli/commands/datasource/add.py @@ -0,0 +1,24 @@ +import click + +from databao_cli.shared.cli_utils import get_project_or_raise, handle_feature_errors + + +@click.command(name="add") +@click.option( + "-d", + "--domain", + type=click.STRING, + default="root", + 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 + """ + from databao_cli.workflows.datasource.add import add_workflow + + 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 new file mode 100644 index 00000000..c58b85bc --- /dev/null +++ b/src/databao_cli/commands/datasource/check.py @@ -0,0 +1,32 @@ +import click + +from databao_cli.shared.cli_utils import get_project_or_raise, handle_feature_errors + + +@click.command(name="check") +@click.argument( + "domains", + type=click.STRING, + nargs=-1, +) +@click.pass_context +@handle_feature_errors +def check(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.features.datasource.check import check_impl + from databao_cli.workflows.datasource.check import print_connection_check_results + + project_layout = get_project_or_raise(ctx.obj["project_dir"]) + results = check_impl(project_layout, requested_domains=list(domains) if domains else None) + + if all(len(v) == 0 for v in results.values()): + click.echo("No datasource found") + else: + for domain, datasource_results in results.items(): + print_connection_check_results(domain, datasource_results) diff --git a/src/databao_cli/commands/datasource/check_datasource_connection.py b/src/databao_cli/commands/datasource/check_datasource_connection.py deleted file mode 100644 index 7aad5411..00000000 --- a/src/databao_cli/commands/datasource/check_datasource_connection.py +++ /dev/null @@ -1,53 +0,0 @@ -import os - -import click -from databao_context_engine import ( - CheckDatasourceConnectionResult, - DatabaoContextDomainManager, - DatasourceId, -) - -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.echo(f"{fq_datasource_name}: {status}") - - -def check_datasource_connection_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) - - if all([len(domain_results) == 0 for domain_results in results.values()]): - click.echo("No datasource found") - return - - for domain, datasource_results in results.items(): - print_connection_check_results(domain, datasource_results) - - -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/commands/index.py b/src/databao_cli/commands/index.py index 0e8c4406..b6f2b5b0 100644 --- a/src/databao_cli/commands/index.py +++ b/src/databao_cli/commands/index.py @@ -1,17 +1,33 @@ -from databao_context_engine import ChunkEmbeddingMode, DatabaoContextDomainManager, DatasourceId, IndexDatasourceResult +import click -from databao_cli.project.layout import ProjectLayout +from databao_cli.shared.cli_utils import handle_feature_errors -def index_impl( - project_layout: ProjectLayout, domain: str, datasources_config_files: list[str] | None -) -> list[IndexDatasourceResult]: - dce_project_dir = project_layout.domains_dir / domain +@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 +@handle_feature_errors +def index(ctx: click.Context, domain: str, datasources_config_files: tuple[str, ...]) -> None: + """Index built contexts into the embeddings database. - datasource_ids = [DatasourceId.from_string_repr(p) for p in datasources_config_files] if datasources_config_files else None + 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.features.index import index_impl + from databao_cli.shared.cli_utils import get_project_or_raise - 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 + 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.") diff --git a/src/databao_cli/commands/init.py b/src/databao_cli/commands/init.py index bbcb795c..51f75899 100644 --- a/src/databao_cli/commands/init.py +++ b/src/databao_cli/commands/init.py @@ -1,72 +1,44 @@ +import sys from pathlib import Path -from databao_context_engine import InitDomainError, init_dce_domain - -from databao_cli.project.layout import ProjectLayout, find_project - - -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()}" - ) +import click + +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.workflows.datasource.add import add_workflow as 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) diff --git a/src/databao_cli/commands/mcp.py b/src/databao_cli/commands/mcp.py index 614a7d28..e06a96a1 100644 --- a/src/databao_cli/commands/mcp.py +++ b/src/databao_cli/commands/mcp.py @@ -1,10 +1,50 @@ """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 -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) +@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 +@handle_feature_errors +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.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) diff --git a/src/databao_cli/commands/status.py b/src/databao_cli/commands/status.py index 7737396c..492b4b21 100644 --- a/src/databao_cli/commands/status.py +++ b/src/databao_cli/commands/status.py @@ -1,46 +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.shared.cli_utils import handle_feature_errors -from databao_cli.project.layout import find_project +@click.command() +@click.pass_context +@handle_feature_errors +def status(ctx: click.Context) -> None: + """Display project status and system-wide information.""" + from databao_cli.features.status import status_impl -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) + 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/commands/datasource/add_datasource_config.py b/src/databao_cli/workflows/datasource/add.py similarity index 56% rename from src/databao_cli/commands/datasource/add_datasource_config.py rename to src/databao_cli/workflows/datasource/add.py index b1bb83e2..6517998c 100644 --- a/src/databao_cli/commands/datasource/add_datasource_config.py +++ b/src/databao_cli/workflows/datasource/add.py @@ -1,3 +1,5 @@ +"""Interactive CLI workflow for adding a datasource.""" + import os import click @@ -7,40 +9,47 @@ DatasourceType, ) -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 +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 -def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None: 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: + 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 {datasource_id.relative_path_to_config_file()}. " + 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, ) - created_datasource = domain_manager.create_datasource_config_interactively( - datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True - ) + 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)}" + 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) @@ -53,5 +62,4 @@ def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> 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..217bf904 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -2,14 +2,15 @@ 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 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 2d2a3310..c170bb4e 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -2,14 +2,15 @@ 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 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( 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: