sqlsaber 0.13.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

sqlsaber/agents/anthropic.py

@@ -21,6 +21,8 @@ from sqlsaber.config.settings import Config
 from sqlsaber.database.connection import BaseDatabaseConnection
 from sqlsaber.memory.manager import MemoryManager
 from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import tool_registry
+from sqlsaber.tools.instructions import InstructionBuilder
 
 
 class AnthropicSQLAgent(BaseSQLAgent):
@@ -51,89 +53,11 @@ class AnthropicSQLAgent(BaseSQLAgent):
         self._last_results = None
         self._last_query = None
 
-        # Define tools in the new format
-        self.tools: list[ToolDefinition] = [
-            ToolDefinition(
-                name="list_tables",
-                description="Get a list of all tables in the database with row counts. Use this first to discover available tables.",
-                input_schema={
-                    "type": "object",
-                    "properties": {},
-                    "required": [],
-                },
-            ),
-            ToolDefinition(
-                name="introspect_schema",
-                description="Introspect database schema to understand table structures.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "table_pattern": {
-                            "type": "string",
-                            "description": "Optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%')",
-                        }
-                    },
-                    "required": [],
-                },
-            ),
-            ToolDefinition(
-                name="execute_sql",
-                description="Execute a SQL query against the database.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "query": {
-                            "type": "string",
-                            "description": "SQL query to execute",
-                        },
-                        "limit": {
-                            "type": "integer",
-                            "description": f"Maximum number of rows to return (default: {AnthropicSQLAgent.DEFAULT_SQL_LIMIT})",
-                            "default": AnthropicSQLAgent.DEFAULT_SQL_LIMIT,
-                        },
-                    },
-                    "required": ["query"],
-                },
-            ),
-            ToolDefinition(
-                name="plot_data",
-                description="Create a plot of query results.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "y_values": {
-                            "type": "array",
-                            "items": {"type": ["number", "null"]},
-                            "description": "Y-axis data points (required)",
-                        },
-                        "x_values": {
-                            "type": "array",
-                            "items": {"type": ["number", "null"]},
-                            "description": "X-axis data points (optional, will use indices if not provided)",
-                        },
-                        "plot_type": {
-                            "type": "string",
-                            "enum": ["line", "scatter", "histogram"],
-                            "description": "Type of plot to create (default: line)",
-                            "default": "line",
-                        },
-                        "title": {
-                            "type": "string",
-                            "description": "Title for the plot",
-                        },
-                        "x_label": {
-                            "type": "string",
-                            "description": "Label for X-axis",
-                        },
-                        "y_label": {
-                            "type": "string",
-                            "description": "Label for Y-axis",
-                        },
-                    },
-                    "required": ["y_values"],
-                },
-            ),
-        ]
+        # Get tool definitions from registry
+        self.tools: list[ToolDefinition] = tool_registry.get_tool_definitions()
+
+        # Initialize instruction builder
+        self.instruction_builder = InstructionBuilder(tool_registry)
 
         # Build system prompt with memories if available
         self.system_prompt = self._build_system_prompt()
@@ -157,31 +81,9 @@ class AnthropicSQLAgent(BaseSQLAgent):
     def _get_sql_assistant_instructions(self) -> str:
         """Get the detailed SQL assistant instructions."""
         db_type = self._get_database_type_name()
-        instructions = f"""You are also a helpful SQL assistant that helps users query their {db_type} database.
-
-Your responsibilities:
-1. Understand user's natural language requests, think and convert them to SQL
-2. Use the provided tools efficiently to explore database schema
-3. Generate appropriate SQL queries
-4. Execute queries safely - queries that modify the database are not allowed
-5. Format and explain results clearly
-6. Create visualizations when requested or when they would be helpful
-
-IMPORTANT - Schema Discovery Strategy:
-1. ALWAYS start with 'list_tables' to see available tables and row counts
-2. Based on the user's query, identify which specific tables are relevant
-3. Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables
-4. Timestamp columns must be converted to text when you write queries
-
-Guidelines:
-- Use list_tables first, then introspect_schema for specific tables only
-- Use table patterns like 'sample%' or '%experiment%' to filter related tables
-- Use proper JOIN syntax and avoid cartesian products
-- Include appropriate WHERE clauses to limit results
-- Explain what the query does in simple terms
-- Handle errors gracefully and suggest fixes
-- Be security conscious - use parameterized queries when needed
-"""
+
+        # Build dynamic instructions from available tools
+        instructions = self.instruction_builder.build_instructions(db_type=db_type)
 
         # Add memory context if database name is available
         if self.database_name:
@@ -189,7 +91,7 @@ Guidelines:
                 self.database_name
             )
             if memory_context.strip():
-                instructions += memory_context
+                instructions += "\n\n" + memory_context
 
         return instructions
 
@@ -199,16 +101,19 @@ Guidelines:
             return None
 
         memory = self.memory_manager.add_memory(self.database_name, content)
-        # Rebuild system prompt with new memory
+        # Rebuild system prompt with new memory (includes dynamic instructions)
         self.system_prompt = self._build_system_prompt()
         return memory.id
 
-    async def execute_sql(self, query: str, limit: int | None = None) -> str:
-        """Execute a SQL query against the database with streaming support."""
-        # Call parent implementation for core functionality
-        result = await super().execute_sql(query, limit)
+    async def _execute_sql_with_tracking(
+        self, query: str, limit: int | None = None
+    ) -> str:
+        """Execute SQL and track results for streaming."""
+        # Get the execute_sql tool and run it
+        tool = tool_registry.get_tool("execute_sql")
+        result = await tool.execute(query=query, limit=limit)
 
-        # Parse result to extract data for streaming (AnthropicSQLAgent specific)
+        # Parse result to extract data for streaming
         try:
             result_data = json.loads(result)
             if result_data.get("success") and "results" in result_data:
@@ -228,7 +133,14 @@ Guidelines:
         self, tool_name: str, tool_input: dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        # Use parent implementation for core tools
+        # Special handling for execute_sql to track results
+        if tool_name == "execute_sql":
+            return await self._execute_sql_with_tracking(
+                tool_input.get("query", ""),
+                tool_input.get("limit", self.DEFAULT_SQL_LIMIT),
+            )
+
+        # Use parent implementation for all other tools
         return await super().process_tool_call(tool_name, tool_input)
 
     def _convert_user_message_to_message(
@@ -450,6 +362,16 @@ Guidelines:
         self._last_query = None
 
         try:
+            # Ensure conversation is active for persistence
+            await self._ensure_conversation()
+
+            # Store user message in conversation history and persistence
+            if use_history:
+                self.conversation_history.append(
+                    {"role": "user", "content": user_query}
+                )
+                await self._store_user_message(user_query)
+
             # Build messages with history if requested
             messages = []
             if use_history:
@@ -461,8 +383,9 @@ Guidelines:
                 instructions = self._get_sql_assistant_instructions()
                 messages.append(Message(MessageRole.USER, instructions))
 
-            # Add current user message
-            messages.append(Message(MessageRole.USER, user_query))
+            # Add current user message if not already in messages from history
+            if not use_history:
+                messages.append(Message(MessageRole.USER, user_query))
 
             # Create initial request and get response
             request = self._create_message_request(messages)
@@ -484,9 +407,12 @@ Guidelines:
                     return
 
                 # Add assistant's response to conversation
-                collected_content.append(
-                    {"role": "assistant", "content": response.content}
-                )
+                assistant_content = {"role": "assistant", "content": response.content}
+                collected_content.append(assistant_content)
+
+                # Store the assistant message immediately (not from collected_content)
+                if use_history:
+                    await self._store_assistant_message(response.content)
 
                 # Execute tools and get results
                 tool_results = []
@@ -499,9 +425,19 @@ Guidelines:
                         tool_results = event
 
                 # Continue conversation with tool results
-                collected_content.append({"role": "user", "content": tool_results})
+                tool_content = {"role": "user", "content": tool_results}
+                collected_content.append(tool_content)
+
+                # Store the tool message immediately and update history
                 if use_history:
-                    self.conversation_history.extend(collected_content)
+                    # Only add the NEW messages to history (not the accumulated ones)
+                    # collected_content has [assistant1, tool1, assistant2, tool2, ...]
+                    # We only want to add the last 2 items that were just added
+                    new_messages_for_history = collected_content[
+                        -2:
+                    ]  # Last assistant + tool pair
+                    self.conversation_history.extend(new_messages_for_history)
+                    await self._store_tool_message(tool_results)
 
                 if cancellation_token is not None and cancellation_token.is_set():
                     return
@@ -541,6 +477,10 @@ Guidelines:
                     {"role": "assistant", "content": response.content}
                 )
 
+                # Store final assistant message in persistence (only if not tool_use)
+                if response.stop_reason != "tool_use":
+                    await self._store_assistant_message(response.content)
+
         except asyncio.CancelledError:
             return
         except Exception as e:

sqlsaber/agents/base.py

@@ -5,8 +5,7 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator
 
-from uniplot import histogram, plot
-
+from sqlsaber.conversation.manager import ConversationManager
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
     CSVConnection,
@@ -16,6 +15,7 @@ from sqlsaber.database.connection import (
 )
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import SQLTool, tool_registry
 
 
 class BaseSQLAgent(ABC):
@@ -26,6 +26,14 @@ class BaseSQLAgent(ABC):
         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()
+
     @abstractmethod
     async def query_stream(
         self,
@@ -42,8 +50,12 @@ class BaseSQLAgent(ABC):
         """
         pass
 
-    def clear_history(self):
+    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:
@@ -59,228 +71,117 @@ class BaseSQLAgent(ABC):
         else:
             return "database"  # Fallback
 
-    async def introspect_schema(self, table_pattern: str | None = None) -> str:
-        """Introspect database schema to understand table structures."""
-        try:
-            # Pass table_pattern to get_schema_info for efficient filtering at DB level
-            schema_info = await self.schema_manager.get_schema_info(table_pattern)
-
-            # Format the schema information
-            formatted_info = {}
-            for table_name, table_info in schema_info.items():
-                formatted_info[table_name] = {
-                    "columns": {
-                        col_name: {
-                            "type": col_info["data_type"],
-                            "nullable": col_info["nullable"],
-                            "default": col_info["default"],
-                        }
-                        for col_name, col_info in table_info["columns"].items()
-                    },
-                    "primary_keys": table_info["primary_keys"],
-                    "foreign_keys": [
-                        f"{fk['column']} -> {fk['references']['table']}.{fk['references']['column']}"
-                        for fk in table_info["foreign_keys"]
-                    ],
-                }
-
-            return json.dumps(formatted_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
-
-    async def list_tables(self) -> str:
-        """List all tables in the database with basic information."""
-        try:
-            tables_info = await self.schema_manager.list_tables()
-            return json.dumps(tables_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error listing tables: {str(e)}"})
-
-    async def execute_sql(self, query: str, limit: int | None = None) -> str:
-        """Execute a SQL query against the database."""
-        try:
-            # Security check - only allow SELECT queries unless write is enabled
-            write_error = self._validate_write_operation(query)
-            if write_error:
-                return json.dumps(
-                    {
-                        "error": write_error,
-                    }
-                )
-
-            # Add LIMIT if not present and it's a SELECT query
-            query = self._add_limit_to_query(query, limit)
-
-            # Execute the query (wrapped in a transaction for safety)
-            results = await self.db.execute_query(query)
-
-            # Format results
-            actual_limit = limit if limit is not None else len(results)
-
-            return json.dumps(
-                {
-                    "success": True,
-                    "row_count": len(results),
-                    "results": results[:actual_limit],  # Extra safety for limit
-                    "truncated": len(results) > actual_limit,
-                }
-            )
-
-        except Exception as e:
-            error_msg = str(e)
-
-            # Provide helpful error messages
-            suggestions = []
-            if "column" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check column names using the schema introspection tool"
-                )
-            elif "table" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check table names using the schema introspection tool"
-                )
-            elif "syntax error" in error_msg.lower():
-                suggestions.append(
-                    "Review SQL syntax, especially JOIN conditions and WHERE clauses"
-                )
-
-            return json.dumps({"error": error_msg, "suggestions": suggestions})
+    def _init_tools(self) -> None:
+        """Initialize SQL tools with database connection."""
+        # Get all SQL tools and set their database connection
+        for tool_name in tool_registry.list_tools(category="sql"):
+            tool = tool_registry.get_tool(tool_name)
+            if isinstance(tool, SQLTool):
+                tool.set_connection(self.db)
 
     async def process_tool_call(
         self, tool_name: str, tool_input: dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        if tool_name == "list_tables":
-            return await self.list_tables()
-        elif tool_name == "introspect_schema":
-            return await self.introspect_schema(tool_input.get("table_pattern"))
-        elif tool_name == "execute_sql":
-            return await self.execute_sql(
-                tool_input["query"], tool_input.get("limit", 100)
-            )
-        elif tool_name == "plot_data":
-            return await self.plot_data(
-                y_values=tool_input["y_values"],
-                x_values=tool_input.get("x_values"),
-                plot_type=tool_input.get("plot_type", "line"),
-                title=tool_input.get("title"),
-                x_label=tool_input.get("x_label"),
-                y_label=tool_input.get("y_label"),
-            )
-        else:
+        try:
+            tool = tool_registry.get_tool(tool_name)
+            return await tool.execute(**tool_input)
+        except KeyError:
             return json.dumps({"error": f"Unknown tool: {tool_name}"})
+        except Exception as e:
+            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.
 
-    def _validate_write_operation(self, query: str) -> str | None:
-        """Validate if a write operation is allowed.
+        Args:
+            conversation_id: ID of the conversation to restore
 
         Returns:
-            None if operation is allowed, error message if not allowed.
+            True if successfully restored, False otherwise
         """
-        query_upper = query.strip().upper()
-
-        # Check for write operations
-        write_keywords = [
-            "INSERT",
-            "UPDATE",
-            "DELETE",
-            "DROP",
-            "CREATE",
-            "ALTER",
-            "TRUNCATE",
-        ]
-        is_write_query = any(query_upper.startswith(kw) for kw in write_keywords)
-
-        if is_write_query:
-            return (
-                "Write operations are not allowed. Only SELECT queries are permitted."
-            )
+        success = await self._conv_manager.restore_conversation_to_agent(
+            conversation_id, self.conversation_history
+        )
 
-        return None
+        if success:
+            # Set up for continuing this conversation
+            self._conversation_id = conversation_id
+            self._msg_index = len(self.conversation_history)
 
-    def _add_limit_to_query(self, query: str, limit: int = 100) -> str:
-        """Add LIMIT clause to SELECT queries if not present."""
-        query_upper = query.strip().upper()
-        if query_upper.startswith("SELECT") and "LIMIT" not in query_upper:
-            return f"{query.rstrip(';')} LIMIT {limit};"
-        return query
+        return success
 
-    async def plot_data(
-        self,
-        y_values: list[float],
-        x_values: list[float] | None = None,
-        plot_type: str = "line",
-        title: str | None = None,
-        x_label: str | None = None,
-        y_label: str | None = None,
-    ) -> str:
-        """Create a terminal plot using uniplot.
+    async def list_conversations(self, limit: int = 50) -> list:
+        """List conversations for this agent's database.
 
         Args:
-            y_values: Y-axis data points
-            x_values: X-axis data points (optional)
-            plot_type: Type of plot - "line", "scatter", or "histogram"
-            title: Plot title
-            x_label: X-axis label
-            y_label: Y-axis label
+            limit: Maximum number of conversations to return
 
         Returns:
-            JSON string with success status and plot details
+            List of conversation data
         """
-        try:
-            # Validate inputs
-            if not y_values:
-                return json.dumps({"error": "No data provided for plotting"})
-
-            # Convert to floats if needed
-            try:
-                y_values = [float(v) if v is not None else None for v in y_values]
-                if x_values:
-                    x_values = [float(v) if v is not None else None for v in x_values]
-            except (ValueError, TypeError) as e:
-                return json.dumps({"error": f"Invalid data format: {str(e)}"})
-
-            # Create the plot
-            if plot_type == "histogram":
-                # For histogram, we only need y_values
-                histogram(
-                    y_values,
-                    title=title,
-                    bins=min(20, len(set(y_values))),  # Adaptive bin count
-                )
-                plot_info = {
-                    "type": "histogram",
-                    "data_points": len(y_values),
-                    "title": title or "Histogram",
-                }
-            elif plot_type in ["line", "scatter"]:
-                # For line/scatter plots
-                plot_kwargs = {
-                    "ys": y_values,
-                    "title": title,
-                    "lines": plot_type == "line",
-                }
-
-                if x_values:
-                    plot_kwargs["xs"] = x_values
-                if x_label:
-                    plot_kwargs["x_unit"] = x_label
-                if y_label:
-                    plot_kwargs["y_unit"] = y_label
-
-                plot(**plot_kwargs)
-
-                plot_info = {
-                    "type": plot_type,
-                    "data_points": len(y_values),
-                    "title": title or f"{plot_type.capitalize()} Plot",
-                    "has_x_values": x_values is not None,
-                }
-            else:
-                return json.dumps({"error": f"Unsupported plot type: {plot_type}"})
-
-            return json.dumps(
-                {"success": True, "plot_rendered": True, "plot_info": plot_info}
-            )
-
-        except Exception as e:
-            return json.dumps({"error": f"Error creating plot: {str(e)}"})
+        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/interactive.py

@@ -136,11 +136,15 @@ class InteractiveSession:
                 if not user_query:
                     continue
 
-                if user_query in ["/exit", "/quit"]:
+                if (
+                    user_query in ["/exit", "/quit"]
+                    or user_query.startswith("/exit")
+                    or user_query.startswith("/quit")
+                ):
                     break
 
                 if user_query == "/clear":
-                    self.agent.clear_history()
+                    await self.agent.clear_history()
                     self.console.print("[green]Conversation history cleared.[/green]\n")
                     continue
 

sqlsaber/conversation/__init__.py

@@ -0,0 +1,12 @@
+"""Conversation history storage for SQLSaber."""
+
+from .manager import ConversationManager
+from .models import Conversation, ConversationMessage
+from .storage import ConversationStorage
+
+__all__ = [
+    "Conversation",
+    "ConversationMessage",
+    "ConversationStorage",
+    "ConversationManager",
+]