Spaces:
Sleeping
Sleeping
File size: 7,777 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | 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
|