sqlsaber 0.1.0__py3-none-any.whl

sqlsaber/cli/interactive.py

@@ -0,0 +1,93 @@
+"""Interactive mode handling for the CLI."""
+
+import questionary
+from rich.console import Console
+from rich.panel import Panel
+
+from sqlsaber.agents.base import BaseSQLAgent
+from sqlsaber.cli.display import DisplayManager
+from sqlsaber.cli.streaming import StreamingQueryHandler
+
+
+class InteractiveSession:
+    """Manages interactive CLI sessions."""
+
+    def __init__(self, console: Console, agent: BaseSQLAgent):
+        self.console = console
+        self.agent = agent
+        self.display = DisplayManager(console)
+        self.streaming_handler = StreamingQueryHandler(console)
+
+    def show_welcome_message(self):
+        """Display welcome message for interactive mode."""
+        self.console.print(
+            Panel.fit(
+                "[bold green]SQLSaber - Use the agent Luke![/bold green]\n\n"
+                "Type your queries in natural language.\n\n"
+                "Press Esc-Enter or Meta-Enter to submit your query.\n\n"
+                "Type 'exit' or 'quit' to leave.",
+                border_style="green",
+            )
+        )
+
+        self.console.print(
+            "[dim]Commands: 'clear' to reset conversation, 'exit' or 'quit' to leave[/dim]"
+        )
+        self.console.print(
+            "[dim]Memory: Start a message with '#' to add it as a memory for this database[/dim]\n"
+        )
+
+    async def run(self):
+        """Run the interactive session loop."""
+        self.show_welcome_message()
+
+        while True:
+            try:
+                user_query = await questionary.text(
+                    ">",
+                    qmark="",
+                    multiline=True,
+                    instruction="",
+                ).ask_async()
+
+                if user_query.lower() in ["exit", "quit", "q"]:
+                    break
+
+                if user_query.lower() == "clear":
+                    self.agent.clear_history()
+                    self.console.print("[green]Conversation history cleared.[/green]\n")
+                    continue
+
+                if memory_text := user_query.strip():
+                    # Check if query starts with # for memory addition
+                    if memory_text.startswith("#"):
+                        memory_content = memory_text[1:].strip()  # Remove # and trim
+                        if memory_content:
+                            # Add memory
+                            memory_id = self.agent.add_memory(memory_content)
+                            if memory_id:
+                                self.console.print(
+                                    f"[green]✓ Memory added:[/green] {memory_content}"
+                                )
+                                self.console.print(
+                                    f"[dim]Memory ID: {memory_id}[/dim]\n"
+                                )
+                            else:
+                                self.console.print(
+                                    "[yellow]Could not add memory (no database context)[/yellow]\n"
+                                )
+                        else:
+                            self.console.print(
+                                "[yellow]Empty memory content after '#'[/yellow]\n"
+                            )
+                        continue
+
+                    await self.streaming_handler.execute_streaming_query(
+                        user_query, self.agent
+                    )
+                    self.display.show_newline()  # Empty line for readability
+
+            except KeyboardInterrupt:
+                self.console.print("\n[yellow]Use 'exit' or 'quit' to leave.[/yellow]")
+            except Exception as e:
+                self.console.print(f"[bold red]Error:[/bold red] {str(e)}")

sqlsaber/cli/memory.py

@@ -0,0 +1,239 @@
+"""Memory management CLI commands."""
+
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.memory.manager import MemoryManager
+
+# Global instances for CLI commands
+console = Console()
+config_manager = DatabaseConfigManager()
+memory_manager = MemoryManager()
+
+# Create the memory management CLI app
+memory_app = typer.Typer(
+    name="memory",
+    help="Manage database-specific memories",
+    add_completion=True,
+)
+
+
+def _get_database_name(database: Optional[str] = None) -> str:
+    """Get the database name to use, either specified or default."""
+    if database:
+        db_config = config_manager.get_database(database)
+        if not db_config:
+            console.print(
+                f"[bold red]Error:[/bold red] Database connection '{database}' not found."
+            )
+            raise typer.Exit(1)
+        return database
+    else:
+        db_config = config_manager.get_default_database()
+        if not db_config:
+            console.print(
+                "[bold red]Error:[/bold red] No database connections configured."
+            )
+            console.print("Use 'sqlsaber db add <name>' to add a database connection.")
+            raise typer.Exit(1)
+        return db_config.name
+
+
+@memory_app.command("add")
+def add_memory(
+    content: str = typer.Argument(..., help="Memory content to add"),
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """Add a new memory for the specified database."""
+    database_name = _get_database_name(database)
+
+    try:
+        memory = memory_manager.add_memory(database_name, content)
+        console.print(f"[green]✓ Memory added for database '{database_name}'[/green]")
+        console.print(f"[dim]Memory ID:[/dim] {memory.id}")
+        console.print(f"[dim]Content:[/dim] {memory.content}")
+    except Exception as e:
+        console.print(f"[bold red]Error adding memory:[/bold red] {e}")
+        raise typer.Exit(1)
+
+
+@memory_app.command("list")
+def list_memories(
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """List all memories for the specified database."""
+    database_name = _get_database_name(database)
+
+    memories = memory_manager.get_memories(database_name)
+
+    if not memories:
+        console.print(
+            f"[yellow]No memories found for database '{database_name}'[/yellow]"
+        )
+        console.print("Use 'sqlsaber memory add \"<content>\"' to add memories")
+        return
+
+    table = Table(title=f"Memories for Database: {database_name}")
+    table.add_column("ID", style="cyan", width=36)
+    table.add_column("Content", style="white")
+    table.add_column("Created", style="dim")
+
+    for memory in memories:
+        # Truncate content if it's too long for display
+        display_content = memory.content
+        if len(display_content) > 80:
+            display_content = display_content[:77] + "..."
+
+        table.add_row(memory.id, display_content, memory.formatted_timestamp())
+
+    console.print(table)
+    console.print(f"\n[dim]Total memories: {len(memories)}[/dim]")
+
+
+@memory_app.command("show")
+def show_memory(
+    memory_id: str = typer.Argument(..., help="Memory ID to show"),
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """Show the full content of a specific memory."""
+    database_name = _get_database_name(database)
+
+    memory = memory_manager.get_memory_by_id(database_name, memory_id)
+
+    if not memory:
+        console.print(
+            f"[bold red]Error:[/bold red] Memory with ID '{memory_id}' not found for database '{database_name}'"
+        )
+        raise typer.Exit(1)
+
+    console.print(f"[bold]Memory ID:[/bold] {memory.id}")
+    console.print(f"[bold]Database:[/bold] {memory.database}")
+    console.print(f"[bold]Created:[/bold] {memory.formatted_timestamp()}")
+    console.print("[bold]Content:[/bold]")
+    console.print(f"{memory.content}")
+
+
+@memory_app.command("remove")
+def remove_memory(
+    memory_id: str = typer.Argument(..., help="Memory ID to remove"),
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """Remove a specific memory by ID."""
+    database_name = _get_database_name(database)
+
+    # First check if memory exists
+    memory = memory_manager.get_memory_by_id(database_name, memory_id)
+    if not memory:
+        console.print(
+            f"[bold red]Error:[/bold red] Memory with ID '{memory_id}' not found for database '{database_name}'"
+        )
+        raise typer.Exit(1)
+
+    # Show memory content before removal
+    console.print("[yellow]Removing memory:[/yellow]")
+    console.print(f"[dim]Content:[/dim] {memory.content}")
+
+    if memory_manager.remove_memory(database_name, memory_id):
+        console.print(
+            f"[green]✓ Memory removed from database '{database_name}'[/green]"
+        )
+    else:
+        console.print(
+            f"[bold red]Error:[/bold red] Failed to remove memory '{memory_id}'"
+        )
+        raise typer.Exit(1)
+
+
+@memory_app.command("clear")
+def clear_memories(
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+    force: bool = typer.Option(
+        False,
+        "--force",
+        "-f",
+        help="Skip confirmation prompt",
+    ),
+):
+    """Clear all memories for the specified database."""
+    database_name = _get_database_name(database)
+
+    # Count memories first
+    memories_count = len(memory_manager.get_memories(database_name))
+
+    if memories_count == 0:
+        console.print(
+            f"[yellow]No memories to clear for database '{database_name}'[/yellow]"
+        )
+        return
+
+    if not force:
+        # Show confirmation
+        console.print(
+            f"[yellow]About to clear {memories_count} memories for database '{database_name}'[/yellow]"
+        )
+        confirm = typer.confirm("Are you sure you want to proceed?")
+        if not confirm:
+            console.print("Operation cancelled")
+            return
+
+    cleared_count = memory_manager.clear_memories(database_name)
+    console.print(
+        f"[green]✓ Cleared {cleared_count} memories for database '{database_name}'[/green]"
+    )
+
+
+@memory_app.command("summary")
+def memory_summary(
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """Show memory summary for the specified database."""
+    database_name = _get_database_name(database)
+
+    summary = memory_manager.get_memories_summary(database_name)
+
+    console.print(f"[bold]Memory Summary for Database: {summary['database']}[/bold]")
+    console.print(f"[dim]Total memories:[/dim] {summary['total_memories']}")
+
+    if summary["total_memories"] > 0:
+        console.print("\n[bold]Recent memories:[/bold]")
+        for memory in summary["memories"][-5:]:  # Show last 5 memories
+            console.print(f"[dim]{memory['timestamp']}[/dim] - {memory['content']}")
+
+
+def create_memory_app() -> typer.Typer:
+    """Return the memory management CLI app."""
+    return memory_app

sqlsaber/cli/models.py

@@ -0,0 +1,231 @@
+"""Model management CLI commands."""
+
+import asyncio
+from typing import Dict, List
+
+import httpx
+import questionary
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from sqlsaber.config.settings import Config
+
+# Global instances for CLI commands
+console = Console()
+
+# Create the model management CLI app
+models_app = typer.Typer(
+    name="models",
+    help="Select and manage models",
+    add_completion=True,
+)
+
+
+class ModelManager:
+    """Manages AI model configuration and fetching."""
+
+    DEFAULT_MODEL = "anthropic:claude-sonnet-4-20250514"
+    MODELS_API_URL = "https://models.dev/api.json"
+
+    async def fetch_available_models(self) -> List[Dict]:
+        """Fetch available models from models.dev API."""
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                response = await client.get(self.MODELS_API_URL)
+                response.raise_for_status()
+                data = response.json()
+
+                # Filter for Anthropic models only
+                anthropic_models = []
+                anthropic_data = data.get("anthropic", {})
+
+                if "models" in anthropic_data:
+                    for model_id, model_info in anthropic_data["models"].items():
+                        # Convert to our format (anthropic:model-name)
+                        formatted_id = f"anthropic:{model_id}"
+
+                        # Extract cost information for display
+                        cost_info = model_info.get("cost", {})
+                        cost_display = ""
+                        if cost_info:
+                            input_cost = cost_info.get("input", 0)
+                            output_cost = cost_info.get("output", 0)
+                            cost_display = f"${input_cost}/{output_cost} per 1M tokens"
+
+                        # Extract context length
+                        limit_info = model_info.get("limit", {})
+                        context_length = limit_info.get("context", 0)
+
+                        anthropic_models.append(
+                            {
+                                "id": formatted_id,
+                                "name": model_info.get("name", model_id),
+                                "description": cost_display,
+                                "context_length": context_length,
+                                "knowledge": model_info.get("knowledge", ""),
+                            }
+                        )
+
+                # Sort by name for better display
+                anthropic_models.sort(key=lambda x: x["name"])
+                return anthropic_models
+        except Exception as e:
+            console.print(f"[red]Error fetching models: {e}[/red]")
+            return []
+
+    def get_current_model(self) -> str:
+        """Get the currently configured model."""
+        config = Config()
+        return config.model_name
+
+    def set_model(self, model_id: str) -> bool:
+        """Set the current model."""
+        try:
+            config = Config()
+            config.set_model(model_id)
+            return True
+        except Exception as e:
+            console.print(f"[red]Error setting model: {e}[/red]")
+            return False
+
+    def reset_model(self) -> bool:
+        """Reset to default model."""
+        return self.set_model(self.DEFAULT_MODEL)
+
+
+model_manager = ModelManager()
+
+
+@models_app.command("list")
+def list_models():
+    """List available AI models."""
+
+    async def fetch_and_display():
+        console.print("[blue]Fetching available models...[/blue]")
+        models = await model_manager.fetch_available_models()
+
+        if not models:
+            console.print(
+                "[yellow]No models available or failed to fetch models[/yellow]"
+            )
+            return
+
+        table = Table(title="Available Anthropic Models")
+        table.add_column("ID", style="cyan")
+        table.add_column("Name", style="green")
+        table.add_column("Description", style="white")
+        table.add_column("Context", style="yellow", justify="right")
+        table.add_column("Current", style="bold red", justify="center")
+
+        current_model = model_manager.get_current_model()
+
+        for model in models:
+            is_current = "✓" if model["id"] == current_model else ""
+            context_str = (
+                f"{model['context_length']:,}" if model["context_length"] else "N/A"
+            )
+
+            # Truncate description if too long
+            description = (
+                model["description"][:50] + "..."
+                if len(model["description"]) > 50
+                else model["description"]
+            )
+
+            table.add_row(
+                model["id"],
+                model["name"],
+                description,
+                context_str,
+                is_current,
+            )
+
+        console.print(table)
+        console.print(f"\n[dim]Current model: {current_model}[/dim]")
+
+    asyncio.run(fetch_and_display())
+
+
+@models_app.command("set")
+def set_model():
+    """Set the AI model to use."""
+
+    async def interactive_set():
+        console.print("[blue]Fetching available models...[/blue]")
+        models = await model_manager.fetch_available_models()
+
+        if not models:
+            console.print("[red]Failed to fetch models. Cannot set model.[/red]")
+            raise typer.Exit(1)
+
+        # Create choices for questionary
+        choices = []
+        for model in models:
+            # Format: "ID - Name (Description)"
+            choice_text = f"{model['id']} - {model['name']}"
+            if model["description"]:
+                choice_text += f" ({model['description'][:50]}{'...' if len(model['description']) > 50 else ''})"
+
+            choices.append({"name": choice_text, "value": model["id"]})
+
+        # Get current model to set as default
+        current_model = model_manager.get_current_model()
+        default_index = 0
+        for i, choice in enumerate(choices):
+            if choice["value"] == current_model:
+                default_index = i
+                break
+
+        selected_model = await questionary.select(
+            "Select a model:",
+            choices=choices,
+            use_shortcuts=True,
+            use_search_filter=True,
+            use_jk_keys=False,  # Disable j/k keys when using search filter
+            default=choices[default_index] if choices else None,
+        ).ask_async()
+
+        if selected_model:
+            if model_manager.set_model(selected_model):
+                console.print(f"[green]✓ Model set to: {selected_model}[/green]")
+            else:
+                console.print("[red]✗ Failed to set model[/red]")
+                raise typer.Exit(1)
+        else:
+            console.print("[yellow]Operation cancelled[/yellow]")
+
+    asyncio.run(interactive_set())
+
+
+@models_app.command("current")
+def current_model():
+    """Show the currently configured model."""
+    current = model_manager.get_current_model()
+    console.print(f"Current model: [cyan]{current}[/cyan]")
+
+
+@models_app.command("reset")
+def reset_model():
+    """Reset to the default model."""
+
+    async def interactive_reset():
+        if await questionary.confirm(
+            f"Reset to default model ({ModelManager.DEFAULT_MODEL})?"
+        ).ask_async():
+            if model_manager.reset_model():
+                console.print(
+                    f"[green]✓ Model reset to default: {ModelManager.DEFAULT_MODEL}[/green]"
+                )
+            else:
+                console.print("[red]✗ Failed to reset model[/red]")
+                raise typer.Exit(1)
+        else:
+            console.print("[yellow]Operation cancelled[/yellow]")
+
+    asyncio.run(interactive_reset())
+
+
+def create_models_app() -> typer.Typer:
+    """Return the model management CLI app."""
+    return models_app

sqlsaber/cli/streaming.py

@@ -0,0 +1,94 @@
+"""Streaming query handling for the CLI."""
+
+from rich.console import Console
+
+from sqlsaber.agents.base import BaseSQLAgent
+from sqlsaber.cli.display import DisplayManager
+
+
+class StreamingQueryHandler:
+    """Handles streaming query execution and display."""
+
+    def __init__(self, console: Console):
+        self.console = console
+        self.display = DisplayManager(console)
+
+    async def execute_streaming_query(self, user_query: str, agent: BaseSQLAgent):
+        """Execute a query with streaming display."""
+
+        has_content = False
+        explanation_started = False
+        status = self.console.status(
+            "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
+        )
+        status.start()
+
+        try:
+            async for event in agent.query_stream(user_query):
+                if event.type == "tool_use":
+                    # Stop any ongoing status, but don't mark has_content yet
+                    self._stop_status(status)
+
+                    if event.data["status"] == "started":
+                        # If explanation was streaming, add newline before tool use
+                        if explanation_started:
+                            self.display.show_newline()
+                        self.display.show_tool_started(event.data["name"])
+                    elif event.data["status"] == "executing":
+                        self.display.show_tool_executing(
+                            event.data["name"], event.data["input"]
+                        )
+
+                elif event.type == "text":
+                    # Always stop status when text streaming starts
+                    self._stop_status(status)
+
+                    if not explanation_started:
+                        explanation_started = True
+                        has_content = True
+
+                    # Print text as it streams
+                    self.display.show_text_stream(event.data)
+
+                elif event.type == "query_result":
+                    if event.data["results"]:
+                        self.display.show_query_results(event.data["results"])
+
+                elif event.type == "tool_result":
+                    # Handle tool results - particularly list_tables and introspect_schema
+                    if event.data.get("tool_name") == "list_tables":
+                        self.display.show_table_list(event.data["result"])
+                        has_content = True
+                    elif event.data.get("tool_name") == "introspect_schema":
+                        self.display.show_schema_info(event.data["result"])
+                        has_content = True
+
+                elif event.type == "processing":
+                    # Show status when processing tool results
+                    if explanation_started:
+                        self.display.show_newline()  # Add newline after explanation text
+                    self._stop_status(status)
+                    status = self.display.show_processing(event.data)
+                    status.start()
+                    has_content = True
+
+                elif event.type == "error":
+                    if not has_content:
+                        self._stop_status(status)
+                        has_content = True
+                    self.display.show_error(event.data)
+
+        finally:
+            # Make sure status is stopped
+            self._stop_status(status)
+
+            # Add a newline after streaming completes if explanation was shown
+            if explanation_started:
+                self.display.show_newline()  # Empty line for better readability
+
+    def _stop_status(self, status):
+        """Safely stop a status spinner."""
+        try:
+            status.stop()
+        except Exception:
+            pass  # Status might already be stopped

sqlsaber/config/__init__.py

@@ -0,0 +1,7 @@
+"""Configuration module for SQLSaber."""
+
+from .settings import Config
+
+__all__ = [
+    "Config",
+]