Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ dependencies = [
"databao-context-engine[snowflake]~=0.6.0",
"prettytable>=3.10.0",
"databao-agent~=0.2.0",
"questionary>=2.1.1",
"streamlit[snowflake]>=1.53.0",
"uuid6>=2024.7.10",
"pyyaml>=6.0",
Expand Down
10 changes: 7 additions & 3 deletions src/databao_cli/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@
import click
from click import Context

from databao_cli.labels import LABELS
from databao_cli.log.logging import configure_logging
from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project
from databao_cli.utils import ask_confirm, register_labels


@click.group()
Expand All @@ -24,6 +26,8 @@ def cli(ctx: Context, verbose: bool, project_dir: Path | None) -> None:
ctx.ensure_object(dict)
ctx.obj["project_dir"] = project_path

register_labels(LABELS)

configure_logging(find_project(project_path), verbose=verbose)


Expand All @@ -49,7 +53,7 @@ def init(ctx: Context) -> None:
try:
project_layout = init_impl(project_dir)
except ProjectDirDoesnotExistError:
if click.confirm(
if ask_confirm(
f"The directory {project_dir.resolve()} does not exist. Do you want to create it?",
default=True,
):
Expand All @@ -69,12 +73,12 @@ def init(ctx: Context) -> None:
# except RuntimeError as e:
# click.echo(str(e), err=True)

if not click.confirm("\nDo you want to configure a domain now?"):
if not ask_confirm("Do you want to configure a domain now?", default=False):
return

add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN)

while click.confirm("\nDo you want to add more datasources?"):
while ask_confirm("Do you want to add more datasources?", default=False):
add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN)


Expand Down
6 changes: 3 additions & 3 deletions src/databao_cli/commands/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@


def app_impl(ctx: click.Context) -> None:
click.echo("Starting Databao UI...")
click.echo("Starting Databao web interface...")

try:
bootstrap_streamlit_app(
Expand All @@ -20,7 +20,7 @@ def app_impl(ctx: click.Context) -> None:
hide_build_context_hint=ctx.obj.get("hide_build_context_hint", False),
)
except subprocess.CalledProcessError as e:
click.echo(f"Error running Streamlit: {e}", err=True)
click.echo(f"Error starting web interface: {e}", err=True)
sys.exit(1)
except KeyboardInterrupt:
click.echo("\nShutting down Databao...")
click.echo("\nShutting down...")
20 changes: 10 additions & 10 deletions src/databao_cli/commands/ask.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def dataframe_to_prettytable(df: pd.DataFrame, max_rows: int = DEFAULT_MAX_DISPL


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."""
"""Initialize the Databao agent using a Context Engine project at the given path."""
project = ProjectLayout(project_path)

status = databao_project_status(project)
Expand All @@ -49,12 +49,12 @@ def initialize_agent_from_dce(project_path: Path, model: str | None, temperature

if status == DatabaoProjectStatus.NO_DATASOURCES:
click.echo(
f"No datasources configured in project at {project.project_dir}. Add datasources first.",
f"No data sources configured in project at {project.project_dir}. Add data sources first",
err=True,
)
sys.exit(1)

click.echo(f"Using DCE project: {project.project_dir}")
click.echo(f"Using Context Engine project: {project.project_dir}")

_domain = create_domain(project.root_domain_dir)

Expand Down Expand Up @@ -101,18 +101,18 @@ def display_result(thread: Thread) -> None:

def _print_help() -> None:
"""Print help message for interactive mode."""
click.echo("Databao REPL")
click.echo("Databao interactive chat")
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(" \\help - Show this help message")
click.echo(" \\clear - Clear conversation history")
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")
click.echo("\nDatabao interactive chat")
click.echo("Type \\help for available commands.\n")

writer = _create_cli_writer() if show_thinking else None

Expand Down Expand Up @@ -148,14 +148,14 @@ def run_interactive_mode(agent: Agent, show_thinking: bool) -> None:
stream_ask=show_thinking,
writer=writer,
)
click.echo("Conversation cleared.\n")
click.echo("Conversation history cleared.\n")
continue

if command == "help":
_print_help()
continue

click.echo(f"Unknown command: {user_input}. Type \\help for available commands.\n")
click.echo(f"Unknown command: {user_input}\nType \\help for available commands\n")
continue

# Process as a question
Expand Down
79 changes: 66 additions & 13 deletions src/databao_cli/commands/context_engine_cli.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import sys
from typing import Any

import click
import questionary
from databao_context_engine import Choice, UserInputCallback


Expand All @@ -15,19 +17,70 @@ def prompt(
show_default: bool = default_value is not None and default_value != ""
final_type = click.Choice(type.choices) if isinstance(type, Choice) else str

# click goes infinite loop if user gives emptry string as an input AND default_value is None
# in order to exit this loop we need to set default value to '' (so it gets accepted)
#
# Code snippet from click:
# while True:
# value = prompt_func(prompt)
# if value:
# break
# elif default is not None:
# value = default
# break
default_value = default_value if default_value else "" if final_type is str else None
return click.prompt(text=text, default=default_value, hide_input=is_secret, type=final_type, show_default=show_default)
# Determine if this field is optional
is_optional = "(Optional)" in text

if isinstance(type, Choice):
is_interactive = sys.stdin.isatty() and sys.stdout.isatty()
if is_interactive:
from databao_cli.labels import LABELS

choices = [questionary.Choice(title=LABELS.get(choice, choice), value=choice) for choice in type.choices]
result = questionary.select(
text,
choices=choices,
default=default_value if default_value is not None and default_value in type.choices else None,
).ask()
if result is None:
raise click.Abort()
return result
else:
return click.prompt(
text=text,
default=default_value,
hide_input=is_secret,
type=click.Choice(type.choices),
show_default=show_default,
)

if default_value:
final_default = default_value
elif is_optional:
final_default = ""
else:
final_default = None

is_interactive = sys.stdin.isatty() and sys.stdout.isatty()
if is_interactive and final_type is str:
prompt_func = questionary.password if is_secret else questionary.text

if not is_optional and final_default is None:
while True:
result = prompt_func(text, default=final_default or "").ask()
if result is None:
raise click.Abort()
value = str(result).strip()
if value:
return value
click.echo("This field is required and cannot be empty. Please try again.")
else:
result = prompt_func(text, default=final_default or "").ask()
if result is None:
raise click.Abort()
return str(result)
else:
if final_type is str and not is_optional and final_default is None:
while True:
value = click.prompt(
text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
)
if value and value.strip():
return value
click.echo("This field is required and cannot be empty. Please try again.")
else:
return click.prompt(
text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
)

def confirm(self, text: str) -> bool:
return click.confirm(text=text)
31 changes: 21 additions & 10 deletions src/databao_cli/commands/datasource/add_datasource_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
DatabaoContextPluginLoader,
DatasourceType,
)
from pydantic import ValidationError

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.utils import ask_confirm, ask_select, ask_text


def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None:
Expand All @@ -21,37 +23,46 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain

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_name = ask_text("Datasource name?")

datasource_id = domain_manager.datasource_config_exists(datasource_name=datasource_name)
if datasource_id is not None:
click.confirm(
ask_confirm(
f"A config file already exists for this datasource {datasource_id.relative_path_to_config_file()}. "
f"Do you want to overwrite it?",
abort=True,
default=False,
)
created_datasource = domain_manager.create_datasource_config_interactively(
datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True
)

while True:
try:
created_datasource = domain_manager.create_datasource_config_interactively(
datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True
)
break
except ValidationError as e:
click.echo(click.style("\nValidation error:", fg="red", bold=True))
for error in e.errors():
field_path = ".".join(str(loc) for loc in error["loc"])
click.echo(click.style(f" • {field_path}: {error['msg']}", fg="red"))
click.echo("\nPlease try again with correct values.\n")

datasource_id = created_datasource.datasource.id
click.echo(
f"{os.linesep}We've created a new config file for your datasource at: "
f"{domain_manager.get_config_file_path_for_datasource(datasource_id)}"
)
if click.confirm("\nDo you want to check the connection to this new datasource?"):
if ask_confirm("Do you want to check the connection to this new datasource?", default=True):
results = domain_manager.check_datasource_connection(datasource_ids=[datasource_id])
print_connection_check_results(domain, results)


def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType:
all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types])
config_type = click.prompt(
config_type = ask_select(
"What type of datasource do you want to add?",
type=click.Choice(all_datasource_types),
default=all_datasource_types[0] if len(all_datasource_types) == 1 else None,
choices=all_datasource_types,
default=None,
)
Comment on lines +62 to 66
click.echo(f"Selected type: {config_type}")

return DatasourceType(full_type=config_type)
16 changes: 8 additions & 8 deletions src/databao_cli/commands/status.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,19 +28,19 @@ def status_impl(project_dir: Path) -> str:

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"Context Engine version: {command_info.version}",
f"Agent version: {version('databao-agent')}",
f"Context Engine storage directory: {command_info.dce_path}",
f"Context Engine plugins: {command_info.plugin_ids}",
"",
f"OS name: {sys.platform}",
f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}",
f"OS: {sys.platform}",
f"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}")
info_lines.append(f"Domain directory: {domain_info.project_path.resolve()}")
info_lines.append(f"Domain ID: {domain_info.project_id!s}")

return os.linesep.join(info_lines)
35 changes: 35 additions & 0 deletions src/databao_cli/labels.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
LABELS = {
"athena": "Amazon Athena",
"bigquery": "BigQuery",
"clickhouse": "ClickHouse",
"duckdb": "DuckDB",
"mssql": "Microsoft SQL Server",
"mysql": "MySQL",
"parquet": "Parquet",
"postgres": "PostgreSQL",
"snowflake": "Snowflake",
"sqlite": "SQLite",
"BigQueryDefaultAuth": "Default auth",
"BigQueryServiceAccountJsonAuth": "Service account JSON credentials",
"BigQueryServiceAccountKeyFileAuth": "Service account key file",
"SnowflakeKeyPairAuth": "Key pair",
"SnowflakePasswordAuth": "Password",
"SnowflakeSSOAuth": "SSO",
"connection.auth.type": "Authentication type",
"connection.host": "Host",
"connection.port": "Port",
"connection.database": "Database",
"connection.schema": "Schema",
"connection.username": "Username",
"connection.password": "Password",
"connection.account": "Account",
"connection.warehouse": "Warehouse",
"connection.role": "Role",
"connection.path": "File path",
"connection.project": "Project",
"connection.dataset": "Dataset",
"connection.location": "Location",
"connection.auth.credentials_file": "Credentials file",
"connection.auth.key_file": "Key file",
"connection.auth.token": "Token",
}
8 changes: 4 additions & 4 deletions src/databao_cli/mcp/tools/databao_ask/agent_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,19 @@ def create_agent_for_tool(
executor: str = "claude_code",
cache: Cache | None = None,
) -> Agent:
"""Create a Databao agent from a DCE project, configured for MCP tool use.
"""Create a Databao agent from a Context Engine project, configured for MCP tool use.

Raises ValueError if the project is not ready (no datasources, no build).
"""
project = ProjectLayout(project_dir)

status = databao_project_status(project)
if status == DatabaoProjectStatus.NOT_INITIALIZED:
raise ValueError("Databao project is not initialized. Run 'databao init' first.")
raise ValueError("Databao project is not initialized. Run 'databao init' first")
if status == DatabaoProjectStatus.NO_DATASOURCES:
raise ValueError("No datasources configured. Run 'databao datasource add' first.")
raise ValueError("No data sources configured. Run 'databao datasource add' first")
if not has_build_output(project):
raise ValueError("Project has no build output. Run 'databao build' first.")
raise ValueError("Project has no build output. Run 'databao build' first")

domain = create_domain(project.root_domain_dir)

Expand Down
Loading
Loading