rakam-systems-agent 0.1.1rc7__py3-none-any.whl

rakam_systems_agent/components/chat_history/sql_chat_history.py

@@ -0,0 +1,446 @@
+"""SQL-Based Chat History Manager.
+
+This module provides a ChatHistoryComponent implementation that stores
+chat history in a SQLite database. Suitable for production deployments
+requiring persistent, structured storage.
+"""
+
+from __future__ import annotations
+import json
+import os
+import sqlite3
+from typing import Any, Dict, List, Optional
+
+from rakam_systems_core.ai_core.interfaces.chat_history import ChatHistoryComponent
+
+# Optional pydantic-ai integration
+try:
+    from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelMessage
+    from pydantic_core import to_jsonable_python
+    PYDANTIC_AI_AVAILABLE = True
+except ImportError:
+    PYDANTIC_AI_AVAILABLE = False
+    ModelMessagesTypeAdapter = None  # type: ignore
+    ModelMessage = None  # type: ignore
+    to_jsonable_python = None  # type: ignore
+
+
+class SQLChatHistory(ChatHistoryComponent):
+    """Chat history manager using SQLite database storage.
+
+    This implementation stores all chat histories in a SQLite database.
+    It's suitable for:
+    - Production deployments
+    - Multi-instance applications (with proper connection handling)
+    - Applications requiring structured queries
+    - Medium to large scale applications
+
+    Config options:
+        db_path: Path to the SQLite database file (default: "./chat_history.db")
+
+    Example:
+        >>> history = SQLChatHistory(config={"db_path": "./data/chats.db"})
+        >>> history.add_message("chat123", {"role": "user", "content": "Hello"})
+        >>> history.add_message("chat123", {"role": "assistant", "content": "Hi there!"})
+        >>> messages = history.get_chat_history("chat123")
+    """
+
+    def __init__(
+        self,
+        name: str = "sql_chat_history",
+        config: Optional[Dict[str, Any]] = None,
+        db_path: Optional[str] = None,
+    ) -> None:
+        """Initialize the SQL chat history manager.
+
+        Args:
+            name: Component name for identification.
+            config: Configuration dictionary. Supports:
+                - db_path: Path to SQLite database file
+            db_path: Direct path override (takes precedence over config).
+        """
+        super().__init__(name, config)
+
+        # Get db path from argument, config, or default
+        self.db_path = db_path or self.config.get(
+            "db_path", "./chat_history.db"
+        )
+
+    def setup(self) -> None:
+        """Initialize database and create tables."""
+        self._initialize_database()
+        super().setup()
+
+    def shutdown(self) -> None:
+        """Cleanup resources."""
+        super().shutdown()
+
+    def _get_connection(self) -> sqlite3.Connection:
+        """Get a database connection with foreign keys enabled.
+
+        Returns:
+            SQLite connection object.
+        """
+        conn = sqlite3.connect(self.db_path)
+        conn.execute('PRAGMA foreign_keys = ON;')
+        return conn
+
+    def _initialize_database(self) -> None:
+        """Initialize SQLite database and create necessary tables.
+
+        Creates the chats and messages tables if they don't exist.
+
+        Raises:
+            Exception: If database initialization fails.
+        """
+        # Ensure directory exists
+        db_dir = os.path.dirname(self.db_path)
+        if db_dir:
+            os.makedirs(db_dir, exist_ok=True)
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+
+            # Create tables
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS chats (
+                    chat_id TEXT PRIMARY KEY
+                )
+            ''')
+
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    chat_id TEXT NOT NULL,
+                    message_order INTEGER NOT NULL,
+                    message_data TEXT NOT NULL,
+                    FOREIGN KEY (chat_id) REFERENCES chats (chat_id) ON DELETE CASCADE
+                )
+            ''')
+
+            # Create index for faster lookups
+            cursor.execute('''
+                CREATE INDEX IF NOT EXISTS idx_messages_chat_id 
+                ON messages (chat_id, message_order)
+            ''')
+
+            conn.commit()
+
+    def _ensure_initialized(self) -> None:
+        """Ensure the component is initialized before operations."""
+        if not self.initialized:
+            self.setup()
+
+    def add_message(self, chat_id: str, message: Dict[str, Any]) -> None:
+        """Add a single message to a chat session.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+            message: Message object (dict with role, content, timestamp, etc.).
+        """
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+
+            # Ensure chat exists
+            cursor.execute(
+                'INSERT OR IGNORE INTO chats (chat_id) VALUES (?)',
+                (chat_id,)
+            )
+
+            # Get next message order
+            cursor.execute(
+                'SELECT COALESCE(MAX(message_order), -1) + 1 FROM messages WHERE chat_id = ?',
+                (chat_id,)
+            )
+            next_order = cursor.fetchone()[0]
+
+            # Insert message
+            message_json = json.dumps(message, ensure_ascii=False)
+            cursor.execute(
+                '''
+                INSERT INTO messages (chat_id, message_order, message_data)
+                VALUES (?, ?, ?)
+                ''',
+                (chat_id, next_order, message_json)
+            )
+
+            conn.commit()
+
+    def set_messages(self, chat_id: str, messages: List[Dict[str, Any]]) -> None:
+        """Set/replace all messages for a chat session.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+            messages: List of message objects to store.
+        """
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+
+            # Ensure chat exists
+            cursor.execute(
+                'INSERT OR IGNORE INTO chats (chat_id) VALUES (?)',
+                (chat_id,)
+            )
+
+            # Delete existing messages
+            cursor.execute(
+                'DELETE FROM messages WHERE chat_id = ?', (chat_id,))
+
+            # Insert new messages with order
+            for order, message in enumerate(messages):
+                message_json = json.dumps(message, ensure_ascii=False)
+                cursor.execute(
+                    '''
+                    INSERT INTO messages (chat_id, message_order, message_data)
+                    VALUES (?, ?, ?)
+                    ''',
+                    (chat_id, order, message_json)
+                )
+
+            conn.commit()
+
+    def get_chat_history(self, chat_id: str) -> List[Dict[str, Any]]:
+        """Retrieve all messages for a chat session.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+
+        Returns:
+            List of message objects, or empty list if chat doesn't exist.
+        """
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''
+                SELECT message_data
+                FROM messages
+                WHERE chat_id = ?
+                ORDER BY message_order ASC
+                ''',
+                (chat_id,)
+            )
+            rows = cursor.fetchall()
+
+            return [json.loads(row[0]) for row in rows]
+
+    def get_all_chat_ids(self) -> List[str]:
+        """Get all chat IDs currently stored.
+
+        Returns:
+            List of all chat session identifiers.
+        """
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute('SELECT chat_id FROM chats')
+            return [row[0] for row in cursor.fetchall()]
+
+    def delete_chat_history(self, chat_id: str) -> bool:
+        """Delete all messages for a chat session.
+
+        Args:
+            chat_id: Unique identifier for the chat session to delete.
+
+        Returns:
+            True if deletion was successful, False if chat_id didn't exist.
+        """
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+
+            # Check if chat exists
+            cursor.execute('SELECT 1 FROM chats WHERE chat_id = ?', (chat_id,))
+            exists = cursor.fetchone() is not None
+
+            if not exists:
+                return False
+
+            # Delete messages (cascades from foreign key, but explicit is safer)
+            cursor.execute(
+                'DELETE FROM messages WHERE chat_id = ?', (chat_id,))
+            cursor.execute('DELETE FROM chats WHERE chat_id = ?', (chat_id,))
+
+            conn.commit()
+            return True
+
+    def clear_all(self) -> None:
+        """Delete all chat histories."""
+        self._ensure_initialized()
+
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute('DELETE FROM messages')
+            cursor.execute('DELETE FROM chats')
+            conn.commit()
+
+    def get_readable_chat_history(
+        self,
+        chat_id: str,
+        user_role: str = "user",
+        assistant_role: str = "assistant",
+    ) -> List[Dict[str, Any]]:
+        """Get chat history in a human-readable format.
+
+        This method transforms the raw message format into a display-friendly
+        format with 'from', 'message', and optional 'timestamp' keys.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+            user_role: The role name for user messages (default: "user").
+            assistant_role: The role name for assistant messages (default: "assistant").
+
+        Returns:
+            List of formatted message dictionaries with:
+                - 'from': "user" or "assistant"
+                - 'message': The message content
+                - 'timestamp': Message timestamp (if available)
+        """
+        self._ensure_initialized()
+
+        messages = self.get_chat_history(chat_id)
+        readable_messages = []
+
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            timestamp = msg.get("timestamp")
+
+            # Determine the 'from' field based on role
+            if role == user_role:
+                from_field = "user"
+            elif role == assistant_role:
+                from_field = "assistant"
+            else:
+                # Skip system messages or unknown roles
+                continue
+
+            formatted = {
+                "from": from_field,
+                "message": content,
+            }
+
+            if timestamp:
+                formatted["timestamp"] = timestamp
+
+            readable_messages.append(formatted)
+
+        return readable_messages
+
+    # ==================== Pydantic-AI Integration ====================
+
+    def get_message_history(self, chat_id: str) -> Optional[List[Any]]:
+        """Get chat history in pydantic-ai compatible format.
+
+        This method converts the stored database history to pydantic-ai's
+        ModelMessage format, ready to be passed to agent.run() or 
+        agent.run_stream() as message_history.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+
+        Returns:
+            List of ModelMessage objects for pydantic-ai, or None if:
+                - Chat doesn't exist or is empty
+                - pydantic-ai is not installed
+
+        Example:
+            >>> history = SQLChatHistory()
+            >>> message_history = history.get_message_history("chat123")
+            >>> result = await agent.run("Hello", message_history=message_history)
+        """
+        if not PYDANTIC_AI_AVAILABLE:
+            raise ImportError(
+                "pydantic-ai is not installed. Install it with: pip install pydantic-ai"
+            )
+
+        self._ensure_initialized()
+        raw_history = self.get_chat_history(chat_id)
+
+        if not raw_history:
+            return None
+
+        return ModelMessagesTypeAdapter.validate_python(raw_history)
+
+    def save_messages(self, chat_id: str, messages: List[Any]) -> None:
+        """Save pydantic-ai messages to history.
+
+        This method converts pydantic-ai's ModelMessage objects to JSON
+        and stores them. Typically called with result.all_messages() after
+        an agent run.
+
+        Args:
+            chat_id: Unique identifier for the chat session.
+            messages: List of pydantic-ai ModelMessage objects 
+                     (e.g., from result.all_messages()).
+
+        Example:
+            >>> result = await agent.run("Hello", message_history=history.get_message_history("chat123"))
+            >>> history.save_messages("chat123", result.all_messages())
+        """
+        if not PYDANTIC_AI_AVAILABLE:
+            raise ImportError(
+                "pydantic-ai is not installed. Install it with: pip install pydantic-ai"
+            )
+
+        self._ensure_initialized()
+
+        # Convert pydantic-ai messages to JSON-serializable format
+        json_messages = to_jsonable_python(messages)
+        self.set_messages(chat_id, json_messages)
+
+
+if __name__ == "__main__":
+    import tempfile
+
+    # Create a temporary database for testing
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "test_chat_history.db")
+
+        # Example usage
+        history = SQLChatHistory(
+            config={"db_path": db_path}
+        )
+
+        # Add messages
+        history.add_message("chat123", {
+            "role": "user",
+            "content": "Hello!",
+            "timestamp": "2025-03-18 10:00:00"
+        })
+        history.add_message("chat123", {
+            "role": "assistant",
+            "content": "Hi there! How can I help?",
+            "timestamp": "2025-03-18 10:00:05"
+        })
+
+        # Retrieve history
+        print("Chat history:", history.get_chat_history("chat123"))
+        print("All chat IDs:", history.get_all_chat_ids())
+        print("Readable format:", history.get_readable_chat_history("chat123"))
+
+        # Test set_messages
+        history.set_messages("chat456", [
+            {"role": "user", "content": "Test message",
+                "timestamp": "2025-03-18 10:02:00"}
+        ])
+        print("Chat history for chat456:", history.get_chat_history("chat456"))
+        print("All chat IDs after adding chat456:", history.get_all_chat_ids())
+
+        # Clean up
+        deleted = history.delete_chat_history("chat123")
+        print(f"Deleted chat123: {deleted}")
+        print("After deletion:", history.get_all_chat_ids())
+
+        # Clear all
+        history.clear_all()
+        print("After clear_all:", history.get_all_chat_ids())
+
+        print("\nAll tests passed!")