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

sqlsaber/cli/threads.py

@@ -0,0 +1,301 @@
+"""Threads CLI: list, show, and resume threads (pydantic-ai message snapshots)."""
+
+import asyncio
+import json
+import time
+from typing import Annotated
+
+import cyclopts
+from pydantic_ai.messages import ModelMessage
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.table import Table
+
+from sqlsaber.agents import build_sqlsaber_agent
+from sqlsaber.cli.display import DisplayManager
+from sqlsaber.cli.interactive import InteractiveSession
+from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.database.connection import DatabaseConnection
+from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
+from sqlsaber.threads import ThreadStorage
+
+# Globals consistent with other CLI modules
+console = Console()
+config_manager = DatabaseConfigManager()
+
+
+threads_app = cyclopts.App(
+    name="threads",
+    help="Manage SQLsaber threads",
+)
+
+
+def _human_readable(timestamp: float | None) -> str:
+    if not timestamp:
+        return "-"
+    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(timestamp))
+
+
+def _render_transcript(
+    console: Console, all_msgs: list[ModelMessage], last_n: int | None = None
+) -> None:
+    """Render conversation turns from ModelMessage[] using DisplayManager."""
+    dm = DisplayManager(console)
+
+    # Locate indices of user prompts
+    user_indices: list[int] = []
+    for idx, message in enumerate(all_msgs):
+        for part in getattr(message, "parts", []):
+            if getattr(part, "part_kind", "") == "user-prompt":
+                user_indices.append(idx)
+                break
+
+    # Build turn slices as (start_idx, end_idx)
+    slices: list[tuple[int, int]] = []
+    if user_indices:
+        for i, start_idx in enumerate(user_indices):
+            end_idx = (
+                user_indices[i + 1] if i + 1 < len(user_indices) else len(all_msgs)
+            )
+            slices.append((start_idx, end_idx))
+
+    if last_n is not None and last_n > 0 and slices:
+        slices = slices[-last_n:]
+
+    def _render_user(message: ModelMessage) -> None:
+        for part in getattr(message, "parts", []):
+            if getattr(part, "part_kind", "") == "user-prompt":
+                content = getattr(part, "content", None)
+                text: str | None = None
+                if isinstance(content, str):
+                    text = content
+                elif isinstance(content, list):  # multimodal
+                    parts: list[str] = []
+                    for seg in content:
+                        if isinstance(seg, str):
+                            parts.append(seg)
+                        else:
+                            try:
+                                parts.append(json.dumps(seg, ensure_ascii=False))
+                            except Exception:
+                                parts.append(str(seg))
+                    text = "\n".join([s for s in parts if s]) or None
+                if text:
+                    console.print(
+                        Panel.fit(Markdown(text), title="User", border_style="cyan")
+                    )
+                    return
+        console.print(Panel.fit("(no content)", title="User", border_style="cyan"))
+
+    def _render_response(message: ModelMessage) -> None:
+        for part in getattr(message, "parts", []):
+            kind = getattr(part, "part_kind", "")
+            if kind == "text":
+                text = getattr(part, "content", "")
+                if isinstance(text, str) and text.strip():
+                    console.print(
+                        Panel.fit(
+                            Markdown(text), title="Assistant", border_style="green"
+                        )
+                    )
+            elif kind in ("tool-call", "builtin-tool-call"):
+                name = getattr(part, "tool_name", "tool")
+                args = getattr(part, "args", None)
+                args_dict: dict = {}
+                if isinstance(args, dict):
+                    args_dict = args
+                elif isinstance(args, str):
+                    try:
+                        parsed = json.loads(args)
+                        if isinstance(parsed, dict):
+                            args_dict = parsed
+                    except Exception:
+                        args_dict = {}
+                dm.show_tool_executing(name, args_dict)
+            elif kind in ("tool-return", "builtin-tool-return"):
+                name = getattr(part, "tool_name", "tool")
+                content = getattr(part, "content", None)
+                if isinstance(content, (dict, list)):
+                    content_str = json.dumps(content, ensure_ascii=False)
+                elif isinstance(content, str):
+                    content_str = content
+                else:
+                    content_str = json.dumps({"return_value": str(content)})
+                if name == "list_tables":
+                    dm.show_table_list(content_str)
+                elif name == "introspect_schema":
+                    dm.show_schema_info(content_str)
+                elif name == "execute_sql":
+                    try:
+                        data = json.loads(content_str)
+                        if (
+                            isinstance(data, dict)
+                            and data.get("success")
+                            and data.get("results")
+                        ):
+                            dm.show_query_results(data["results"])  # type: ignore[arg-type]
+                        else:
+                            console.print(
+                                Panel.fit(
+                                    content_str,
+                                    title=f"Tool result: {name}",
+                                    border_style="yellow",
+                                )
+                            )
+                    except Exception:
+                        console.print(
+                            Panel.fit(
+                                content_str,
+                                title=f"Tool result: {name}",
+                                border_style="yellow",
+                            )
+                        )
+                else:
+                    console.print(
+                        Panel.fit(
+                            content_str,
+                            title=f"Tool result: {name}",
+                            border_style="yellow",
+                        )
+                    )
+        # Thinking parts omitted
+
+    for start_idx, end_idx in slices or [(0, len(all_msgs))]:
+        if start_idx < len(all_msgs):
+            _render_user(all_msgs[start_idx])
+        for i in range(start_idx + 1, end_idx):
+            _render_response(all_msgs[i])
+        console.print("")
+
+
+@threads_app.command(name="list")
+def list_threads(
+    database: Annotated[
+        str | None,
+        cyclopts.Parameter(["--database", "-d"], help="Filter by database name"),
+    ] = None,
+    limit: Annotated[
+        int,
+        cyclopts.Parameter(["--limit", "-n"], help="Max threads to return"),
+    ] = 50,
+):
+    """List threads (optionally filtered by database)."""
+    store = ThreadStorage()
+    threads = asyncio.run(store.list_threads(database_name=database, limit=limit))
+    if not threads:
+        console.print("No threads found.")
+        return
+    table = Table(title="Threads")
+    table.add_column("ID", style="cyan")
+    table.add_column("Database", style="magenta")
+    table.add_column("Title", style="green")
+    table.add_column("Last Activity", style="dim")
+    table.add_column("Model", style="yellow")
+    for t in threads:
+        table.add_row(
+            t.id,
+            t.database_name or "-",
+            (t.title or "-")[:60],
+            _human_readable(getattr(t, "last_activity_at", None)),
+            t.model_name or "-",
+        )
+    console.print(table)
+
+
+@threads_app.command
+def show(
+    thread_id: Annotated[str, cyclopts.Parameter(help="Thread ID")],
+):
+    """Show thread metadata and render the full transcript."""
+    store = ThreadStorage()
+    thread = asyncio.run(store.get_thread(thread_id))
+    if not thread:
+        console.print(f"[red]Thread not found:[/red] {thread_id}")
+        return
+    msgs = asyncio.run(store.get_thread_messages(thread_id))
+    console.print(f"[bold]Thread: {thread.id}[/bold]")
+    console.print("")
+    console.print(f"Database: {thread.database_name}")
+    console.print(f"Title: {thread.title}")
+    console.print(f"Last activity: {_human_readable(thread.last_activity_at)}")
+    console.print(f"Model: {thread.model_name}")
+    console.print("")
+
+    _render_transcript(console, msgs, None)
+
+
+@threads_app.command
+def resume(
+    thread_id: Annotated[str, cyclopts.Parameter(help="Thread ID to resume")],
+    database: Annotated[
+        str | None,
+        cyclopts.Parameter(["--database", "-d"], help="Database name or DSN override"),
+    ] = None,
+):
+    """Render transcript, then resume thread in interactive mode."""
+    store = ThreadStorage()
+
+    async def _run() -> None:
+        thread = await store.get_thread(thread_id)
+        if not thread:
+            console.print(f"[red]Thread not found:[/red] {thread_id}")
+            return
+        db_selector = database or thread.database_name
+        if not db_selector:
+            console.print(
+                "[red]No database specified or stored with this thread.[/red]"
+            )
+            return
+        try:
+            resolved = resolve_database(db_selector, config_manager)
+            connection_string = resolved.connection_string
+            db_name = resolved.name
+        except DatabaseResolutionError as e:
+            console.print(f"[red]Database resolution error:[/red] {e}")
+            return
+
+        db_conn = DatabaseConnection(connection_string)
+        try:
+            agent = build_sqlsaber_agent(db_conn, db_name)
+            history = await store.get_thread_messages(thread_id)
+            console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
+            _render_transcript(console, history, None)
+            session = InteractiveSession(
+                console=console,
+                agent=agent,
+                db_conn=db_conn,
+                database_name=db_name,
+                initial_thread_id=thread_id,
+                initial_history=history,
+            )
+            await session.run()
+        finally:
+            await db_conn.close()
+            console.print("\n[green]Goodbye![/green]")
+
+    asyncio.run(_run())
+
+
+@threads_app.command
+def prune(
+    days: Annotated[
+        int,
+        cyclopts.Parameter(
+            ["--days", "-n"], help="Delete threads older than this many days"
+        ),
+    ] = 30,
+):
+    """Prune old threads by last activity timestamp."""
+    store = ThreadStorage()
+
+    async def _run() -> None:
+        deleted = await store.prune_threads(older_than_days=days)
+        console.print(f"[green]✓ Pruned {deleted} thread(s).[/green]")
+
+    asyncio.run(_run())
+
+
+def create_threads_app() -> cyclopts.App:
+    """Return the threads sub-app (for registration)."""
+    return threads_app

sqlsaber/database/schema.py

@@ -1,7 +1,7 @@
 """Database schema introspection utilities."""
 
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, TypedDict
 
 import aiosqlite
 
@@ -12,7 +12,35 @@ from sqlsaber.database.connection import (
     PostgreSQLConnection,
     SQLiteConnection,
 )
-from sqlsaber.models.types import SchemaInfo
+
+
+class ColumnInfo(TypedDict):
+    """Type definition for column information."""
+
+    data_type: str
+    nullable: bool
+    default: str | None
+    max_length: int | None
+    precision: int | None
+    scale: int | None
+
+
+class ForeignKeyInfo(TypedDict):
+    """Type definition for foreign key information."""
+
+    column: str
+    references: dict[str, str]  # {"table": "schema.table", "column": "column_name"}
+
+
+class SchemaInfo(TypedDict):
+    """Type definition for schema information."""
+
+    schema: str
+    name: str
+    type: str
+    columns: dict[str, ColumnInfo]
+    primary_keys: list[str]
+    foreign_keys: list[ForeignKeyInfo]
 
 
 class BaseSchemaIntrospector(ABC):

sqlsaber/threads/__init__.py

@@ -0,0 +1,5 @@
+"""Thread storage for pydantic-ai message histories."""
+
+from .storage import Thread, ThreadStorage
+
+__all__ = ["Thread", "ThreadStorage"]

sqlsaber/threads/storage.py

@@ -0,0 +1,303 @@
+"""SQLite storage for pydantic-ai thread snapshots.
+
+Each thread represents a session (interactive or non-interactive) and stores the
+complete pydantic-ai message history as a snapshot. On every completed run in the
+same session, we overwrite the snapshot with the new full history.
+
+This design intentionally avoids per-run append logs and mirrors pydantic-ai's
+recommended approach of serializing ModelMessage[] with ModelMessagesTypeAdapter.
+"""
+
+import asyncio
+import logging
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+import aiosqlite
+import platformdirs
+from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
+
+logger = logging.getLogger(__name__)
+
+
+SCHEMA_SQL = """
+CREATE TABLE IF NOT EXISTS threads (
+    id TEXT PRIMARY KEY,
+    database_name TEXT,
+    title TEXT,
+    created_at REAL NOT NULL,
+    ended_at REAL,
+    last_activity_at REAL NOT NULL,
+    model_name TEXT,
+    messages_json BLOB NOT NULL,
+    extra_metadata TEXT
+);
+
+CREATE INDEX IF NOT EXISTS idx_threads_dbname ON threads(database_name);
+CREATE INDEX IF NOT EXISTS idx_threads_activity ON threads(last_activity_at);
+"""
+
+
+class Thread:
+    """Thread metadata."""
+
+    def __init__(
+        self,
+        id: str,
+        database_name: str | None,
+        title: str | None,
+        created_at: float,
+        ended_at: float | None,
+        last_activity_at: float,
+        model_name: str | None,
+    ) -> None:
+        self.id = id
+        self.database_name = database_name
+        self.title = title
+        self.created_at = created_at
+        self.ended_at = ended_at
+        self.last_activity_at = last_activity_at
+        self.model_name = model_name
+
+
+class ThreadStorage:
+    """Handles SQLite storage for pydantic-ai thread snapshots."""
+
+    def __init__(self) -> None:
+        self.db_path = Path(platformdirs.user_config_dir("sqlsaber")) / "threads.db"
+        self._lock = asyncio.Lock()
+        self._initialized = False
+
+    async def _init_db(self) -> None:
+        if self._initialized:
+            return
+        try:
+            self.db_path.parent.mkdir(parents=True, exist_ok=True)
+            async with aiosqlite.connect(self.db_path) as db:
+                await db.executescript(SCHEMA_SQL)
+                await db.commit()
+            self._initialized = True
+            logger.debug("Initialized threads database at %s", self.db_path)
+        except Exception as e:  # pragma: no cover - best-effort persistence
+            logger.warning("Failed to initialize threads DB: %s", e)
+
+    async def save_snapshot(
+        self,
+        *,
+        messages_json: bytes,
+        database_name: str | None,
+        thread_id: str | None = None,
+        extra_metadata: str | None = None,
+    ) -> str:
+        """Create or update a thread snapshot."""
+        await self._init_db()
+        now = time.time()
+
+        if thread_id is None:
+            thread_id = str(uuid.uuid4())
+            try:
+                async with self._lock, aiosqlite.connect(self.db_path) as db:
+                    await db.execute(
+                        """
+                        INSERT INTO threads (
+                            id, database_name, created_at,
+                            last_activity_at, messages_json, extra_metadata
+                        ) VALUES (?, ?, ?, ?, ?, ?)
+                        """,
+                        (
+                            thread_id,
+                            database_name,
+                            now,
+                            now,
+                            messages_json,
+                            extra_metadata,
+                        ),
+                    )
+                    await db.commit()
+                logger.debug("Created thread %s", thread_id)
+                return thread_id
+            except Exception as e:  # pragma: no cover
+                logger.warning("Failed to create thread: %s", e)
+                return thread_id
+        else:
+            try:
+                async with self._lock, aiosqlite.connect(self.db_path) as db:
+                    await db.execute(
+                        """
+                        UPDATE threads
+                        SET last_activity_at = ?,
+                            messages_json = ?,
+                            extra_metadata = COALESCE(?, extra_metadata)
+                        WHERE id = ?
+                        """,
+                        (
+                            now,
+                            messages_json,
+                            extra_metadata,
+                            thread_id,
+                        ),
+                    )
+                    await db.commit()
+                logger.debug("Updated thread %s snapshot", thread_id)
+                return thread_id
+            except Exception as e:  # pragma: no cover
+                logger.warning("Failed to update thread %s: %s", thread_id, e)
+                return thread_id
+
+    async def save_metadata(
+        self,
+        *,
+        thread_id: str,
+        title: str | None = None,
+        model_name: str | None = None,
+    ) -> bool:
+        """Update thread metadata (title/model/extra). Only provided fields are updated."""
+        await self._init_db()
+
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    """
+                    UPDATE threads
+                    SET title = ?, model_name = ?
+                    WHERE id = ?
+                    """,
+                    (title, model_name, thread_id),
+                )
+                await db.commit()
+            return True
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to update metadata for thread %s: %s", thread_id, e)
+            return False
+
+    async def end_thread(self, thread_id: str) -> bool:
+        await self._init_db()
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    "UPDATE threads SET ended_at = ?, last_activity_at = ? WHERE id = ?",
+                    (time.time(), time.time(), thread_id),
+                )
+                await db.commit()
+            return True
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to end thread %s: %s", thread_id, e)
+            return False
+
+    async def get_thread(self, thread_id: str) -> Thread | None:
+        await self._init_db()
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(
+                    """
+                    SELECT id, database_name, title, created_at, ended_at,
+                           last_activity_at, model_name
+                    FROM threads WHERE id = ?
+                    """,
+                    (thread_id,),
+                ) as cur:
+                    row = await cur.fetchone()
+                    if not row:
+                        return None
+                    return Thread(
+                        id=row[0],
+                        database_name=row[1],
+                        title=row[2],
+                        created_at=row[3],
+                        ended_at=row[4],
+                        last_activity_at=row[5],
+                        model_name=row[6],
+                    )
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to get thread %s: %s", thread_id, e)
+            return None
+
+    async def get_thread_messages(self, thread_id: str) -> list[ModelMessage]:
+        """Load the full message history for a thread as ModelMessage[]."""
+        await self._init_db()
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(
+                    "SELECT messages_json FROM threads WHERE id = ?",
+                    (thread_id,),
+                ) as cur:
+                    row = await cur.fetchone()
+                    if not row:
+                        return []
+                    messages_blob: bytes = row[0]
+                    return ModelMessagesTypeAdapter.validate_json(messages_blob)
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to load thread %s messages: %s", thread_id, e)
+            return []
+
+    async def list_threads(
+        self, *, database_name: str | None = None, limit: int = 50
+    ) -> list[Thread]:
+        await self._init_db()
+        try:
+            query = (
+                "SELECT id, database_name, title, created_at, ended_at, last_activity_at, model_name"
+                " FROM threads"
+            )
+            params: list[Any] = []
+            if database_name:
+                query += " WHERE database_name = ?"
+                params.append(database_name)
+            query += " ORDER BY last_activity_at DESC LIMIT ?"
+            params.append(limit)
+
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(query, params) as cur:
+                    threads: list[Thread] = []
+                    async for row in cur:
+                        threads.append(
+                            Thread(
+                                id=row[0],
+                                database_name=row[1],
+                                title=row[2],
+                                created_at=row[3],
+                                ended_at=row[4],
+                                last_activity_at=row[5],
+                                model_name=row[6],
+                            )
+                        )
+                    return threads
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to list threads: %s", e)
+            return []
+
+    async def delete_thread(self, thread_id: str) -> bool:
+        await self._init_db()
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                cur = await db.execute("DELETE FROM threads WHERE id = ?", (thread_id,))
+                await db.commit()
+                return cur.rowcount > 0
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to delete thread %s: %s", thread_id, e)
+            return False
+
+    async def prune_threads(self, older_than_days: int = 30) -> int:
+        """Delete threads whose last_activity_at is older than the cutoff.
+
+        Args:
+            older_than_days: Threads with last_activity_at older than this many days ago will be deleted.
+
+        Returns:
+            The number of rows deleted (best-effort; 0 on failure).
+        """
+        await self._init_db()
+        cutoff = time.time() - older_than_days * 24 * 3600
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                cur = await db.execute(
+                    "DELETE FROM threads WHERE last_activity_at < ?",
+                    (cutoff,),
+                )
+                await db.commit()
+                return cur.rowcount or 0
+        except Exception as e:  # pragma: no cover
+            logger.warning("Failed to prune threads: %s", e)
+            return 0

sqlsaber/tools/__init__.py

@@ -7,7 +7,6 @@ from .registry import ToolRegistry, register_tool, tool_registry
 
 # Import concrete tools to register them
 from .sql_tools import ExecuteSQLTool, IntrospectSchemaTool, ListTablesTool, SQLTool
-from .visualization_tools import PlotDataTool
 
 __all__ = [
     "Tool",
@@ -21,5 +20,4 @@ __all__ = [
     "ListTablesTool",
     "IntrospectSchemaTool",
     "ExecuteSQLTool",
-    "PlotDataTool",
 ]

sqlsaber/tools/base.py

@@ -1,7 +1,6 @@
 """Base class for SQLSaber tools."""
 
 from abc import ABC, abstractmethod
-from types import SimpleNamespace
 from typing import Any
 
 from .enums import ToolCategory, WorkflowPosition
@@ -44,17 +43,6 @@ class Tool(ABC):
         """
         pass
 
-    def to_definition(self):
-        """Convert this tool to a ToolDefinition-like object with attributes.
-
-        Tests expect attribute access (definition.name), so return a SimpleNamespace.
-        """
-        return SimpleNamespace(
-            name=self.name,
-            description=self.description,
-            input_schema=self.input_schema,
-        )
-
     @property
     def category(self) -> ToolCategory:
         """Return the tool category. Override to customize."""