Multimedika/Bot_Development

1import redis2import os3import json4 5from datetime import datetime6from dotenv import load_dotenv7 8from fastapi.responses import JSONResponse9from typing import Optional, List, Dict10from llama_index.storage.chat_store.redis import RedisChatStore11from pymongo.mongo_client import MongoClient12from llama_index.core.memory import ChatMemoryBuffer13from service.dto import ChatMessage14 15load_dotenv()16 17 18class ChatStore:19    def __init__(self):20        self.redis_client = redis.Redis(21            # host="redis-10365.c244.us-east-1-2.ec2.redns.redis-cloud.com",22            host = os.getenv("REDIS_HOST"),23            port=os.getenv("REDIS_PORT"),24            username = os.getenv("REDIS_USERNAME"),25            password=os.getenv("REDIS_PASSWORD"),26        )27 28        uri = os.getenv("MONGO_URI")29        self.client = MongoClient(uri)30 31    def initialize_memory_bot(self, session_id):32        # Decode Redis keys to work with strings33        redis_keys = [key.decode('utf-8') for key in self.redis_client.keys()]34        35        chat_store = RedisChatStore(36            redis_client=self.redis_client, ttl=86400  # Time-to-live set for 1 hour37        )38        db = self.client["bot_database"]39        40        # Check if the session exists in Redis or MongoDB41        if session_id in redis_keys:42            # If the session already exists in Redis, create the memory buffer using Redis43            memory = ChatMemoryBuffer.from_defaults(44                token_limit=3000, chat_store=chat_store, chat_store_key=session_id45            )46        elif session_id in db.list_collection_names():47            # If the session exists in MongoDB but not Redis, fetch messages from MongoDB48            self.add_chat_history_to_redis(session_id)  # Add chat history to Redis49            # Then create the memory buffer using Redis50            memory = ChatMemoryBuffer.from_defaults(51                token_limit=3000, chat_store=chat_store, chat_store_key=session_id52            )53        else:54            # If the session doesn't exist in either Redis or MongoDB, create an empty memory buffer55            memory = ChatMemoryBuffer.from_defaults(56                token_limit=3000, chat_store=chat_store, chat_store_key=session_id57            )58        59        return memory60 61    def get_messages(self, session_id: str) -> List[dict]:62        """Get messages for a session_id."""63        items = self.redis_client.lrange(session_id, 0, -1)64        if len(items) == 0:65            return []66 67        # Decode and parse each item into a dictionary68        return [json.loads(m.decode("utf-8")) for m in items]69    70    def get_last_message(self, session_id: str) -> Optional[Dict]:71        """Get the last message for a session_id."""72        last_message = self.redis_client.lindex(session_id, -1)73        74        if last_message is None:75            return None  # Return None if there are no messages76        77        # Decode and parse the last message into a dictionary78        return json.loads(last_message.decode("utf-8"))79 80    def get_last_message_mongodb(self, session_id: str):81        db = self.client["bot_database"]82        collection = db[session_id]83 84        # Get the last document by sorting by _id in descending order85        last_document = collection.find().sort("_id", -1).limit(1)86        87        # Iterasi last_document dan kembalikan isi content jika ada88        for doc in last_document:89            return str(doc.get('content', ""))  # kembalikan content atau string kosong jika tidak ada90        91        # Jika tidak ada dokumen, kembalikan string kosong92        return ""93 94    def delete_last_message(self, session_id: str) -> Optional[ChatMessage]:95        """Delete last message for a session_id."""96        return self.redis_client.rpop(session_id)97 98    def delete_messages(self, session_id: str) -> Optional[List[ChatMessage]]:99        """Delete messages for a session_id."""100        self.redis_client.delete(session_id)101        db = self.client["bot_database"]102        db.session_id.drop()103        return None104 105    def clean_message(self, session_id: str) -> Optional[ChatMessage]:106        """Delete specific message for a session_id."""107        current_list = self.redis_client.lrange(session_id, 0, -1)108 109        indices_to_delete = []110        for index, item in enumerate(current_list):111            data = json.loads(item)  # Parse JSON string to dict112 113            # Logic to determine if item should be removed114            if (data.get("role") == "assistant" and data.get("content") is None) or (115                data.get("role") == "tool"116            ):117                indices_to_delete.append(index)118 119        # Remove elements by their indices in reverse order120        for index in reversed(indices_to_delete):121            self.redis_client.lrem(122                session_id, 1, current_list[index]123            )  # Remove the element from the list in Redis124 125    def get_keys(self) -> List[str]:126        """Get all keys."""127        try:128            return [key.decode("utf-8") for key in self.redis_client.keys("*")]129 130        except Exception as e:131            return JSONResponse(status_code=400, content="the error when get keys")132 133    def add_message(self, session_id: str, message: Optional[ChatMessage]) -> None:134        """Add a message for a session_id."""135        item = json.dumps(self._message_to_dict(message))136        self.redis_client.rpush(session_id, item)137 138    def _message_to_dict(self, message: Optional[ChatMessage]) -> dict:139        # Convert the ChatMessage instance into a dictionary with necessary adjustments140        message_dict = message.model_dump()141        # Convert any datetime fields to ISO format, if needed142        if isinstance(message_dict.get('timestamp'), datetime):143            message_dict['timestamp'] = message_dict['timestamp'].isoformat()144        return message_dict145 146    def add_chat_history_to_redis(self, session_id: str) -> None:147        """Fetch chat history from MongoDB and add it to Redis."""148        db = self.client["bot_database"]149        collection = db[session_id]150 151        try:152            chat_history = collection.find()153            chat_history_list = [154                {155                    key: message[key]156                    for key in message157                    if key not in ["_id", "timestamp"] and message[key] is not None158                }159                for message in chat_history160                if message is not None161            ]162 163            for message in chat_history_list:164                # Convert MongoDB document to the format you need165                item = json.dumps(166                    self._message_to_dict(ChatMessage(**message))167                )  # Convert message to dict168                # Push to Redis169                self.redis_client.rpush(session_id, item)170            self.redis_client.expire(session_id, time=86400)171 172        except Exception as e:173            return JSONResponse(status_code=500, content="Add Database Error")174 175    def get_all_messages_mongodb(self, session_id):176        """Get all messages for a session_id from MongoDB."""177        try:178            db = self.client["bot_database"]179            collection = db[session_id]180 181            # Retrieve all documents from the collection182            documents = collection.find()183 184            # Convert the cursor to a list and exclude the _id field185            documents_list = [186                {key: doc[key] for key in doc if key !="_id" and doc[key] is not None}187                for doc in documents188            ]189 190            return documents_list191        192        except Exception as e:193            return JSONResponse(status_code=500, content=f"An error occurred while retrieving messages: {e}")