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

sqlsaber/conversation/manager.py

@@ -0,0 +1,224 @@
+"""Manager for conversation storage operations."""
+
+import logging
+import uuid
+from typing import Any
+
+from .models import Conversation, ConversationMessage
+from .storage import ConversationStorage
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationManager:
+    """High-level manager for conversation storage operations."""
+
+    def __init__(self):
+        """Initialize conversation manager."""
+        self._storage = ConversationStorage()
+
+    async def start_conversation(self, database_name: str) -> str:
+        """Start a new conversation.
+
+        Args:
+            database_name: Name of the database for this conversation
+
+        Returns:
+            Conversation ID
+        """
+        try:
+            return await self._storage.create_conversation(database_name)
+        except Exception as e:
+            logger.warning(f"Failed to start conversation: {e}")
+            return str(uuid.uuid4())
+
+    async def add_user_message(
+        self, conversation_id: str, content: str | dict[str, Any], index: int
+    ) -> str:
+        """Add a user message to the conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            content: Message content
+            index: Sequential index in conversation
+
+        Returns:
+            Message ID
+        """
+        try:
+            return await self._storage.add_message(
+                conversation_id, "user", content, index
+            )
+        except Exception as e:
+            logger.warning(f"Failed to add user message: {e}")
+            return str(uuid.uuid4())
+
+    async def add_assistant_message(
+        self,
+        conversation_id: str,
+        content: list[dict[str, Any]] | dict[str, Any],
+        index: int,
+    ) -> str:
+        """Add an assistant message to the conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            content: Message content (typically ContentBlock list)
+            index: Sequential index in conversation
+
+        Returns:
+            Message ID
+        """
+        try:
+            return await self._storage.add_message(
+                conversation_id, "assistant", content, index
+            )
+        except Exception as e:
+            logger.warning(f"Failed to add assistant message: {e}")
+            return str(uuid.uuid4())
+
+    async def add_tool_message(
+        self,
+        conversation_id: str,
+        content: list[dict[str, Any]] | dict[str, Any],
+        index: int,
+    ) -> str:
+        """Add a tool/system message to the conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            content: Message content (typically tool results)
+            index: Sequential index in conversation
+
+        Returns:
+            Message ID
+        """
+        try:
+            return await self._storage.add_message(
+                conversation_id, "tool", content, index
+            )
+        except Exception as e:
+            logger.warning(f"Failed to add tool message: {e}")
+            return str(uuid.uuid4())
+
+    async def end_conversation(self, conversation_id: str) -> bool:
+        """End a conversation.
+
+        Args:
+            conversation_id: ID of the conversation to end
+
+        Returns:
+            True if successfully ended, False otherwise
+        """
+        try:
+            return await self._storage.end_conversation(conversation_id)
+        except Exception as e:
+            logger.warning(f"Failed to end conversation: {e}")
+            return False
+
+    async def get_conversation(self, conversation_id: str) -> Conversation | None:
+        """Get a conversation by ID.
+
+        Args:
+            conversation_id: ID of the conversation
+
+        Returns:
+            Conversation object or None if not found
+        """
+        try:
+            return await self._storage.get_conversation(conversation_id)
+        except Exception as e:
+            logger.warning(f"Failed to get conversation: {e}")
+            return None
+
+    async def get_conversation_messages(
+        self, conversation_id: str
+    ) -> list[ConversationMessage]:
+        """Get all messages for a conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+
+        Returns:
+            List of messages ordered by index
+        """
+        try:
+            return await self._storage.get_conversation_messages(conversation_id)
+        except Exception as e:
+            logger.warning(f"Failed to get conversation messages: {e}")
+            return []
+
+    async def list_conversations(
+        self, database_name: str | None = None, limit: int = 50
+    ) -> list[Conversation]:
+        """List conversations.
+
+        Args:
+            database_name: Optional database name filter
+            limit: Maximum number of conversations to return
+
+        Returns:
+            List of conversations ordered by start time (newest first)
+        """
+        try:
+            return await self._storage.list_conversations(database_name, limit)
+        except Exception as e:
+            logger.warning(f"Failed to list conversations: {e}")
+            return []
+
+    async def delete_conversation(self, conversation_id: str) -> bool:
+        """Delete a conversation.
+
+        Args:
+            conversation_id: ID of the conversation to delete
+
+        Returns:
+            True if successfully deleted, False otherwise
+        """
+        try:
+            return await self._storage.delete_conversation(conversation_id)
+        except Exception as e:
+            logger.warning(f"Failed to delete conversation: {e}")
+            return False
+
+    async def get_database_names(self) -> list[str]:
+        """Get list of database names with conversations.
+
+        Returns:
+            List of unique database names
+        """
+        try:
+            return await self._storage.get_database_names()
+        except Exception as e:
+            logger.warning(f"Failed to get database names: {e}")
+            return []
+
+    async def restore_conversation_to_agent(
+        self, conversation_id: str, agent_history: list[dict[str, Any]]
+    ) -> bool:
+        """Restore a conversation's messages to an agent's in-memory history.
+
+        Args:
+            conversation_id: ID of the conversation to restore
+            agent_history: Agent's conversation_history list to populate
+
+        Returns:
+            True if successfully restored, False otherwise
+        """
+        try:
+            messages = await self.get_conversation_messages(conversation_id)
+
+            # Clear existing history
+            agent_history.clear()
+
+            # Convert messages back to agent format
+            for msg in messages:
+                if msg.role in ("user", "assistant", "tool"):
+                    agent_history.append({"role": msg.role, "content": msg.content})
+
+            logger.debug(f"Restored {len(messages)} messages to agent history")
+            return True
+
+        except Exception as e:
+            logger.warning(f"Failed to restore conversation to agent: {e}")
+            return False

sqlsaber/conversation/models.py

@@ -0,0 +1,120 @@
+"""Data models for conversation storage."""
+
+import json
+import time
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class Conversation:
+    """Represents a conversation session."""
+
+    id: str
+    database_name: str
+    started_at: float
+    ended_at: float | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "id": self.id,
+            "database_name": self.database_name,
+            "started_at": self.started_at,
+            "ended_at": self.ended_at,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Conversation":
+        """Create from dictionary."""
+        return cls(
+            id=data["id"],
+            database_name=data["database_name"],
+            started_at=data["started_at"],
+            ended_at=data.get("ended_at"),
+        )
+
+    def formatted_start_time(self) -> str:
+        """Get human-readable start timestamp."""
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.started_at))
+
+    def formatted_end_time(self) -> str | None:
+        """Get human-readable end timestamp."""
+        if self.ended_at is None:
+            return None
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.ended_at))
+
+    def duration_seconds(self) -> float | None:
+        """Get conversation duration in seconds."""
+        if self.ended_at is None:
+            return None
+        return self.ended_at - self.started_at
+
+
+@dataclass
+class ConversationMessage:
+    """Represents a single message in a conversation."""
+
+    id: str
+    conversation_id: str
+    role: str
+    content: dict[str, Any] | str
+    index_in_conv: int
+    created_at: float
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "id": self.id,
+            "conversation_id": self.conversation_id,
+            "role": self.role,
+            "content": self.content,
+            "index_in_conv": self.index_in_conv,
+            "created_at": self.created_at,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ConversationMessage":
+        """Create from dictionary."""
+        return cls(
+            id=data["id"],
+            conversation_id=data["conversation_id"],
+            role=data["role"],
+            content=data["content"],
+            index_in_conv=data["index_in_conv"],
+            created_at=data["created_at"],
+        )
+
+    def formatted_timestamp(self) -> str:
+        """Get human-readable timestamp."""
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.created_at))
+
+    def content_json(self) -> str:
+        """Get content as JSON string for storage."""
+        return json.dumps(self.content)
+
+    @classmethod
+    def from_storage_data(
+        cls,
+        id_: str,
+        conversation_id: str,
+        role: str,
+        content_json: str,
+        index_in_conv: int,
+        created_at: float,
+    ) -> "ConversationMessage":
+        """Create from SQLite storage data."""
+        try:
+            content = json.loads(content_json)
+        except json.JSONDecodeError:
+            # Fallback to string content for malformed JSON
+            content = content_json
+
+        return cls(
+            id=id_,
+            conversation_id=conversation_id,
+            role=role,
+            content=content,
+            index_in_conv=index_in_conv,
+            created_at=created_at,
+        )

sqlsaber/conversation/storage.py

@@ -0,0 +1,362 @@
+"""SQLite storage implementation for conversation history."""
+
+import asyncio
+import json
+import logging
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+import aiosqlite
+import platformdirs
+
+from .models import Conversation, ConversationMessage
+
+logger = logging.getLogger(__name__)
+
+# Database schema
+SCHEMA_SQL = """
+-- Conversations table
+CREATE TABLE IF NOT EXISTS conversations (
+    id TEXT PRIMARY KEY,
+    database_name TEXT NOT NULL,
+    started_at REAL NOT NULL,
+    ended_at REAL
+);
+
+-- Messages table
+CREATE TABLE IF NOT EXISTS messages (
+    id TEXT PRIMARY KEY,
+    conversation_id TEXT NOT NULL,
+    role TEXT NOT NULL,
+    content TEXT NOT NULL,
+    index_in_conv INTEGER NOT NULL,
+    created_at REAL NOT NULL,
+    FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
+);
+
+-- Indices for performance
+CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conversation_id, index_in_conv);
+CREATE INDEX IF NOT EXISTS idx_conv_dbname ON conversations(database_name);
+
+-- Store schema version for future migrations
+CREATE TABLE IF NOT EXISTS schema_version (
+    version INTEGER PRIMARY KEY
+);
+
+INSERT OR IGNORE INTO schema_version (version) VALUES (1);
+"""
+
+
+class ConversationStorage:
+    """Handles SQLite storage and retrieval of conversation history."""
+
+    _DB_VERSION = 1
+
+    def __init__(self):
+        """Initialize conversation storage."""
+        self.db_path = (
+            Path(platformdirs.user_config_dir("sqlsaber")) / "conversations.db"
+        )
+        self._lock = asyncio.Lock()
+        self._initialized: bool = False
+
+    async def _init_db(self) -> None:
+        """Initialize the database with schema if needed."""
+        if self._initialized:
+            return
+
+        try:
+            # Ensure parent directory exists
+            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(f"Initialized conversation database at {self.db_path}")
+
+        except Exception as e:
+            logger.warning(f"Failed to initialize conversation database: {e}")
+            # Don't raise - let the system continue without persistence
+
+    async def create_conversation(self, database_name: str) -> str:
+        """Create a new conversation record.
+
+        Args:
+            database_name: Name of the database for this conversation
+
+        Returns:
+            Conversation ID (UUID)
+        """
+        await self._init_db()
+
+        conversation_id = str(uuid.uuid4())
+        started_at = time.time()
+
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    """
+                    INSERT INTO conversations (id, database_name, started_at)
+                    VALUES (?, ?, ?)
+                    """,
+                    (conversation_id, database_name, started_at),
+                )
+                await db.commit()
+
+            logger.debug(
+                f"Created conversation {conversation_id} for db {database_name}"
+            )
+            return conversation_id
+
+        except Exception as e:
+            logger.warning(f"Failed to create conversation: {e}")
+            return conversation_id  # Return ID anyway to allow in-memory operation
+
+    async def add_message(
+        self,
+        conversation_id: str,
+        role: str,
+        content: str | list[dict[str, Any]] | dict[str, Any],
+        index_in_conv: int,
+    ) -> str:
+        """Add a message to a conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            role: Message role (user, assistant, system, tool)
+            content: Message content (will be JSON-serialized)
+            index_in_conv: Sequential index of message in conversation
+
+        Returns:
+            Message ID (UUID)
+        """
+        await self._init_db()
+
+        message_id = str(uuid.uuid4())
+        created_at = time.time()
+
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    """
+                    INSERT INTO messages (id, conversation_id, role, content, index_in_conv, created_at)
+                    VALUES (?, ?, ?, ?, ?, ?)
+                    """,
+                    (
+                        message_id,
+                        conversation_id,
+                        role,
+                        json.dumps(content),
+                        index_in_conv,
+                        created_at,
+                    ),
+                )
+                await db.commit()
+
+            logger.debug(
+                f"Added message {message_id} to conversation {conversation_id}"
+            )
+            return message_id
+
+        except Exception as e:
+            logger.warning(f"Failed to add message: {e}")
+            return message_id  # Return ID anyway
+
+    async def end_conversation(self, conversation_id: str) -> bool:
+        """Mark a conversation as ended.
+
+        Args:
+            conversation_id: ID of the conversation to end
+
+        Returns:
+            True if successfully updated, False otherwise
+        """
+        await self._init_db()
+
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                await db.execute(
+                    "UPDATE conversations SET ended_at = ? WHERE id = ?",
+                    (time.time(), conversation_id),
+                )
+                await db.commit()
+
+            logger.debug(f"Ended conversation {conversation_id}")
+            return True
+
+        except Exception as e:
+            logger.warning(f"Failed to end conversation: {e}")
+            return False
+
+    async def get_conversation(self, conversation_id: str) -> Conversation | None:
+        """Get a conversation by ID.
+
+        Args:
+            conversation_id: ID of the conversation
+
+        Returns:
+            Conversation object or None if not found
+        """
+        await self._init_db()
+
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(
+                    "SELECT id, database_name, started_at, ended_at FROM conversations WHERE id = ?",
+                    (conversation_id,),
+                ) as cursor:
+                    row = await cursor.fetchone()
+                    if not row:
+                        return None
+
+                    return Conversation(
+                        id=row[0],
+                        database_name=row[1],
+                        started_at=row[2],
+                        ended_at=row[3],
+                    )
+
+        except Exception as e:
+            logger.warning(f"Failed to get conversation {conversation_id}: {e}")
+            return None
+
+    async def get_conversation_messages(
+        self, conversation_id: str
+    ) -> list[ConversationMessage]:
+        """Get all messages for a conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+
+        Returns:
+            List of ConversationMessage objects ordered by index
+        """
+        await self._init_db()
+
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(
+                    """
+                    SELECT id, conversation_id, role, content, index_in_conv, created_at
+                    FROM messages
+                    WHERE conversation_id = ?
+                    ORDER BY index_in_conv
+                    """,
+                    (conversation_id,),
+                ) as cursor:
+                    messages = []
+                    async for row in cursor:
+                        message = ConversationMessage.from_storage_data(
+                            id_=row[0],
+                            conversation_id=row[1],
+                            role=row[2],
+                            content_json=row[3],
+                            index_in_conv=row[4],
+                            created_at=row[5],
+                        )
+                        messages.append(message)
+
+                    return messages
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to get messages for conversation {conversation_id}: {e}"
+            )
+            return []
+
+    async def list_conversations(
+        self, database_name: str | None = None, limit: int = 50
+    ) -> list[Conversation]:
+        """List conversations, optionally filtered by database.
+
+        Args:
+            database_name: Optional database name filter
+            limit: Maximum number of conversations to return
+
+        Returns:
+            List of Conversation objects ordered by start time (newest first)
+        """
+        await self._init_db()
+
+        try:
+            query = """
+                SELECT id, database_name, started_at, ended_at
+                FROM conversations
+            """
+            params = []
+
+            if database_name:
+                query += " WHERE database_name = ?"
+                params.append(database_name)
+
+            query += " ORDER BY started_at DESC LIMIT ?"
+            params.append(limit)
+
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(query, params) as cursor:
+                    conversations = []
+                    async for row in cursor:
+                        conversation = Conversation(
+                            id=row[0],
+                            database_name=row[1],
+                            started_at=row[2],
+                            ended_at=row[3],
+                        )
+                        conversations.append(conversation)
+
+                    return conversations
+
+        except Exception as e:
+            logger.warning(f"Failed to list conversations: {e}")
+            return []
+
+    async def delete_conversation(self, conversation_id: str) -> bool:
+        """Delete a conversation and all its messages.
+
+        Args:
+            conversation_id: ID of the conversation to delete
+
+        Returns:
+            True if successfully deleted, False otherwise
+        """
+        await self._init_db()
+
+        try:
+            async with self._lock, aiosqlite.connect(self.db_path) as db:
+                # Delete conversation (messages will be deleted by CASCADE)
+                cursor = await db.execute(
+                    "DELETE FROM conversations WHERE id = ?", (conversation_id,)
+                )
+                await db.commit()
+
+                deleted = cursor.rowcount > 0
+                if deleted:
+                    logger.debug(f"Deleted conversation {conversation_id}")
+                return deleted
+
+        except Exception as e:
+            logger.warning(f"Failed to delete conversation {conversation_id}: {e}")
+            return False
+
+    async def get_database_names(self) -> list[str]:
+        """Get list of all database names that have conversations.
+
+        Returns:
+            List of unique database names
+        """
+        await self._init_db()
+
+        try:
+            async with aiosqlite.connect(self.db_path) as db:
+                async with db.execute(
+                    "SELECT DISTINCT database_name FROM conversations ORDER BY database_name"
+                ) as cursor:
+                    return [row[0] async for row in cursor]
+
+        except Exception as e:
+            logger.warning(f"Failed to get database names: {e}")
+            return []