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

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,7 +21,7 @@ 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,
@@ -31,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 [
@@ -55,7 +58,7 @@ class InteractiveSession:
     def __init__(
         self,
         console: Console,
-        agent: Agent,
+        sqlsaber_agent: "SQLSaberAgent",
         db_conn,
         database_name: str,
         *,
@@ -63,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)
@@ -176,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,
             )
@@ -191,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(),
@@ -206,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
@@ -269,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("#"):
@@ -276,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/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

sqlsaber/database/csv.py

@@ -0,0 +1,133 @@
+"""CSV database connection using DuckDB backend."""
+
+import asyncio
+from pathlib import Path
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+
+import duckdb
+
+from .base import DEFAULT_QUERY_TIMEOUT, BaseDatabaseConnection, QueryTimeoutError
+from .duckdb import DuckDBSchemaIntrospector
+
+
+def _execute_duckdb_transaction(
+    conn: duckdb.DuckDBPyConnection, query: str, args: tuple[Any, ...]
+) -> list[dict[str, Any]]:
+    """Run a DuckDB query inside a transaction and return list of dicts."""
+    conn.execute("BEGIN TRANSACTION")
+    try:
+        if args:
+            conn.execute(query, args)
+        else:
+            conn.execute(query)
+
+        if conn.description is None:
+            rows: list[dict[str, Any]] = []
+        else:
+            columns = [col[0] for col in conn.description]
+            data = conn.fetchall()
+            rows = [dict(zip(columns, row)) for row in data]
+
+        conn.execute("ROLLBACK")
+        return rows
+    except Exception:
+        conn.execute("ROLLBACK")
+        raise
+
+
+class CSVConnection(BaseDatabaseConnection):
+    """CSV file connection using DuckDB per query."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+
+        raw_path = connection_string.replace("csv:///", "", 1)
+        self.csv_path = raw_path.split("?", 1)[0]
+
+        self.delimiter = ","
+        self.encoding = "utf-8"
+        self.has_header = True
+
+        parsed = urlparse(connection_string)
+        if parsed.query:
+            params = parse_qs(parsed.query)
+            self.delimiter = params.get("delimiter", [self.delimiter])[0]
+            self.encoding = params.get("encoding", [self.encoding])[0]
+            self.has_header = params.get("header", ["true"])[0].lower() == "true"
+
+        self.table_name = Path(self.csv_path).stem or "csv_table"
+
+    async def get_pool(self):
+        """CSV connections do not maintain a pool."""
+        return None
+
+    async def close(self):
+        """No persistent resources to close for CSV connections."""
+        pass
+
+    def _quote_identifier(self, identifier: str) -> str:
+        escaped = identifier.replace('"', '""')
+        return f'"{escaped}"'
+
+    def _quote_literal(self, value: str) -> str:
+        escaped = value.replace("'", "''")
+        return f"'{escaped}'"
+
+    def _normalized_encoding(self) -> str | None:
+        encoding = (self.encoding or "").strip()
+        if not encoding or encoding.lower() == "utf-8":
+            return None
+        return encoding.replace("-", "").replace("_", "").upper()
+
+    def _create_view(self, conn: duckdb.DuckDBPyConnection) -> None:
+        header_literal = "TRUE" if self.has_header else "FALSE"
+        option_parts = [f"HEADER={header_literal}"]
+
+        if self.delimiter:
+            option_parts.append(f"DELIM={self._quote_literal(self.delimiter)}")
+
+        encoding = self._normalized_encoding()
+        if encoding:
+            option_parts.append(f"ENCODING={self._quote_literal(encoding)}")
+
+        options_sql = ""
+        if option_parts:
+            options_sql = ", " + ", ".join(option_parts)
+
+        base_relation_sql = (
+            f"read_csv_auto({self._quote_literal(self.csv_path)}{options_sql})"
+        )
+
+        create_view_sql = (
+            f"CREATE VIEW {self._quote_identifier(self.table_name)} AS "
+            f"SELECT * FROM {base_relation_sql}"
+        )
+        conn.execute(create_view_sql)
+
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
+        args_tuple = tuple(args) if args else tuple()
+
+        def _run_query() -> list[dict[str, Any]]:
+            conn = duckdb.connect(":memory:")
+            try:
+                self._create_view(conn)
+                return _execute_duckdb_transaction(conn, query, args_tuple)
+            finally:
+                conn.close()
+
+        try:
+            return await asyncio.wait_for(
+                asyncio.to_thread(_run_query), timeout=effective_timeout
+            )
+        except asyncio.TimeoutError as exc:
+            raise QueryTimeoutError(effective_timeout or 0) from exc
+
+
+class CSVSchemaIntrospector(DuckDBSchemaIntrospector):
+    """CSV-specific schema introspection using DuckDB backend."""
+
+    pass