Spaces:
Sleeping
Sleeping
| import redis | |
| import os | |
| import json | |
| from datetime import datetime | |
| from dotenv import load_dotenv | |
| 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 | |
| load_dotenv() | |
| class ChatStore: | |
| def __init__(self): | |
| self.redis_client = redis.Redis( | |
| # host="redis-10365.c244.us-east-1-2.ec2.redns.redis-cloud.com", | |
| host = os.getenv("REDIS_HOST"), | |
| port=os.getenv("REDIS_PORT"), | |
| username = os.getenv("REDIS_USERNAME"), | |
| password=os.getenv("REDIS_PASSWORD"), | |
| ) | |
| uri = os.getenv("MONGO_URI") | |
| self.client = MongoClient(uri) | |
| def initialize_memory_bot(self, session_id): | |
| # Decode Redis keys to work with strings | |
| redis_keys = [key.decode('utf-8') for key in self.redis_client.keys()] | |
| chat_store = RedisChatStore( | |
| redis_client=self.redis_client, ttl=86400 # Time-to-live set for 1 hour | |
| ) | |
| db = self.client["bot_database"] | |
| # Check if the session exists in Redis or MongoDB | |
| if session_id in redis_keys: | |
| # If the session already exists in Redis, create the memory buffer using Redis | |
| memory = ChatMemoryBuffer.from_defaults( | |
| token_limit=3000, chat_store=chat_store, chat_store_key=session_id | |
| ) | |
| elif session_id in db.list_collection_names(): | |
| # If the session exists in MongoDB but not Redis, fetch messages from MongoDB | |
| self.add_chat_history_to_redis(session_id) # Add chat history to Redis | |
| # Then create the memory buffer using Redis | |
| memory = ChatMemoryBuffer.from_defaults( | |
| token_limit=3000, chat_store=chat_store, chat_store_key=session_id | |
| ) | |
| else: | |
| # If the session doesn't exist in either Redis or MongoDB, create an empty memory buffer | |
| 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 get_last_message_mongodb(self, session_id: str): | |
| db = self.client["bot_database"] | |
| collection = db[session_id] | |
| # Get the last document by sorting by _id in descending order | |
| last_document = collection.find().sort("_id", -1).limit(1) | |
| # Iterasi last_document dan kembalikan isi content jika ada | |
| for doc in last_document: | |
| return str(doc.get('content', "")) # kembalikan content atau string kosong jika tidak ada | |
| # Jika tidak ada dokumen, kembalikan string kosong | |
| return "" | |
| 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: | |
| return [key.decode("utf-8") for key in self.redis_client.keys("*")] | |
| except Exception as e: | |
| return JSONResponse(status_code=400, content="the error when get keys") | |
| def add_message(self, session_id: str, message: Optional[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: Optional[ChatMessage]) -> dict: | |
| # Convert the ChatMessage instance into a dictionary with necessary adjustments | |
| message_dict = message.model_dump() | |
| # Convert any datetime fields to ISO format, if needed | |
| if isinstance(message_dict.get('timestamp'), datetime): | |
| message_dict['timestamp'] = message_dict['timestamp'].isoformat() | |
| return message_dict | |
| 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 | |
| ] | |
| return documents_list | |
| except Exception as e: | |
| return JSONResponse(status_code=500, content=f"An error occurred while retrieving messages: {e}") |