Store conversation info in redis memory.
| 17 | |
| 18 | |
| 19 | class RedisMemory(BaseMemory): |
| 20 | """Store conversation info in redis memory.""" |
| 21 | |
| 22 | expire_time: int = ONE_HOUR |
| 23 | redis_key_prefix: str |
| 24 | redis_client: Redis |
| 25 | |
| 26 | class Config: |
| 27 | """Configuration for this pydantic object.""" |
| 28 | |
| 29 | arbitrary_types_allowed = True |
| 30 | |
| 31 | def load_memory( |
| 32 | self, key: Optional[str] = None, default: Optional[Any] = None, **kwargs |
| 33 | ) -> Any: |
| 34 | """Get the key's corresponding value from redis.""" |
| 35 | if not key.startswith(self.redis_key_prefix): |
| 36 | key = self.redis_key_prefix + f":{key}" |
| 37 | pickled = self.redis_client.get(key) |
| 38 | if not pickled: |
| 39 | return default |
| 40 | return pickle.loads(pickled) |
| 41 | |
| 42 | def load_conversation(self, **kwargs: Dict[str, Any]) -> ChatMessageHistory: |
| 43 | """Return chat message history.""" |
| 44 | redis_key = self.redis_key_prefix + f":{ChatMessageHistory.__name__}" |
| 45 | return ChatMessageHistory(messages=self.load_memory(redis_key, [])) |
| 46 | |
| 47 | def save_memory(self, key: str, value: Any) -> None: |
| 48 | """Save the key value pair to redis.""" |
| 49 | if not key.startswith(self.redis_key_prefix): |
| 50 | key = self.redis_key_prefix + f":{key}" |
| 51 | pickled = pickle.dumps(value) |
| 52 | self.redis_client.set(key, pickled, ex=self.expire_time) |
| 53 | |
| 54 | def save_conversation( |
| 55 | self, message: str, message_type: MessageType, **kwargs |
| 56 | ) -> None: |
| 57 | """Save context from this conversation to redis.""" |
| 58 | redis_key = self.redis_key_prefix + f":{ChatMessageHistory.__name__}" |
| 59 | pickled = self.redis_client.get(redis_key) |
| 60 | if pickled: |
| 61 | messages: list[BaseMessage] = pickle.loads(pickled) |
| 62 | else: |
| 63 | messages = [] |
| 64 | if message_type == MessageType.AIMessage: |
| 65 | messages.append(AIMessage(content=message)) |
| 66 | elif message_type == MessageType.UserMessage: |
| 67 | messages.append(UserMessage(content=message)) |
| 68 | elif message_type == MessageType.FunctionMessage: |
| 69 | messages.append(FunctionMessage(content=message, name=kwargs["name"])) |
| 70 | elif message_type == MessageType.SystemMessage: |
| 71 | messages.append(SystemMessage(content=message)) |
| 72 | else: |
| 73 | raise ValueError(f"Unsupported message type: {message_type}") |
| 74 | self.save_memory(redis_key, messages) |
| 75 | |
| 76 | def clear(self) -> None: |
no outgoing calls