core/chat/chatstore.py

import redis
import os
import json
from fastapi.responses import JSONResponse
from typing import Optional, List, Dict
from llama_index.storage.chat_store.redis import RedisChatStore
from pymongo.mongo_client import MongoClient
from llama_index.core.memory import ChatMemoryBuffer
from service.dto import ChatMessage


class ChatStore:
    def __init__(self):
        self.redis_client = redis.Redis(
            host="redis-10365.c244.us-east-1-2.ec2.redns.redis-cloud.com",
            port=10365,
            password=os.environ.get("REDIS_PASSWORD"),
        )

        uri = os.getenv("MONGO_URI")
        self.client = MongoClient(uri)

    def initialize_memory_bot(self, session_id):
        
        chat_store = RedisChatStore(
            redis_client=self.redis_client, ttl=86400  # Time-to-live set for 1 hour
        )

        db = self.client["bot_database"]

        if (
            self.redis_client.exists(session_id)
            or session_id in db.list_collection_names()
        ):
            if session_id not in self.redis_client.keys():
                self.add_chat_history_to_redis(
                    session_id
                )  # Add chat history to Redis if not found
            # Create memory buffer with chat store and session key
            memory = ChatMemoryBuffer.from_defaults(
                token_limit=3000, chat_store=chat_store, chat_store_key=session_id
            )
        else:
            # Handle the case where the session doesn't exist
            memory = ChatMemoryBuffer.from_defaults(
                token_limit=3000, chat_store=chat_store, chat_store_key=session_id
            )

        return memory

    def get_messages(self, session_id: str) -> List[dict]:
        """Get messages for a session_id."""
        items = self.redis_client.lrange(session_id, 0, -1)
        if len(items) == 0:
            return []

        # Decode and parse each item into a dictionary
        return [json.loads(m.decode("utf-8")) for m in items]
    
    def get_last_message(self, session_id: str) -> Optional[Dict]:
        """Get the last message for a session_id."""
        last_message = self.redis_client.lindex(session_id, -1)
        
        if last_message is None:
            return None  # Return None if there are no messages
        
        # Decode and parse the last message into a dictionary
        return json.loads(last_message.decode("utf-8"))


    def delete_last_message(self, session_id: str) -> Optional[ChatMessage]:
        """Delete last message for a session_id."""
        return self.redis_client.rpop(session_id)

    def delete_messages(self, session_id: str) -> Optional[List[ChatMessage]]:
        """Delete messages for a session_id."""
        self.redis_client.delete(session_id)
        db = self.client["bot_database"]
        db.session_id.drop()
        return None

    def clean_message(self, session_id: str) -> Optional[ChatMessage]:
        """Delete specific message for a session_id."""
        current_list = self.redis_client.lrange(session_id, 0, -1)

        indices_to_delete = []
        for index, item in enumerate(current_list):
            data = json.loads(item)  # Parse JSON string to dict

            # Logic to determine if item should be removed
            if (data.get("role") == "assistant" and data.get("content") is None) or (
                data.get("role") == "tool"
            ):
                indices_to_delete.append(index)

        # Remove elements by their indices in reverse order
        for index in reversed(indices_to_delete):
            self.redis_client.lrem(
                session_id, 1, current_list[index]
            )  # Remove the element from the list in Redis

    def get_keys(self) -> List[str]:
        """Get all keys."""
        try:
            print(self.redis_client.keys("*"))
            return [key.decode("utf-8") for key in self.redis_client.keys("*")]

        except Exception as e:
            # Log the error and return JSONResponse for FastAPI
            print(f"An error occurred in update data.: {e}")
            return JSONResponse(status_code=400, content="the error when get keys")

    def add_message(self, session_id: str, message: ChatMessage) -> None:
        """Add a message for a session_id."""
        item = json.dumps(self._message_to_dict(message))
        self.redis_client.rpush(session_id, item)

    def _message_to_dict(self, message: ChatMessage) -> dict:
        return message.model_dump()

    def add_chat_history_to_redis(self, session_id: str) -> None:
        """Fetch chat history from MongoDB and add it to Redis."""
        db = self.client["bot_database"]
        collection = db[session_id]

        try:
            chat_history = collection.find()
            chat_history_list = [
                {
                    key: message[key]
                    for key in message
                    if key not in ["_id", "timestamp"] and message[key] is not None
                }
                for message in chat_history
                if message is not None
            ]

            for message in chat_history_list:
                # Convert MongoDB document to the format you need
                item = json.dumps(
                    self._message_to_dict(ChatMessage(**message))
                )  # Convert message to dict
                # Push to Redis
                self.redis_client.rpush(session_id, item)
            self.redis_client.expire(session_id, time=86400)

        except Exception as e:
            return JSONResponse(status_code=500, content="Add Database Error")

    def get_all_messages_mongodb(self, session_id):
        """Get all messages for a session_id from MongoDB."""
        try:
            db = self.client["bot_database"]
            collection = db[session_id]

            # Retrieve all documents from the collection
            documents = collection.find()

            # Convert the cursor to a list and exclude the _id field
            documents_list = [
                {key: doc[key] for key in doc if key !="_id" and doc[key] is not None}
                for doc in documents
            ]

            # Print the list of documents without the _id field
            print(documents_list)  # Optional: If you want to see the output

            return documents_list
        
        except Exception as e:
            print(f"An error occurred while retrieving messages: {e}")
            return JSONResponse(status_code=500, content=f"An error occurred while retrieving messages: {e}")