nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/cli/commands/indexing.py

@@ -0,0 +1,126 @@
+import json
+import sys
+from typing import Any, Dict
+from nl2sql.cli.reporting import ConsolePresenter
+from nl2sql.indexing.vector_store import VectorStore
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.context import NL2SQLContext
+from nl2sql.cli.common.decorators import handle_cli_errors
+from nl2sql.indexing.orchestrator import IndexingOrchestrator
+
+@handle_cli_errors
+def run_indexing(
+    ctx: NL2SQLContext,
+) -> None:
+    """
+    Runs schema indexing for all registered datasources.
+
+    This command clears the existing vector store and indexes
+    schema chunks for each configured datasource using the
+    indexing orchestrator.
+
+    Args:
+        ctx: The initialized NL2SQLContext.
+
+    Raises:
+        SystemExit: With code 1 if the store could not be cleared or if any
+            datasource failed to index.
+    """
+    presenter = ConsolePresenter()
+    presenter.print_info(f"Indexing schema to: {ctx.vector_store.persist_directory}")
+
+    adapters = ctx.ds_registry.list_adapters()
+    orchestrator = IndexingOrchestrator(ctx)
+    stats = []
+    errors = []
+    empty_stats = []
+
+    presenter.start_interactive_status("Clearing existing data...")
+    try:
+        orchestrator.clear_store()
+    except Exception as e:
+        presenter.stop_interactive_status()
+        presenter.print_error(f"Failed to clear existing data: {e}")
+        sys.exit(1)
+    presenter.stop_interactive_status()
+    presenter.print_success("Cleared existing data.")
+
+    for adapter in adapters:
+        ds_id = adapter.datasource_id
+
+        try:
+            task = presenter.start_task_line(f"Indexing {ds_id}...")
+            schema_stats = orchestrator.index_datasource(adapter)
+            if schema_stats:
+                stats.append(schema_stats)
+            else:
+                empty_stats.append(ds_id)
+            presenter.finish_task_line(task, f"{ds_id} indexed", success=True)
+        except Exception as e:
+            if "task" in locals():
+                presenter.finish_task_line(task, f"{ds_id} failed", success=False)
+            presenter.print_error(f"Failed to index {ds_id}: {e}")
+            errors.append(
+                {"datasource_id": ds_id, "error": str(e)}
+            )
+    
+    if errors:
+        presenter.print_table(errors, "Indexing Errors", columns=["datasource_id", "error"])
+    
+    if stats:
+        summary_rows = []
+        total_chunks = 0
+        totals_by_type: Dict[str, int] = {}
+
+        for s in stats:
+            ds_id = s.get("datasource_id", "unknown")
+            schema_version = s.get("schema_version", "-")
+            chunk_stats = {
+                k: v for k, v in s.items()
+                if k not in ("datasource_id", "schema_version")
+            }
+            ds_total = sum(v for v in chunk_stats.values() if isinstance(v, int))
+            total_chunks += ds_total
+            for key, val in chunk_stats.items():
+                if isinstance(val, int):
+                    totals_by_type[key] = totals_by_type.get(key, 0) + val
+
+            summary_rows.append(
+                {
+                    "datasource_id": ds_id,
+                    "schema_version": schema_version,
+                    "total_chunks": ds_total,
+                    "chunks": json.dumps(chunk_stats, separators=(",", ":")),
+                }
+            )
+
+        presenter.print_table(
+            summary_rows,
+            "Indexing Summary",
+            columns=["datasource_id", "schema_version", "total_chunks", "chunks"],
+        )
+
+        if totals_by_type:
+            totals_str = ", ".join([f"{k}={v}" for k, v in sorted(totals_by_type.items())])
+            presenter.print_info(f"Total chunks indexed: {total_chunks} ({totals_str})")
+
+    total_adapters = len(adapters)
+    succeeded = len(stats)
+    failed = len(errors)
+    skipped = len(empty_stats)
+    presenter.print_info(
+        f"Datasources: total={total_adapters}, succeeded={succeeded}, failed={failed}, empty={skipped}"
+    )
+
+    if empty_stats:
+        presenter.print_warning(f"Datasources with empty stats: {', '.join(empty_stats)}")
+
+    if errors:
+        # Any failure is fatal, not just a total one: a partially populated
+        # index answers later queries from an incomplete schema, which is worse
+        # than a loud stop because it looks like it worked. Scripts chaining
+        # `nl2sql index && nl2sql run ...` need that signal.
+        presenter.print_warning("Indexing completed with errors.")
+        sys.exit(1)
+
+    presenter.print_success("Indexing complete.")

nl2sql/cli/commands/info.py

@@ -0,0 +1,25 @@
+from rich.console import Console
+from rich.table import Table
+from nl2sql.datasources.discovery import discover_adapters
+
+from nl2sql.cli.common.decorators import handle_cli_errors
+
+@handle_cli_errors
+def list_available_adapters() -> None:
+    """Discovers and displays all installed Datasource Adapters."""
+    console = Console()
+    adapters = discover_adapters()
+
+    if not adapters:
+        console.print("[yellow]No adapters found. Please install a driver extra (e.g., nl2sql\[postgres]).[/yellow]")
+        return
+
+    table = Table(title="Installed Datasource Adapters")
+    table.add_column("Adapter ID", style="cyan", no_wrap=True)
+    table.add_column("Class", style="magenta")
+    table.add_column("Status", style="green")
+
+    for name, cls in adapters.items():
+        table.add_row(name, f"{cls.__module__}.{cls.__name__}", "Active")
+
+    console.print(table)

nl2sql/cli/commands/install.py

@@ -0,0 +1,27 @@
+import sys
+import subprocess
+from rich.markup import escape
+from rich.prompt import Confirm
+from nl2sql.cli.config import KNOWN_ADAPTERS
+from nl2sql.cli.console import console, print_success, print_error, print_step
+
+def install_package(package_name: str) -> bool:
+    print_step(f"Installing {package_name}...")
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
+        print_success(f"Installed {package_name}")
+        return True
+    except subprocess.CalledProcessError:
+        print_error(f"Failed to install {package_name}")
+        return False
+
+from nl2sql.cli.common.decorators import handle_cli_errors
+
+@handle_cli_errors
+def install_command(adapter_name: str):
+    target_pkg = adapter_name
+    if adapter_name in KNOWN_ADAPTERS:
+        target_pkg = KNOWN_ADAPTERS[adapter_name]
+    
+    if Confirm.ask(f"Install [cyan]{escape(str(target_pkg))}[/cyan]?"):
+        install_package(target_pkg)

nl2sql/cli/commands/policy.py

@@ -0,0 +1,57 @@
+import typer
+import pathlib
+import sys
+from typing import Optional
+from typing_extensions import Annotated
+from rich.console import Console
+from rich.markup import escape
+from rich.text import Text
+from rich.table import Table
+
+from nl2sql import PolicyAPI
+
+app = typer.Typer(help="Manage RBAC policies and security.")
+console = Console()
+
+from nl2sql.cli.common.decorators import handle_cli_errors
+
+@app.command("validate")
+@handle_cli_errors
+def validate(
+    config: Annotated[Optional[pathlib.Path], typer.Option("--config", help="Path to datasource config")] = None,
+    policies: Annotated[Optional[pathlib.Path], typer.Option("--policies", help="Path to policies.json")] = None,
+    secrets: Annotated[Optional[pathlib.Path], typer.Option("--secrets", help="Path to secrets config")] = None,
+):
+    """
+    Validate policy syntax and integrity against defined datasources.
+    """
+    console.print(f"[bold blue]Validating Policies from:[/bold blue] {escape(str(policies or 'default'))}")
+
+    api = PolicyAPI()
+    report = api.validate_policies(policies_path=policies, datasources_path=config, secrets_path=secrets)
+
+    if report.errors:
+        console.print(f"[bold red]Schema Validation Failed:[/bold red]\\n{escape(str(report.errors[0]))}")
+        sys.exit(1)
+
+    console.print(f"[bold blue]Checking Integrity against Datasources:[/bold blue] {escape(str(config or 'default'))}")
+    if report.available_datasources:
+        console.print(f"[dim]Available Datasources: {report.available_datasources}[/dim]")
+
+    table = Table(title="Policy Integrity Report")
+    table.add_column("Role", style="cyan")
+    table.add_column("Target", style="magenta")
+    table.add_column("Status", style="green")
+    table.add_column("Details", style="white")
+
+    for entry in report.entries:
+        status = "[green]OK[/green]" if entry.status == "OK" else f"[red]{entry.status}[/red]"
+        table.add_row(Text(str(entry.role)), Text(str(entry.target)), status, Text(str(entry.details)))
+
+    console.print(table)
+
+    if not report.ok:
+        console.print("\n[bold red]Integrity Check Failed: Policies reference missing resources.[/bold red]")
+        sys.exit(1)
+    else:
+        console.print("\n[bold green]✓ Policy Integrity Verified[/bold green]")

nl2sql/cli/commands/run.py

@@ -0,0 +1,166 @@
+import sys
+import json
+
+from rich.text import Text
+
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.llm import LLMRegistry
+from nl2sql.indexing.vector_store import VectorStore
+from nl2sql.cli.reporting import ConsolePresenter
+from nl2sql.pipeline.pipeline_runner import PipelineRunner
+from nl2sql.common.settings import settings
+from nl2sql.cli.types import RunConfig
+from nl2sql.cli.common.decorators import handle_cli_errors
+from nl2sql.context import NL2SQLContext
+
+@handle_cli_errors
+def run_pipeline(
+    config: RunConfig, 
+    ctx: NL2SQLContext
+) -> None:
+    """Executes the NL2SQL pipeline."""
+    if not config.query:
+        return
+        
+    presenter = ConsolePresenter()
+    presenter.print_info(f"Query: {config.query}")
+    if config.no_exec:
+        presenter.print_warning("Execution disabled (no_exec). Only SQL/plan output will be shown.")
+    
+    # Instantiate Runner
+    runner = PipelineRunner(ctx)
+    
+    # Setup Monitoring
+    from nl2sql.services.callbacks.monitor import PipelineMonitorCallback
+    monitor = PipelineMonitorCallback(presenter)
+    
+    presenter.start_interactive_status("Thinking...")
+    
+    # Execution
+    result = runner.run(
+        query=config.query,
+        role=config.role,
+        datasource_id=config.ds_id,
+        execute=not config.no_exec,
+        callbacks=[monitor]
+    )
+    
+    presenter.stop_interactive_status()
+    
+    # Handle Result
+    if not result.success:
+        if result.traceback:
+            presenter.print_error(result.traceback)
+        presenter.print_error(result.error or "Unknown Pipeline Error")
+        sys.exit(1)
+        
+    final_state = result.final_state
+    
+
+    query_history = []
+    subgraph_outputs = final_state.get("subgraph_outputs") or {}
+    for _subgraph_id, output in subgraph_outputs.items():
+        if isinstance(output, dict):
+            sub_query = output.get("sub_query")
+            sql_draft = output.get("sql_draft")
+        else:
+            sub_query = getattr(output, "sub_query", None)
+            sql_draft = getattr(output, "sql_draft", None)
+
+        if not sub_query:
+            continue
+
+        if isinstance(sub_query, dict):
+            ds_id = sub_query.get("datasource_id")
+            intent = sub_query.get("intent")
+        else:
+            ds_id = getattr(sub_query, "datasource_id", None)
+            intent = getattr(sub_query, "intent", None)
+
+        query_history.append({
+            "sub_query": intent,
+            "datasource_id": ds_id,
+            "sql": sql_draft,
+        })
+    
+    if config.verbose:
+        reasoning = final_state.get("reasoning", [])
+        presenter.print_execution_tree(config.query, query_history, top_level_reasoning=reasoning)
+        
+        try:
+            with open("last_reasoning.json", "w") as f:
+                dump_data = {
+                    "global_reasoning": reasoning,
+                    "execution_history": query_history
+                }
+                json.dump(dump_data, f, indent=2, default=str)
+            presenter.print_info("Detailed reasoning trace saved to last_reasoning.json")
+        except Exception:
+            pass
+
+    if query_history:
+        datasources_used = sorted({item.get("datasource_id") for item in query_history if item.get("datasource_id")})
+        if datasources_used:
+            presenter.print_info(f"Datasources used: {', '.join(datasources_used)}")
+
+        for item in query_history:
+            ds = item.get("datasource_id", "Unknown")
+            sub_query = item.get("sub_query")
+            sql = item.get("sql")
+
+            if sql:
+                body = Text()
+                if sub_query:
+                    body.append(f"Sub-Query: {sub_query}\n", style="bold")
+                body.append(f"Datasource: {ds}\n\n", style="bold")
+                body.append(str(sql))
+                presenter.print_sql(body)
+    
+    elif final_state.get("sql_draft"):
+        sql_draft_data = final_state.get("sql_draft")
+        sql_draft = sql_draft_data.get("sql") if isinstance(sql_draft_data, dict) else getattr(sql_draft_data, "sql", None)
+        if sql_draft:
+             presenter.print_sql(str(sql_draft))
+
+    errors = final_state.get("errors")
+    if errors:
+        presenter.print_pipeline_errors(errors)
+
+    warnings = final_state.get("warnings") or []
+    for warning in warnings:
+        if isinstance(warning, dict):
+            presenter.print_warning(json.dumps(warning, default=str))
+        else:
+            presenter.print_warning(str(warning))
+
+    answer_payload = None
+    answer_synth = final_state.get("answer_synthesizer_response")
+    if answer_synth:
+        if isinstance(answer_synth, dict):
+            answer_payload = answer_synth.get("final_answer")
+        else:
+            answer_payload = getattr(answer_synth, "final_answer", None)
+    if isinstance(answer_payload, dict) and answer_payload:
+        presenter.print_answer_synthesizer_output(answer_payload)
+
+    final_answer = final_state.get("final_answer")
+    if isinstance(final_answer, dict) and final_answer:
+        presenter.print_answer_synthesizer_output(final_answer)
+    elif type(final_answer) == str and final_answer:
+        presenter.print_final_answer(final_answer)
+    elif type(final_answer) == list and final_answer:
+        presenter.print_execution_result(final_answer)
+        
+    execution = final_state.get("execution")
+    if execution:
+        row_count = execution.get("row_count", 0) if isinstance(execution, dict) else getattr(execution, "row_count", 0)
+        presenter.print_rows_returned(row_count)
+
+
+    if config.show_perf:
+        tree, metrics, node_map = monitor.get_performance_tree()
+        presenter.print_performance_tree(tree, metrics, node_map)
+    
+    
+    from nl2sql.common.metrics import TOKEN_LOG
+    presenter.print_cost_summary(result.duration, TOKEN_LOG)