mem0/embedchain/memory/base.py

import json
import logging
import sqlite3
import uuid
from typing import Any, Optional

from embedchain.constants import SQLITE_PATH
from embedchain.memory.message import ChatMessage
from embedchain.memory.utils import merge_metadata_dict

CHAT_MESSAGE_CREATE_TABLE_QUERY = """
    CREATE TABLE IF NOT EXISTS ec_chat_history (
        app_id TEXT,
        id TEXT,
        session_id TEXT,
        question TEXT,
        answer TEXT,
        metadata TEXT,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (id, app_id, session_id)
    )
"""


class ChatHistory:
    def __init__(self) -> None:
        with sqlite3.connect(SQLITE_PATH, check_same_thread=False) as self.connection:
            self.cursor = self.connection.cursor()
            self.cursor.execute(CHAT_MESSAGE_CREATE_TABLE_QUERY)
            self.connection.commit()

    def add(self, app_id, session_id, chat_message: ChatMessage) -> Optional[str]:
        memory_id = str(uuid.uuid4())
        metadata_dict = merge_metadata_dict(chat_message.human_message.metadata, chat_message.ai_message.metadata)
        if metadata_dict:
            metadata = self._serialize_json(metadata_dict)
        ADD_CHAT_MESSAGE_QUERY = """
            INSERT INTO ec_chat_history (app_id, id, session_id, question, answer, metadata)
            VALUES (?, ?, ?, ?, ?, ?)
        """
        self.cursor.execute(
            ADD_CHAT_MESSAGE_QUERY,
            (
                app_id,
                memory_id,
                session_id,
                chat_message.human_message.content,
                chat_message.ai_message.content,
                metadata if metadata_dict else "{}",
            ),
        )
        self.connection.commit()
        logging.info(f"Added chat memory to db with id: {memory_id}")
        return memory_id

    def delete(self, app_id: str, session_id: Optional[str] = None):
        """
        Delete all chat history for a given app_id and session_id.
        This is useful for deleting chat history for a given user.

        :param app_id: The app_id to delete chat history for
        :param session_id: The session_id to delete chat history for

        :return: None
        """
        if session_id:
            DELETE_CHAT_HISTORY_QUERY = "DELETE FROM ec_chat_history WHERE app_id=? AND session_id=?"
            params = (app_id, session_id)
        else:
            DELETE_CHAT_HISTORY_QUERY = "DELETE FROM ec_chat_history WHERE app_id=?"
            params = (app_id,)

        self.cursor.execute(DELETE_CHAT_HISTORY_QUERY, params)
        self.connection.commit()

    def get(self, app_id, session_id, num_rounds=10, display_format=False) -> list[ChatMessage]:
        """
        Get the most recent num_rounds rounds of conversations
        between human and AI, for a given app_id.
        """

        QUERY = """
            SELECT * FROM ec_chat_history
            WHERE app_id=? AND session_id=?
            ORDER BY created_at DESC
            LIMIT ?
        """
        self.cursor.execute(
            QUERY,
            (app_id, session_id, num_rounds),
        )

        results = self.cursor.fetchall()
        history = []
        for result in results:
            app_id, _, session_id, question, answer, metadata, timestamp = result
            metadata = self._deserialize_json(metadata=metadata)
            # Return list of dict if display_format is True
            if display_format:
                history.append({"human": question, "ai": answer, "metadata": metadata, "timestamp": timestamp})
            else:
                memory = ChatMessage()
                memory.add_user_message(question, metadata=metadata)
                memory.add_ai_message(answer, metadata=metadata)
                history.append(memory)
        return history

    def count(self, app_id: str, session_id: Optional[str] = None):
        """
        Count the number of chat messages for a given app_id and session_id.

        :param app_id: The app_id to count chat history for
        :param session_id: The session_id to count chat history for

        :return: The number of chat messages for a given app_id and session_id
        """
        if session_id:
            QUERY = "SELECT COUNT(*) FROM ec_chat_history WHERE app_id=? AND session_id=?"
            params = (app_id, session_id)
        else:
            QUERY = "SELECT COUNT(*) FROM ec_chat_history WHERE app_id=?"
            params = (app_id,)

        self.cursor.execute(QUERY, params)
        count = self.cursor.fetchone()[0]
        return count

    @staticmethod
    def _serialize_json(metadata: dict[str, Any]):
        return json.dumps(metadata)

    @staticmethod
    def _deserialize_json(metadata: str):
        return json.loads(metadata)

    def close_connection(self):
        self.connection.close()