sqlsaber 0.24.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

sqlsaber/cli/display.py

@@ -8,7 +8,7 @@ rendered with Live.
 import json
 from typing import Sequence, Type
 
-from pydantic_ai.messages import ModelResponsePart, TextPart
+from pydantic_ai.messages import ModelResponsePart, TextPart, ThinkingPart
 from rich.columns import Columns
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.live import Live
@@ -75,7 +75,7 @@ class LiveMarkdownRenderer:
             self.end()
             self.paragraph_break()
 
-        self._start()
+        self._start(kind)
         self._current_kind = kind
 
     def append(self, text: str | None) -> None:
@@ -87,7 +87,13 @@ class LiveMarkdownRenderer:
             self.ensure_segment(TextPart)
 
         self._buffer += text
-        self._live.update(Markdown(self._buffer))
+
+        # Apply dim styling for thinking segments
+        if self._current_kind == ThinkingPart:
+            content = Markdown(self._buffer, style="dim")
+            self._live.update(content)
+        else:
+            self._live.update(Markdown(self._buffer))
 
     def end(self) -> None:
         """Finalize and stop the current Live segment, if any."""
@@ -95,13 +101,17 @@ class LiveMarkdownRenderer:
             return
         # Persist the *final* render exactly once, then shut Live down.
         buf = self._buffer
+        kind = self._current_kind
         self._live.stop()
         self._live = None
         self._buffer = ""
         self._current_kind = None
         # Print the complete markdown to scroll-back for permanent reference
         if buf:
-            self.console.print(Markdown(buf))
+            if kind == ThinkingPart:
+                self.console.print(Text(buf, style="dim"))
+            else:
+                self.console.print(Markdown(buf))
 
     def end_if_active(self) -> None:
         self.end()
@@ -153,10 +163,20 @@ class LiveMarkdownRenderer:
         text = Text(f" {message}", style="yellow")
         return Columns([spinner, text], expand=False)
 
-    def _start(self, initial_markdown: str = "") -> None:
+    def _start(
+        self, kind: Type[ModelResponsePart] | None = None, initial_markdown: str = ""
+    ) -> None:
         if self._live is not None:
             self.end()
         self._buffer = initial_markdown or ""
+
+        # Add visual styling for thinking segments
+        if kind == ThinkingPart:
+            if self.console.is_terminal:
+                self.console.print("[dim]💭 Thinking...[/dim]")
+            else:
+                self.console.print("*Thinking...*\n")
+
         # NOTE: Use transient=True so the live widget disappears on exit,
         # giving a clean transition to the final printed result.
         live = Live(
@@ -219,7 +239,9 @@ class DisplayManager:
             if self.console.is_terminal:
                 self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
                 self.show_newline()
-                syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
+                syntax = Syntax(
+                    query, "sql", background_color="default", word_wrap=True
+                )
                 self.console.print(syntax)
             else:
                 self.console.print("**Executing SQL:**\n")
@@ -271,9 +293,7 @@ class DisplayManager:
                     f"[yellow]... and {len(results) - 20} more rows[/yellow]"
                 )
             else:
-                self.console.print(
-                    f"*... and {len(results) - 20} more rows*\n"
-                )
+                self.console.print(f"*... and {len(results) - 20} more rows*\n")
 
     def show_error(self, error_message: str):
         """Display error message."""

sqlsaber/cli/interactive.py

@@ -3,13 +3,13 @@
 import asyncio
 from pathlib import Path
 from textwrap import dedent
+from typing import TYPE_CHECKING
 
 import platformdirs
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.patch_stdout import patch_stdout
 from prompt_toolkit.styles import Style
-from pydantic_ai import Agent
 from rich.console import Console
 from rich.markdown import Markdown
 from rich.panel import Panel
@@ -21,8 +21,9 @@ from sqlsaber.cli.completers import (
 )
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
-from sqlsaber.database.connection import (
+from sqlsaber.database import (
     CSVConnection,
+    DuckDBConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -30,6 +31,9 @@ from sqlsaber.database.connection import (
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.threads import ThreadStorage
 
+if TYPE_CHECKING:
+    from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
+
 
 def bottom_toolbar():
     return [
@@ -54,7 +58,7 @@ class InteractiveSession:
     def __init__(
         self,
         console: Console,
-        agent: Agent,
+        sqlsaber_agent: "SQLSaberAgent",
         db_conn,
         database_name: str,
         *,
@@ -62,7 +66,7 @@ class InteractiveSession:
         initial_history: list | None = None,
     ):
         self.console = console
-        self.agent = agent
+        self.sqlsaber_agent = sqlsaber_agent
         self.db_conn = db_conn
         self.database_name = database_name
         self.display = DisplayManager(console)
@@ -85,8 +89,12 @@ class InteractiveSession:
             if isinstance(self.db_conn, PostgreSQLConnection)
             else "MySQL"
             if isinstance(self.db_conn, MySQLConnection)
+            else "DuckDB"
+            if isinstance(self.db_conn, DuckDBConnection)
+            else "DuckDB"
+            if isinstance(self.db_conn, CSVConnection)
             else "SQLite"
-            if isinstance(self.db_conn, (SQLiteConnection, CSVConnection))
+            if isinstance(self.db_conn, SQLiteConnection)
             else "database"
         )
 
@@ -171,7 +179,7 @@ class InteractiveSession:
         query_task = asyncio.create_task(
             self.streaming_handler.execute_streaming_query(
                 user_query,
-                self.agent,
+                self.sqlsaber_agent,
                 self.cancellation_token,
                 self.message_history,
             )
@@ -186,11 +194,6 @@ class InteractiveSession:
                     # 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(),
@@ -201,8 +204,8 @@ class InteractiveSession:
                     if self.first_message:
                         await self._threads.save_metadata(
                             thread_id=self._thread_id,
-                            title=title,
-                            model_name=model_name,
+                            title=user_query,
+                            model_name=self.sqlsaber_agent.agent.model.model_name,
                         )
                 except Exception:
                     pass
@@ -264,6 +267,17 @@ class InteractiveSession:
                     self._thread_id = None
                     continue
 
+                # Thinking commands
+                if user_query == "/thinking on":
+                    self.sqlsaber_agent.set_thinking(enabled=True)
+                    self.console.print("[green]✓ Thinking enabled[/green]\n")
+                    continue
+
+                if user_query == "/thinking off":
+                    self.sqlsaber_agent.set_thinking(enabled=False)
+                    self.console.print("[green]✓ Thinking disabled[/green]\n")
+                    continue
+
                 if memory_text := user_query.strip():
                     # Check if query starts with # for memory addition
                     if memory_text.startswith("#"):
@@ -271,9 +285,7 @@ class InteractiveSession:
                         if memory_content:
                             # Add memory via the agent's memory manager
                             try:
-                                mm = getattr(
-                                    self.agent, "_sqlsaber_memory_manager", None
-                                )
+                                mm = self.sqlsaber_agent.memory_manager
                                 if mm and self.database_name:
                                     memory = mm.add_memory(
                                         self.database_name, memory_content

sqlsaber/cli/streaming.py

@@ -8,9 +8,9 @@ rendered via DisplayManager helpers.
 import asyncio
 import json
 from functools import singledispatchmethod
-from typing import AsyncIterable
+from typing import TYPE_CHECKING, AsyncIterable
 
-from pydantic_ai import Agent, RunContext
+from pydantic_ai import RunContext
 from pydantic_ai.messages import (
     AgentStreamEvent,
     FunctionToolCallEvent,
@@ -26,6 +26,9 @@ from rich.console import Console
 
 from sqlsaber.cli.display import DisplayManager
 
+if TYPE_CHECKING:
+    from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
+
 
 class StreamingQueryHandler:
     """
@@ -130,7 +133,7 @@ class StreamingQueryHandler:
     async def execute_streaming_query(
         self,
         user_query: str,
-        agent: Agent,
+        sqlsaber_agent: "SQLSaberAgent",
         cancellation_token: asyncio.Event | None = None,
         message_history: list | None = None,
     ):
@@ -139,21 +142,16 @@ class StreamingQueryHandler:
         try:
             # If Anthropic OAuth, inject SQLsaber instructions before the first user prompt
             prepared_prompt: str | list[str] = user_query
-            is_oauth = bool(getattr(agent, "_sqlsaber_is_oauth", False))
             no_history = not message_history
-            if is_oauth and no_history:
-                ib = getattr(agent, "_sqlsaber_instruction_builder", None)
-                mm = getattr(agent, "_sqlsaber_memory_manager", None)
-                db_type = getattr(agent, "_sqlsaber_db_type", "database")
-                db_name = getattr(agent, "_sqlsaber_database_name", None)
-                instructions = (
-                    ib.build_instructions(db_type=db_type) if ib is not None else ""
-                )
-                mem = (
-                    mm.format_memories_for_prompt(db_name)
-                    if (mm is not None and db_name)
-                    else ""
+            if sqlsaber_agent.is_oauth and no_history:
+                instructions = sqlsaber_agent.instruction_builder.build_instructions(
+                    db_type=sqlsaber_agent.db_type
                 )
+                mem = ""
+                if sqlsaber_agent.database_name:
+                    mem = sqlsaber_agent.memory_manager.format_memories_for_prompt(
+                        sqlsaber_agent.database_name
+                    )
                 parts = [p for p in (instructions, mem) if p and str(p).strip()]
                 if parts:
                     injected = "\n\n".join(parts)
@@ -163,7 +161,7 @@ class StreamingQueryHandler:
             self.display.live.start_status("Crunching data...")
 
             # Run the agent with our event stream handler
-            run = await agent.run(
+            run = await sqlsaber_agent.agent.run(
                 prepared_prompt,
                 message_history=message_history,
                 event_stream_handler=self._event_stream_handler,

sqlsaber/cli/threads.py

@@ -148,7 +148,9 @@ def _render_transcript(
                             )
                         else:
                             if is_redirected:
-                                console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                                console.print(
+                                    f"**Tool result ({name}):**\n\n{content_str}\n"
+                                )
                             else:
                                 console.print(
                                     Panel.fit(
@@ -159,7 +161,9 @@ def _render_transcript(
                                 )
                     except Exception:
                         if is_redirected:
-                            console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                            console.print(
+                                f"**Tool result ({name}):**\n\n{content_str}\n"
+                            )
                         else:
                             console.print(
                                 Panel.fit(
@@ -258,10 +262,10 @@ def resume(
 
     async def _run() -> None:
         # Lazy imports to avoid heavy modules at CLI startup
-        from sqlsaber.agents import build_sqlsaber_agent
+        from sqlsaber.agents import SQLSaberAgent
         from sqlsaber.cli.interactive import InteractiveSession
         from sqlsaber.config.database import DatabaseConfigManager
-        from sqlsaber.database.connection import DatabaseConnection
+        from sqlsaber.database import DatabaseConnection
         from sqlsaber.database.resolver import (
             DatabaseResolutionError,
             resolve_database,
@@ -288,7 +292,7 @@ def resume(
 
         db_conn = DatabaseConnection(connection_string)
         try:
-            agent = build_sqlsaber_agent(db_conn, db_name)
+            sqlsaber_agent = SQLSaberAgent(db_conn, db_name)
             history = await store.get_thread_messages(thread_id)
             if console.is_terminal:
                 console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
@@ -297,7 +301,7 @@ def resume(
             _render_transcript(console, history, None)
             session = InteractiveSession(
                 console=console,
-                agent=agent,
+                sqlsaber_agent=sqlsaber_agent,
                 db_conn=db_conn,
                 database_name=db_name,
                 initial_thread_id=thread_id,

sqlsaber/config/database.py

@@ -18,7 +18,7 @@ class DatabaseConfig:
     """Database connection configuration."""
 
     name: str
-    type: str  # postgresql, mysql, sqlite, csv
+    type: str  # postgresql, mysql, sqlite, duckdb, csv
     host: str | None
     port: int | None
     database: str
@@ -90,6 +90,8 @@ class DatabaseConfig:
 
         elif self.type == "sqlite":
             return f"sqlite:///{self.database}"
+        elif self.type == "duckdb":
+            return f"duckdb:///{self.database}"
         elif self.type == "csv":
             # For CSV files, database field contains the file path
             base_url = f"csv:///{self.database}"

sqlsaber/config/settings.py

@@ -46,7 +46,10 @@ class ModelConfigManager:
     def _load_config(self) -> dict[str, Any]:
         """Load configuration from file."""
         if not self.config_file.exists():
-            return {"model": self.DEFAULT_MODEL}
+            return {
+                "model": self.DEFAULT_MODEL,
+                "thinking_enabled": False,
+            }
 
         try:
             with open(self.config_file, "r") as f:
@@ -54,9 +57,15 @@ class ModelConfigManager:
                 # Ensure we have a model set
                 if "model" not in config:
                     config["model"] = self.DEFAULT_MODEL
+                # Set defaults for thinking if not present
+                if "thinking_enabled" not in config:
+                    config["thinking_enabled"] = False
                 return config
         except (json.JSONDecodeError, IOError):
-            return {"model": self.DEFAULT_MODEL}
+            return {
+                "model": self.DEFAULT_MODEL,
+                "thinking_enabled": False,
+            }
 
     def _save_config(self, config: dict[str, Any]) -> None:
         """Save configuration to file."""
@@ -76,6 +85,17 @@ class ModelConfigManager:
         config["model"] = model
         self._save_config(config)
 
+    def get_thinking_enabled(self) -> bool:
+        """Get whether thinking is enabled."""
+        config = self._load_config()
+        return config.get("thinking_enabled", False)
+
+    def set_thinking_enabled(self, enabled: bool) -> None:
+        """Set whether thinking is enabled."""
+        config = self._load_config()
+        config["thinking_enabled"] = enabled
+        self._save_config(config)
+
 
 class Config:
     """Configuration class for SQLSaber."""
@@ -86,6 +106,9 @@ class Config:
         self.api_key_manager = APIKeyManager()
         self.auth_config_manager = AuthConfigManager()
 
+        # Thinking configuration
+        self.thinking_enabled = self.model_config_manager.get_thinking_enabled()
+
         # Authentication method (API key or Anthropic OAuth)
         self.auth_method = self.auth_config_manager.get_auth_method()
 

sqlsaber/database/__init__.py

@@ -1,9 +1,63 @@
 """Database module for SQLSaber."""
 
-from .connection import DatabaseConnection
+from .base import (
+    DEFAULT_QUERY_TIMEOUT,
+    BaseDatabaseConnection,
+    BaseSchemaIntrospector,
+    ColumnInfo,
+    ForeignKeyInfo,
+    IndexInfo,
+    QueryTimeoutError,
+    SchemaInfo,
+)
+from .csv import CSVConnection, CSVSchemaIntrospector
+from .duckdb import DuckDBConnection, DuckDBSchemaIntrospector
+from .mysql import MySQLConnection, MySQLSchemaIntrospector
+from .postgresql import PostgreSQLConnection, PostgreSQLSchemaIntrospector
 from .schema import SchemaManager
+from .sqlite import SQLiteConnection, SQLiteSchemaIntrospector
+
+
+def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
+    """Factory function to create appropriate database connection based on connection string."""
+    if connection_string.startswith("postgresql://"):
+        return PostgreSQLConnection(connection_string)
+    elif connection_string.startswith("mysql://"):
+        return MySQLConnection(connection_string)
+    elif connection_string.startswith("sqlite:///"):
+        return SQLiteConnection(connection_string)
+    elif connection_string.startswith("duckdb://"):
+        return DuckDBConnection(connection_string)
+    elif connection_string.startswith("csv:///"):
+        return CSVConnection(connection_string)
+    else:
+        raise ValueError(
+            f"Unsupported database type in connection string: {connection_string}"
+        )
+
 
 __all__ = [
+    # Base classes and types
+    "BaseDatabaseConnection",
+    "BaseSchemaIntrospector",
+    "ColumnInfo",
+    "DEFAULT_QUERY_TIMEOUT",
+    "ForeignKeyInfo",
+    "IndexInfo",
+    "QueryTimeoutError",
+    "SchemaInfo",
+    # Concrete implementations
+    "PostgreSQLConnection",
+    "MySQLConnection",
+    "SQLiteConnection",
+    "DuckDBConnection",
+    "CSVConnection",
+    "PostgreSQLSchemaIntrospector",
+    "MySQLSchemaIntrospector",
+    "SQLiteSchemaIntrospector",
+    "DuckDBSchemaIntrospector",
+    "CSVSchemaIntrospector",
+    # Factory function and manager
     "DatabaseConnection",
     "SchemaManager",
 ]

sqlsaber/database/base.py

@@ -0,0 +1,124 @@
+"""Base classes and type definitions for database connections and schema introspection."""
+
+from abc import ABC, abstractmethod
+from typing import Any, TypedDict
+
+# Default query timeout to prevent runaway queries
+DEFAULT_QUERY_TIMEOUT = 30.0  # seconds
+
+
+class QueryTimeoutError(RuntimeError):
+    """Exception raised when a query exceeds its timeout."""
+
+    def __init__(self, seconds: float):
+        self.timeout = seconds
+        super().__init__(f"Query exceeded timeout of {seconds}s")
+
+
+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 IndexInfo(TypedDict):
+    """Type definition for index information."""
+
+    name: str
+    columns: list[str]  # ordered
+    unique: bool
+    type: str | None  # btree, gin, FULLTEXT, etc. None if unknown
+
+
+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]
+    indexes: list[IndexInfo]
+
+
+class BaseDatabaseConnection(ABC):
+    """Abstract base class for database connections."""
+
+    def __init__(self, connection_string: str):
+        self.connection_string = connection_string
+        self._pool = None
+
+    @abstractmethod
+    async def get_pool(self):
+        """Get or create connection pool."""
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """Close the connection pool."""
+        pass
+
+    @abstractmethod
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+
+        Args:
+            query: SQL query to execute
+            *args: Query parameters
+            timeout: Query timeout in seconds (overrides default_timeout)
+        """
+        pass
+
+
+class BaseSchemaIntrospector(ABC):
+    """Abstract base class for database-specific schema introspection."""
+
+    @abstractmethod
+    async def get_tables_info(
+        self, connection, table_pattern: str | None = None
+    ) -> dict[str, Any]:
+        """Get tables information for the specific database type."""
+        pass
+
+    @abstractmethod
+    async def get_columns_info(self, connection, tables: list) -> list:
+        """Get columns information for the specific database type."""
+        pass
+
+    @abstractmethod
+    async def get_foreign_keys_info(self, connection, tables: list) -> list:
+        """Get foreign keys information for the specific database type."""
+        pass
+
+    @abstractmethod
+    async def get_primary_keys_info(self, connection, tables: list) -> list:
+        """Get primary keys information for the specific database type."""
+        pass
+
+    @abstractmethod
+    async def get_indexes_info(self, connection, tables: list) -> list:
+        """Get indexes information for the specific database type."""
+        pass
+
+    @abstractmethod
+    async def list_tables_info(self, connection) -> list[dict[str, Any]]:
+        """Get list of tables with basic information."""
+        pass