Spaces:
Sleeping
Sleeping
| from typing import List, Any, Optional, cast | |
| from pydantic import BaseModel, ConfigDict | |
| from phi.model.base import Model | |
| from phi.model.message import Message | |
| from phi.memory.memory import Memory | |
| from phi.memory.db import MemoryDb | |
| from phi.memory.row import MemoryRow | |
| from phi.utils.log import logger | |
| class MemoryManager(BaseModel): | |
| model: Optional[Model] = None | |
| user_id: Optional[str] = None | |
| # Provide the system prompt for the manager as a string | |
| system_prompt: Optional[str] = None | |
| # Memory Database | |
| db: Optional[MemoryDb] = None | |
| # Do not set the input message here, it will be set by the run method | |
| input_message: Optional[str] = None | |
| model_config = ConfigDict(arbitrary_types_allowed=True) | |
| def update_model(self) -> None: | |
| if self.model is None: | |
| try: | |
| from phi.model.openai import OpenAIChat | |
| except ModuleNotFoundError as e: | |
| logger.exception(e) | |
| logger.error( | |
| "phidata uses `openai` as the default model provider. Please provide a `model` or install `openai`." | |
| ) | |
| exit(1) | |
| self.model = OpenAIChat() | |
| self.model.add_tool(self.add_memory) | |
| self.model.add_tool(self.update_memory) | |
| self.model.add_tool(self.delete_memory) | |
| self.model.add_tool(self.clear_memory) | |
| def get_existing_memories(self) -> Optional[List[MemoryRow]]: | |
| if self.db is None: | |
| return None | |
| return self.db.read_memories(user_id=self.user_id) | |
| def add_memory(self, memory: str) -> str: | |
| """Use this function to add a memory to the database. | |
| Args: | |
| memory (str): The memory to be stored. | |
| Returns: | |
| str: A message indicating if the memory was added successfully or not. | |
| """ | |
| try: | |
| if self.db: | |
| self.db.upsert_memory( | |
| MemoryRow(user_id=self.user_id, memory=Memory(memory=memory, input=self.input_message).to_dict()) | |
| ) | |
| return "Memory added successfully" | |
| except Exception as e: | |
| logger.warning(f"Error storing memory in db: {e}") | |
| return f"Error adding memory: {e}" | |
| def delete_memory(self, id: str) -> str: | |
| """Use this function to delete a memory from the database. | |
| Args: | |
| id (str): The id of the memory to be deleted. | |
| Returns: | |
| str: A message indicating if the memory was deleted successfully or not. | |
| """ | |
| try: | |
| if self.db: | |
| self.db.delete_memory(id=id) | |
| return "Memory deleted successfully" | |
| except Exception as e: | |
| logger.warning(f"Error deleting memory in db: {e}") | |
| return f"Error deleting memory: {e}" | |
| def update_memory(self, id: str, memory: str) -> str: | |
| """Use this function to update a memory in the database. | |
| Args: | |
| id (str): The id of the memory to be updated. | |
| memory (str): The updated memory. | |
| Returns: | |
| str: A message indicating if the memory was updated successfully or not. | |
| """ | |
| try: | |
| if self.db: | |
| self.db.upsert_memory( | |
| MemoryRow( | |
| id=id, user_id=self.user_id, memory=Memory(memory=memory, input=self.input_message).to_dict() | |
| ) | |
| ) | |
| return "Memory updated successfully" | |
| except Exception as e: | |
| logger.warning(f"Error updating memory in db: {e}") | |
| return f"Error updating memory: {e}" | |
| def clear_memory(self) -> str: | |
| """Use this function to clear all memories from the database. | |
| Returns: | |
| str: A message indicating if the memory was cleared successfully or not. | |
| """ | |
| try: | |
| if self.db: | |
| self.db.clear() | |
| return "Memory cleared successfully" | |
| except Exception as e: | |
| logger.warning(f"Error clearing memory in db: {e}") | |
| return f"Error clearing memory: {e}" | |
| def get_system_message(self) -> Message: | |
| # -*- Return a system message for the memory manager | |
| system_prompt_lines = [ | |
| "Your task is to generate a concise memory for the user's message. " | |
| "Create a memory that captures the key information provided by the user, as if you were storing it for future reference. " | |
| "The memory should be a brief, third-person statement that encapsulates the most important aspect of the user's input, without adding any extraneous details. " | |
| "This memory will be used to enhance the user's experience in subsequent conversations.", | |
| "You will also be provided with a list of existing memories. You may:", | |
| " 1. Add a new memory using the `add_memory` tool.", | |
| " 2. Update a memory using the `update_memory` tool.", | |
| " 3. Delete a memory using the `delete_memory` tool.", | |
| " 4. Clear all memories using the `clear_memory` tool. Use this with extreme caution, as it will remove all memories from the database.", | |
| ] | |
| existing_memories = self.get_existing_memories() | |
| if existing_memories and len(existing_memories) > 0: | |
| system_prompt_lines.extend( | |
| [ | |
| "\nExisting memories:", | |
| "<existing_memories>\n" | |
| + "\n".join([f" - id: {m.id} | memory: {m.memory}" for m in existing_memories]) | |
| + "\n</existing_memories>", | |
| ] | |
| ) | |
| return Message(role="system", content="\n".join(system_prompt_lines)) | |
| def run( | |
| self, | |
| message: Optional[str] = None, | |
| **kwargs: Any, | |
| ) -> Optional[str]: | |
| logger.debug("*********** MemoryManager Start ***********") | |
| # Update the Model (set defaults, add logit etc.) | |
| self.update_model() | |
| # Prepare the List of messages to send to the Model | |
| messages_for_model: List[Message] = [self.get_system_message()] | |
| # Add the user prompt message | |
| user_prompt_message = Message(role="user", content=message, **kwargs) if message else None | |
| if user_prompt_message is not None: | |
| messages_for_model += [user_prompt_message] | |
| # Set input message added with the memory | |
| self.input_message = message | |
| # Generate a response from the Model (includes running function calls) | |
| self.model = cast(Model, self.model) | |
| response = self.model.response(messages=messages_for_model) | |
| logger.debug("*********** MemoryManager End ***********") | |
| return response.content | |
| async def arun( | |
| self, | |
| message: Optional[str] = None, | |
| **kwargs: Any, | |
| ) -> Optional[str]: | |
| logger.debug("*********** Async MemoryManager Start ***********") | |
| # Update the Model (set defaults, add logit etc.) | |
| self.update_model() | |
| # Prepare the List of messages to send to the Model | |
| messages_for_model: List[Message] = [self.get_system_message()] | |
| # Add the user prompt message | |
| user_prompt_message = Message(role="user", content=message, **kwargs) if message else None | |
| if user_prompt_message is not None: | |
| messages_for_model += [user_prompt_message] | |
| # Set input message added with the memory | |
| self.input_message = message | |
| # Generate a response from the Model (includes running function calls) | |
| self.model = cast(Model, self.model) | |
| response = await self.model.aresponse(messages=messages_for_model) | |
| logger.debug("*********** Async MemoryManager End ***********") | |
| return response.content | |