sqlsaber 0.16.1__py3-none-any.whl → 0.18.0__py3-none-any.whl

sqlsaber/agents/base.py

@@ -5,7 +5,6 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator
 
-from sqlsaber.conversation.manager import ConversationManager
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
     CSVConnection,
@@ -23,12 +22,6 @@ class BaseSQLAgent(ABC):
     def __init__(self, db_connection: BaseDatabaseConnection):
         self.db = db_connection
         self.schema_manager = SchemaManager(db_connection)
-        self.conversation_history: list[dict[str, Any]] = []
-
-        # Conversation persistence
-        self._conv_manager = ConversationManager()
-        self._conversation_id: str | None = None
-        self._msg_index: int = 0
 
         # Initialize SQL tools with database connection
         self._init_tools()
@@ -49,14 +42,6 @@ class BaseSQLAgent(ABC):
         """
         pass
 
-    async def clear_history(self):
-        """Clear conversation history."""
-        # End current conversation in storage
-        await self._end_conversation()
-
-        # Clear in-memory history
-        self.conversation_history = []
-
     def _get_database_type_name(self) -> str:
         """Get the human-readable database type name."""
         if isinstance(self.db, PostgreSQLConnection):
@@ -91,96 +76,3 @@ class BaseSQLAgent(ABC):
             return json.dumps(
                 {"error": f"Error executing tool '{tool_name}': {str(e)}"}
             )
-
-    # Conversation persistence helpers
-
-    async def _ensure_conversation(self) -> None:
-        """Ensure a conversation is active for storing messages."""
-        if self._conversation_id is None:
-            db_name = getattr(self, "database_name", "unknown")
-            self._conversation_id = await self._conv_manager.start_conversation(db_name)
-            self._msg_index = 0
-
-    async def _store_user_message(self, content: str | dict[str, Any]) -> None:
-        """Store a user message in conversation history."""
-        if self._conversation_id is None:
-            return
-
-        await self._conv_manager.add_user_message(
-            self._conversation_id, content, self._msg_index
-        )
-        self._msg_index += 1
-
-    async def _store_assistant_message(
-        self, content: list[dict[str, Any]] | dict[str, Any]
-    ) -> None:
-        """Store an assistant message in conversation history."""
-        if self._conversation_id is None:
-            return
-
-        await self._conv_manager.add_assistant_message(
-            self._conversation_id, content, self._msg_index
-        )
-        self._msg_index += 1
-
-    async def _store_tool_message(
-        self, content: list[dict[str, Any]] | dict[str, Any]
-    ) -> None:
-        """Store a tool/system message in conversation history."""
-        if self._conversation_id is None:
-            return
-
-        await self._conv_manager.add_tool_message(
-            self._conversation_id, content, self._msg_index
-        )
-        self._msg_index += 1
-
-    async def _end_conversation(self) -> None:
-        """End the current conversation."""
-        if self._conversation_id:
-            await self._conv_manager.end_conversation(self._conversation_id)
-        self._conversation_id = None
-        self._msg_index = 0
-
-    async def restore_conversation(self, conversation_id: str) -> bool:
-        """Restore a conversation from storage to in-memory history.
-
-        Args:
-            conversation_id: ID of the conversation to restore
-
-        Returns:
-            True if successfully restored, False otherwise
-        """
-        success = await self._conv_manager.restore_conversation_to_agent(
-            conversation_id, self.conversation_history
-        )
-
-        if success:
-            # Set up for continuing this conversation
-            self._conversation_id = conversation_id
-            self._msg_index = len(self.conversation_history)
-
-        return success
-
-    async def list_conversations(self, limit: int = 50) -> list:
-        """List conversations for this agent's database.
-
-        Args:
-            limit: Maximum number of conversations to return
-
-        Returns:
-            List of conversation data
-        """
-        db_name = getattr(self, "database_name", None)
-        conversations = await self._conv_manager.list_conversations(db_name, limit)
-
-        return [
-            {
-                "id": conv.id,
-                "database_name": conv.database_name,
-                "started_at": conv.formatted_start_time(),
-                "ended_at": conv.formatted_end_time(),
-                "duration": conv.duration_seconds(),
-            }
-            for conv in conversations
-        ]

sqlsaber/cli/commands.py

@@ -7,22 +7,8 @@ from typing import Annotated
 import cyclopts
 from rich.console import Console
 
-from sqlsaber.agents import build_sqlsaber_agent
-from sqlsaber.cli.auth import create_auth_app
-from sqlsaber.cli.database import create_db_app
-from sqlsaber.cli.interactive import InteractiveSession
-from sqlsaber.cli.memory import create_memory_app
-from sqlsaber.cli.models import create_models_app
-from sqlsaber.cli.streaming import StreamingQueryHandler
+# Lazy imports - only import what's needed for CLI parsing
 from sqlsaber.config.database import DatabaseConfigManager
-from sqlsaber.database.connection import (
-    CSVConnection,
-    DatabaseConnection,
-    MySQLConnection,
-    PostgreSQLConnection,
-    SQLiteConnection,
-)
-from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
 
 
 class CLIError(Exception):
@@ -106,6 +92,21 @@ def query(
     """
 
     async def run_session():
+        # Import heavy dependencies only when actually running a query
+        # This is only done to speed up startup time
+        from sqlsaber.agents import build_sqlsaber_agent
+        from sqlsaber.cli.interactive import InteractiveSession
+        from sqlsaber.cli.streaming import StreamingQueryHandler
+        from sqlsaber.database.connection import (
+            CSVConnection,
+            DatabaseConnection,
+            MySQLConnection,
+            PostgreSQLConnection,
+            SQLiteConnection,
+        )
+        from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
+        from sqlsaber.threads import ThreadStorage
+
         # Check if query_text is None and stdin has data
         actual_query = query_text
         if query_text is None and not sys.stdin.isatty():
@@ -149,7 +150,33 @@ def query(
                 console.print(
                     f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
                 )
-                await streaming_handler.execute_streaming_query(actual_query, agent)
+                run = await streaming_handler.execute_streaming_query(
+                    actual_query, agent
+                )
+                # Persist non-interactive run as a thread snapshot so it can be resumed later
+                try:
+                    if run is not None:
+                        threads = ThreadStorage()
+                        # Extract title and model name
+                        title = actual_query
+                        model_name: str | None = agent.model.model_name
+
+                        thread_id = await threads.save_snapshot(
+                            messages_json=run.all_messages_json(),
+                            database_name=db_name,
+                        )
+                        await threads.save_metadata(
+                            thread_id=thread_id, title=title, model_name=model_name
+                        )
+                        await threads.end_thread(thread_id)
+                        console.print(
+                            f"[dim]You can continue this thread using:[/dim] saber threads resume {thread_id}"
+                        )
+                except Exception:
+                    # best-effort persistence; don't fail the CLI on storage errors
+                    pass
+                finally:
+                    await threads.prune_threads()
             else:
                 # Interactive mode
                 session = InteractiveSession(console, agent, db_conn, db_name)
@@ -168,21 +195,45 @@ def query(
         sys.exit(e.exit_code)
 
 
-# Add authentication management commands
-auth_app = create_auth_app()
-app.command(auth_app, name="auth")
+# Use lazy imports for fast CLI startup time
+@app.command(name="auth")
+def auth(*args, **kwargs):
+    """Manage authentication configuration."""
+    from sqlsaber.cli.auth import create_auth_app
+
+    return create_auth_app()(*args, **kwargs)
+
+
+@app.command(name="db")
+def db(*args, **kwargs):
+    """Manage database connections."""
+    from sqlsaber.cli.database import create_db_app
+
+    return create_db_app()(*args, **kwargs)
+
+
+@app.command(name="memory")
+def memory(*args, **kwargs):
+    """Manage database-specific memories."""
+    from sqlsaber.cli.memory import create_memory_app
+
+    return create_memory_app()(*args, **kwargs)
+
+
+@app.command(name="models")
+def models(*args, **kwargs):
+    """Select and manage models."""
+    from sqlsaber.cli.models import create_models_app
+
+    return create_models_app()(*args, **kwargs)
 
-# Add database management commands after main callback is defined
-db_app = create_db_app()
-app.command(db_app, name="db")
 
-# Add memory management commands
-memory_app = create_memory_app()
-app.command(memory_app, name="memory")
+@app.command(name="threads")
+def threads(*args, **kwargs):
+    """Manage SQLsaber threads."""
+    from sqlsaber.cli.threads import create_threads_app
 
-# Add model management commands
-models_app = create_models_app()
-app.command(models_app, name="models")
+    return create_threads_app()(*args, **kwargs)
 
 
 def main():

sqlsaber/cli/display.py

@@ -209,36 +209,6 @@ class DisplayManager:
         except Exception as e:
             self.show_error(f"Error displaying schema information: {str(e)}")
 
-    def show_plot(self, plot_data: dict):
-        """Display plot information and status."""
-        try:
-            # Parse the result if it's a string
-            if isinstance(plot_data.get("result"), str):
-                result = json.loads(plot_data["result"])
-            else:
-                result = plot_data.get("result", {})
-
-            # Check if there was an error
-            if "error" in result:
-                self.show_error(f"Plot error: {result['error']}")
-                return
-
-            # If plot was successful, show plot info
-            if result.get("success") and result.get("plot_rendered"):
-                plot_info = result.get("plot_info", {})
-                self.console.print(
-                    f"\n[bold green]✓ Plot rendered:[/bold green] {plot_info.get('title', 'Plot')}"
-                )
-                self.console.print(
-                    f"[dim]  Type: {plot_info.get('type', 'unknown')}, "
-                    f"Data points: {plot_info.get('data_points', 0)}[/dim]"
-                )
-
-        except json.JSONDecodeError:
-            self.show_error("Failed to parse plot result")
-        except Exception as e:
-            self.show_error(f"Error displaying plot: {str(e)}")
-
     def show_markdown_response(self, content: list):
         """Display the assistant's response as rich markdown in a panel."""
         if not content:

sqlsaber/cli/interactive.py

@@ -14,13 +14,29 @@ from sqlsaber.cli.completers import (
 )
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
+from sqlsaber.database.connection import (
+    CSVConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
 from sqlsaber.database.schema import SchemaManager
+from sqlsaber.threads import ThreadStorage
 
 
 class InteractiveSession:
     """Manages interactive CLI sessions."""
 
-    def __init__(self, console: Console, agent: Agent, db_conn, database_name: str):
+    def __init__(
+        self,
+        console: Console,
+        agent: Agent,
+        db_conn,
+        database_name: str,
+        *,
+        initial_thread_id: str | None = None,
+        initial_history: list | None = None,
+    ):
         self.console = console
         self.agent = agent
         self.db_conn = db_conn
@@ -30,19 +46,16 @@ class InteractiveSession:
         self.current_task: asyncio.Task | None = None
         self.cancellation_token: asyncio.Event | None = None
         self.table_completer = TableNameCompleter()
-        self.message_history: list | None = []
+        self.message_history: list | None = initial_history or []
+        # Conversation Thread persistence
+        self._threads = ThreadStorage()
+        self._thread_id: str | None = initial_thread_id
+        self.first_message = not self._thread_id
 
     def show_welcome_message(self):
         """Display welcome message for interactive mode."""
         # Show database information
         db_name = self.database_name or "Unknown"
-        from sqlsaber.database.connection import (
-            CSVConnection,
-            MySQLConnection,
-            PostgreSQLConnection,
-            SQLiteConnection,
-        )
-
         db_type = (
             "PostgreSQL"
             if isinstance(self.db_conn, PostgreSQLConnection)
@@ -53,32 +66,36 @@ class InteractiveSession:
             else "database"
         )
 
-        self.console.print(
-            Panel.fit(
-                """
+        if self.first_message:
+            self.console.print(
+                Panel.fit(
+                    """
 ███████  ██████  ██      ███████  █████  ██████  ███████ ██████
 ██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
 ███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
      ██ ██ ▄▄ ██ ██           ██ ██   ██ ██   ██ ██      ██   ██
 ███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
             ▀▀
-"""
+    """
+                )
+            )
+            self.console.print(
+                "\n",
+                "[dim] > Use '/clear' to reset conversation",
+                "[dim] > Use '/exit' or '/quit' to leave[/dim]",
+                "[dim] > Use 'Ctrl+C' to interrupt and return to prompt\n\n",
+                "[dim] > Start message with '#' to add something to agent's memory for this database",
+                "[dim] > Type '@' to get table name completions",
+                "[dim] > Press 'Esc-Enter' or 'Meta-Enter' to submit your question",
+                sep="\n",
             )
-        )
-        self.console.print(
-            "\n",
-            "[dim] ≥ Use '/clear' to reset conversation",
-            "[dim] ≥ Use '/exit' or '/quit' to leave[/dim]",
-            "[dim] ≥ Use 'Ctrl+C' to interrupt and return to prompt\n\n",
-            "[dim] ≥ Start message with '#' to add something to agent's memory for this database",
-            "[dim] ≥ Type '@' to get table name completions",
-            "[dim] ≥ Press 'Esc-Enter' or 'Meta-Enter' to submit your question",
-            sep="\n",
-        )
 
         self.console.print(
             f"[bold blue]\n\nConnected to:[/bold blue] {db_name} ({db_type})\n"
         )
+        # If resuming a thread, show a notice
+        if self._thread_id:
+            self.console.print(f"[dim]Resuming thread:[/dim] {self._thread_id}\n")
 
     async def _update_table_cache(self):
         """Update the table completer cache with fresh data."""
@@ -132,8 +149,29 @@ class InteractiveSession:
                 try:
                     # Use all_messages() so the system prompt and all prior turns are preserved
                     self.message_history = run_result.all_messages()
+
+                    # Extract title (first user prompt) and model name
+                    if not self._thread_id:
+                        title = user_query
+                        model_name = self.agent.model.model_name
+
+                    # Persist snapshot to thread storage (create or overwrite)
+                    self._thread_id = await self._threads.save_snapshot(
+                        messages_json=run_result.all_messages_json(),
+                        database_name=self.database_name,
+                        thread_id=self._thread_id,
+                    )
+                    # Save metadata separately (only if its the first message)
+                    if self.first_message:
+                        await self._threads.save_metadata(
+                            thread_id=self._thread_id,
+                            title=title,
+                            model_name=model_name,
+                        )
                 except Exception:
                     pass
+                finally:
+                    await self._threads.prune_threads()
         finally:
             self.current_task = None
             self.cancellation_token = None
@@ -165,12 +203,26 @@ class InteractiveSession:
                     or user_query.startswith("/exit")
                     or user_query.startswith("/quit")
                 ):
+                    # Print resume hint if there is an active thread
+                    if self._thread_id:
+                        await self._threads.end_thread(self._thread_id)
+                        self.console.print(
+                            f"[dim]You can continue this thread using:[/dim] saber threads resume {self._thread_id}"
+                        )
                     break
 
                 if user_query == "/clear":
                     # Reset local history (pydantic-ai call will receive empty history on next run)
                     self.message_history = []
+                    # End current thread (if any) so the next turn creates a fresh one
+                    try:
+                        if self._thread_id:
+                            await self._threads.end_thread(self._thread_id)
+                    except Exception:
+                        pass
                     self.console.print("[green]Conversation history cleared.[/green]\n")
+                    # Do not print resume hint when clearing; a new thread will be created on next turn
+                    self._thread_id = None
                     continue
 
                 if memory_text := user_query.strip():

sqlsaber/cli/streaming.py

@@ -72,10 +72,6 @@ class StreamingQueryHandler:
                     except json.JSONDecodeError:
                         # If not JSON, ignore here
                         pass
-                elif tool_name == "plot_data":
-                    self.display.show_plot(
-                        {"tool_name": tool_name, "result": content, "input": {}}
-                    )
 
     async def execute_streaming_query(
         self,