mysqlchatstore.py

from typing import Optional, Any
from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from pydantic import Field
import pymysql

from llama_index.core.storage.chat_store import BaseChatStore
from llama_index.core.llms import ChatMessage
from llama_index.core.memory import ChatMemoryBuffer



class MySQLChatStore(BaseChatStore):
    """
    Implementação de um ChatStore que armazena mensagens em uma tabela MySQL,
    unindo a pergunta do usuário e a resposta do assistente na mesma linha.
    """
    table_name: Optional[str] = Field(default="chatstore", description="Nome da tabela MySQL.")

    _session: Optional[sessionmaker] = None
    _async_session: Optional[sessionmaker] = None

    def __init__(self, session: sessionmaker, async_session: sessionmaker, table_name: str):
        super().__init__(table_name=table_name.lower())
        self._session = session
        self._async_session = async_session
        self._initialize()

    @classmethod
    def from_params(cls, host: str, port: str, database: str, user: str, password: str, table_name: str = "chatstore") -> "MySQLChatStore":
        """
        Cria o sessionmaker síncrono e assíncrono, retornando a instância da classe.
        """
        conn_str = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
        async_conn_str = f"mysql+aiomysql://{user}:{password}@{host}:{port}/{database}"
        session, async_session = cls._connect(conn_str, async_conn_str)
        return cls(session=session, async_session=async_session, table_name=table_name)

    @classmethod
    def _connect(cls, connection_string: str, async_connection_string: str) -> tuple[sessionmaker, sessionmaker]:
        """
        Cria e retorna um sessionmaker síncrono e um sessionmaker assíncrono.
        """
        engine = create_engine(connection_string, echo=False)
        session = sessionmaker(bind=engine)

        async_engine = create_async_engine(async_connection_string)
        async_session = sessionmaker(bind=async_engine, class_=AsyncSession)

        return session, async_session

    def _initialize(self):
        """
        Garante que a tabela exista, com colunas para armazenar user_input e response.
        """
        with self._session() as session:
            session.execute(text(f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    chat_store_key VARCHAR(255) NOT NULL,
                    user_input TEXT,
                    response TEXT,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """))
            session.commit()

    def get_keys(self) -> list[str]:
        """
        Retorna todas as chaves armazenadas.
        """
        with self._session() as session:
            result = session.execute(text(f"""
                SELECT DISTINCT chat_store_key FROM {self.table_name}
            """))
            return [row[0] for row in result.fetchall()]

    def get_messages(self, key: str) -> list[ChatMessage]:
        """
        Retorna a conversa inteira (perguntas e respostas), na ordem de inserção (id).
        Cada linha pode conter o user_input, o response ou ambos (caso já respondido).
        """
        with self._session() as session:
            rows = session.execute(text(f"""
                SELECT user_input, response
                FROM {self.table_name}
                WHERE chat_store_key = :key
                ORDER BY id
            """), {"key": key}).fetchall()

            messages = []
            for user_in, resp in rows:
                if user_in is not None:
                    messages.append(ChatMessage(role='user', content=user_in))
                if resp is not None:
                    messages.append(ChatMessage(role='assistant', content=resp))
            return messages

    def set_messages(self, key: str, messages: list[ChatMessage]) -> None:
        """
        Sobrescreve o histórico de mensagens de uma chave (apaga tudo e insere novamente).
        Se quiser somente acrescentar, use add_message.
        
        Aqui, cada pergunta do usuário gera uma nova linha.
        Assim que encontrar uma mensagem de assistente, atualiza essa mesma linha.
        Se houver assistentes sem usuários, insere normalmente.
        """
        with self._session() as session:
            # Limpa histórico anterior
            session.execute(text(f"""
                DELETE FROM {self.table_name} WHERE chat_store_key = :key
            """), {"key": key})

            # Reinsere na ordem
            current_id = None
            for msg in messages:
                if msg.role == 'user':
                    # Cria nova linha com user_input
                    result = session.execute(text(f"""
                        INSERT INTO {self.table_name} (chat_store_key, user_input)
                        VALUES (:key, :ui)
                    """), {"key": key, "ui": msg.content})
                    # Pega o id do insert
                    current_id = result.lastrowid

                else:
                    # Tenta atualizar a última linha se existir
                    if current_id is not None:
                        session.execute(text(f"""
                            UPDATE {self.table_name}
                            SET response = :resp
                            WHERE id = :id
                        """), {"resp": msg.content, "id": current_id})
                        # Depois de atualizar a linha, zera o current_id
                        current_id = None
                    else:
                        # Se não houver pergunta pendente, insere como nova linha
                        session.execute(text(f"""
                            INSERT INTO {self.table_name} (chat_store_key, response)
                            VALUES (:key, :resp)
                        """), {"key": key, "resp": msg.content})

            session.commit()

    def add_message(self, key: str, message: ChatMessage) -> None:
        """
        Acrescenta uma nova mensagem no fluxo. Se for do usuário, insere nova linha;
        se for do assistente, tenta preencher a linha pendente que não tenha resposta.
        """

        with self._session() as session:
            if message.role == 'user':
                # Sempre cria uma nova linha para mensagens de usuário
                insert_stmt = text(f"""
                    INSERT INTO {self.table_name} (chat_store_key, user_input)
                    VALUES (:key, :ui)
                """)
                session.execute(insert_stmt, {
                    "key": key,
                    "ui": message.content
                })
            else:
                # Tenta encontrar a última linha sem resposta
                
                row = session.execute(text(f"""
                    SELECT id
                    FROM {self.table_name}
                    WHERE chat_store_key = :key
                      AND user_input IS NOT NULL
                      AND response IS NULL
                    ORDER BY id DESC
                    LIMIT 1
               """), {"key": key}).fetchone()

                if row:
                    # Atualiza com a resposta
                    msg_id = row[0]
                    
                    update_stmt = text(f"""
                        UPDATE {self.table_name}
                        SET response = :resp
                        WHERE id = :id
                    """)
                    session.execute(update_stmt, {
                        "resp": message.content,
                        "id": msg_id
                    })
                else:
                    # Se não achar linha pendente, insere como nova
                   
                    insert_stmt = text(f"""
                        INSERT INTO {self.table_name} (chat_store_key, response)
                        VALUES (:key, :resp)
                    """)
                    session.execute(insert_stmt, {
                        "key": key,
                        "resp": message.content
                    })

            session.commit()
            


    def delete_messages(self, key: str) -> None:
        """
        Remove todas as linhas associadas a 'key'.
        """
        with self._session() as session:
            session.execute(text(f"""
                DELETE FROM {self.table_name} WHERE chat_store_key = :key
            """), {"key": key})
            session.commit()

    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
        """
        Apaga a última mensagem da conversa (considerando a ordem de inserção).
        Se a última linha tiver pergunta e resposta, remove primeiro a resposta;
        caso não exista resposta, remove a linha inteira.
        """
        with self._session() as session:
            # Localiza a última linha
            row = session.execute(text(f"""
                SELECT id, user_input, response
                FROM {self.table_name}
                WHERE chat_store_key = :key
                ORDER BY id DESC
                LIMIT 1
            """), {"key": key}).fetchone()

            if not row:
                return None

            row_id, user_in, resp = row

            # Se a linha tiver somente pergunta, apagamos a linha inteira.
            # Se tiver também a resposta, apagamos só a parte do assistente.
            if user_in and resp:
                # Remove a resposta
                session.execute(text(f"""
                    UPDATE {self.table_name}
                    SET response = NULL
                    WHERE id = :id
                """), {"id": row_id})
                session.commit()
                return ChatMessage(role='assistant', content=resp)
            else:
                # Deleta a linha inteira
                session.execute(text(f"""
                    DELETE FROM {self.table_name}
                    WHERE id = :id
                """), {"id": row_id})
                session.commit()

                if user_in:
                    return ChatMessage(role='user', content=user_in)
                elif resp:
                    return ChatMessage(role='assistant', content=resp)
                else:
                    return None

    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
        """
        Deleta a mensagem com base na ordem total do histórico. O índice 'idx' é
        calculado após reconstruir a lista de ChatMessages (user e assistant).
        """
        messages = self.get_messages(key)
        if idx < 0 or idx >= len(messages):
            return None

        removed = messages[idx]

        # Agora precisamos traduzir 'idx' para saber qual registro no banco será modificado.
        # É mais simples recriar todos os dados com set_messages sem a mensagem em 'idx':
        messages.pop(idx)
        self.set_messages(key, messages)

        return removed