rp/pr/memory/conversation_memory.py

import json
import sqlite3
import time
from typing import Any, Dict, List, Optional


class ConversationMemory:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._initialize_memory()

    def _initialize_memory(self):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS conversation_history (
                conversation_id TEXT PRIMARY KEY,
                session_id TEXT,
                started_at REAL NOT NULL,
                ended_at REAL,
                message_count INTEGER DEFAULT 0,
                summary TEXT,
                topics TEXT,
                metadata TEXT
            )
        """
        )

        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS conversation_messages (
                message_id TEXT PRIMARY KEY,
                conversation_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp REAL NOT NULL,
                tool_calls TEXT,
                metadata TEXT,
                FOREIGN KEY (conversation_id) REFERENCES conversation_history(conversation_id)
            )
        """
        )

        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_conv_session ON conversation_history(session_id)
        """
        )
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_conv_started ON conversation_history(started_at DESC)
        """
        )
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_msg_conversation ON conversation_messages(conversation_id)
        """
        )
        cursor.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_msg_timestamp ON conversation_messages(timestamp)
        """
        )

        conn.commit()
        conn.close()

    def create_conversation(
        self,
        conversation_id: str,
        session_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT INTO conversation_history
            (conversation_id, session_id, started_at, metadata)
            VALUES (?, ?, ?, ?)
        """,
            (
                conversation_id,
                session_id,
                time.time(),
                json.dumps(metadata) if metadata else None,
            ),
        )

        conn.commit()
        conn.close()

    def add_message(
        self,
        conversation_id: str,
        message_id: str,
        role: str,
        content: str,
        tool_calls: Optional[List[Dict[str, Any]]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT INTO conversation_messages
            (message_id, conversation_id, role, content, timestamp, tool_calls, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
            (
                message_id,
                conversation_id,
                role,
                content,
                time.time(),
                json.dumps(tool_calls) if tool_calls else None,
                json.dumps(metadata) if metadata else None,
            ),
        )

        cursor.execute(
            """
            UPDATE conversation_history
            SET message_count = message_count + 1
            WHERE conversation_id = ?
        """,
            (conversation_id,),
        )

        conn.commit()
        conn.close()

    def get_conversation_messages(
        self, conversation_id: str, limit: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        if limit:
            cursor.execute(
                """
                SELECT message_id, role, content, timestamp, tool_calls, metadata
                FROM conversation_messages
                WHERE conversation_id = ?
                ORDER BY timestamp DESC
                LIMIT ?
            """,
                (conversation_id, limit),
            )
        else:
            cursor.execute(
                """
                SELECT message_id, role, content, timestamp, tool_calls, metadata
                FROM conversation_messages
                WHERE conversation_id = ?
                ORDER BY timestamp ASC
            """,
                (conversation_id,),
            )

        messages = []
        for row in cursor.fetchall():
            messages.append(
                {
                    "message_id": row[0],
                    "role": row[1],
                    "content": row[2],
                    "timestamp": row[3],
                    "tool_calls": json.loads(row[4]) if row[4] else None,
                    "metadata": json.loads(row[5]) if row[5] else None,
                }
            )

        conn.close()
        return messages

    def update_conversation_summary(
        self, conversation_id: str, summary: str, topics: Optional[List[str]] = None
    ):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            """
            UPDATE conversation_history
            SET summary = ?, topics = ?, ended_at = ?
            WHERE conversation_id = ?
        """,
            (
                summary,
                json.dumps(topics) if topics else None,
                time.time(),
                conversation_id,
            ),
        )

        conn.commit()
        conn.close()

    def search_conversations(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            """
            SELECT DISTINCT h.conversation_id, h.session_id, h.started_at,
                   h.message_count, h.summary, h.topics
            FROM conversation_history h
            LEFT JOIN conversation_messages m ON h.conversation_id = m.conversation_id
            WHERE h.summary LIKE ? OR h.topics LIKE ? OR m.content LIKE ?
            ORDER BY h.started_at DESC
            LIMIT ?
        """,
            (f"%{query}%", f"%{query}%", f"%{query}%", limit),
        )

        conversations = []
        for row in cursor.fetchall():
            conversations.append(
                {
                    "conversation_id": row[0],
                    "session_id": row[1],
                    "started_at": row[2],
                    "message_count": row[3],
                    "summary": row[4],
                    "topics": json.loads(row[5]) if row[5] else [],
                }
            )

        conn.close()
        return conversations

    def get_recent_conversations(
        self, limit: int = 10, session_id: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        if session_id:
            cursor.execute(
                """
                SELECT conversation_id, session_id, started_at, ended_at,
                       message_count, summary, topics
                FROM conversation_history
                WHERE session_id = ?
                ORDER BY started_at DESC
                LIMIT ?
            """,
                (session_id, limit),
            )
        else:
            cursor.execute(
                """
                SELECT conversation_id, session_id, started_at, ended_at,
                       message_count, summary, topics
                FROM conversation_history
                ORDER BY started_at DESC
                LIMIT ?
            """,
                (limit,),
            )

        conversations = []
        for row in cursor.fetchall():
            conversations.append(
                {
                    "conversation_id": row[0],
                    "session_id": row[1],
                    "started_at": row[2],
                    "ended_at": row[3],
                    "message_count": row[4],
                    "summary": row[5],
                    "topics": json.loads(row[6]) if row[6] else [],
                }
            )

        conn.close()
        return conversations

    def delete_conversation(self, conversation_id: str) -> bool:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute(
            "DELETE FROM conversation_messages WHERE conversation_id = ?",
            (conversation_id,),
        )
        cursor.execute(
            "DELETE FROM conversation_history WHERE conversation_id = ?",
            (conversation_id,),
        )

        deleted = cursor.rowcount > 0
        conn.commit()
        conn.close()

        return deleted

    def get_statistics(self) -> Dict[str, Any]:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        cursor = conn.cursor()

        cursor.execute("SELECT COUNT(*) FROM conversation_history")
        total_conversations = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM conversation_messages")
        total_messages = cursor.fetchone()[0]

        cursor.execute("SELECT SUM(message_count) FROM conversation_history")
        cursor.fetchone()[0] or 0

        cursor.execute(
            """
            SELECT AVG(message_count) FROM conversation_history WHERE message_count > 0
        """
        )
        avg_messages = cursor.fetchone()[0] or 0

        conn.close()

        return {
            "total_conversations": total_conversations,
            "total_messages": total_messages,
            "average_messages_per_conversation": round(avg_messages, 2),
        }