Spaces:
Sleeping
Sleeping
| import redis | |
| import os | |
| import json | |
| from fastapi import HTTPException | |
| from uuid import uuid4 | |
| from typing import Optional, List | |
| from llama_index.storage.chat_store.redis import RedisChatStore | |
| 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"), | |
| ) | |
| def generate_uuid(use_hex=False): | |
| if use_hex: | |
| return str(uuid4().hex) | |
| else: | |
| return str(uuid4()) | |
| def initialize_memory_bot(self, session_id=None): | |
| if session_id is None: | |
| session_id = self.generate_uuid() | |
| # chat_store = SimpleChatStore() | |
| chat_store = RedisChatStore( | |
| redis_client=self.redis_client | |
| ) # Need to be configured | |
| 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 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, key: str) -> Optional[List[ChatMessage]]: | |
| """Delete messages for a key.""" | |
| self.redis_client.delete(key) | |
| 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 raise HTTPException for FastAPI | |
| print(f"An error occurred in update data.: {e}") | |
| raise HTTPException( | |
| status_code=400, detail="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() |