Memory#
There are several use cases where it is valuable to maintain a store of useful facts that can be intelligently added to the context of the agent just before a specific step. The typically use case here is a RAG pattern where a query is used to retrieve relevant information from a database that is then added to the agent’s context.
AgentChat provides a Memory protocol that can be extended to provide this functionality. The key methods are query, update_context, add, clear, and close.
add: add new entries to the memory storequery: retrieve relevant information from the memory storeupdate_context: mutate an agent’s internalmodel_contextby adding the retrieved information (used in theAssistantAgentclass)clear: clear all entries from the memory storeclose: clean up any resources used by the memory store
ListMemory Example#
{py:class}~autogen_core.memory.ListMemory is provided as an example implementation of the {py:class}~autogen_core.memory.Memory protocol. It is a simple list-based memory implementation that maintains memories in chronological order, appending the most recent memories to the model’s context. The implementation is designed to be straightforward and predictable, making it easy to understand and debug. In the following example, we will use ListMemory to maintain a memory bank of user preferences and demonstrate how it can be used to provide consistent context for agent responses over time.
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType
from autogen_ext.models.openai import OpenAIChatCompletionClient
# Initialize user memory
user_memory = ListMemory()
# Add user preferences to memory
await user_memory.add(MemoryContent(content="The weather should be in metric units", mime_type=MemoryMimeType.TEXT))
await user_memory.add(MemoryContent(content="Meal recipe must be vegan", mime_type=MemoryMimeType.TEXT))
async def get_weather(city: str, units: str = "imperial") -> str:
if units == "imperial":
return f"The weather in {city} is 73 °F and Sunny."
elif units == "metric":
return f"The weather in {city} is 23 °C and Sunny."
else:
return f"Sorry, I don't know the weather in {city}."
assistant_agent = AssistantAgent(
name="assistant_agent",
model_client=OpenAIChatCompletionClient(
model="gpt-4o-2024-08-06",
),
tools=[get_weather],
memory=[user_memory],
)
# Run the agent with a task.
stream = assistant_agent.run_stream(task="What is the weather in New York?")
await Console(stream)
---------- user ----------
What is the weather in New York?
---------- assistant_agent ----------
[MemoryContent(content='The weather should be in metric units', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None)]
---------- assistant_agent ----------
[FunctionCall(id='call_GpimUGED5zUbfxORaGo2JD6F', arguments='{"city":"New York","units":"metric"}', name='get_weather')]
---------- assistant_agent ----------
[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_GpimUGED5zUbfxORaGo2JD6F', is_error=False)]
---------- assistant_agent ----------
The weather in New York is 23 °C and Sunny.
TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='What is the weather in New York?', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, metadata={}, content=[MemoryContent(content='The weather should be in metric units', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None)], type='MemoryQueryEvent'), ToolCallRequestEvent(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=123, completion_tokens=20), metadata={}, content=[FunctionCall(id='call_GpimUGED5zUbfxORaGo2JD6F', arguments='{"city":"New York","units":"metric"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='assistant_agent', models_usage=None, metadata={}, content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_GpimUGED5zUbfxORaGo2JD6F', is_error=False)], type='ToolCallExecutionEvent'), ToolCallSummaryMessage(source='assistant_agent', models_usage=None, metadata={}, content='The weather in New York is 23 °C and Sunny.', type='ToolCallSummaryMessage')], stop_reason=None)
We can inspect that the assistant_agent model_context is actually updated with the retrieved memory entries. The transform method is used to format the retrieved memory entries into a string that can be used by the agent. In this case, we simply concatenate the content of each memory entry into a single string.
await assistant_agent._model_context.get_messages()
[UserMessage(content='What is the weather in New York?', source='user', type='UserMessage'),
SystemMessage(content='\nRelevant memory content (in chronological order):\n1. The weather should be in metric units\n2. Meal recipe must be vegan\n', type='SystemMessage'),
AssistantMessage(content=[FunctionCall(id='call_GpimUGED5zUbfxORaGo2JD6F', arguments='{"city":"New York","units":"metric"}', name='get_weather')], thought=None, source='assistant_agent', type='AssistantMessage'),
FunctionExecutionResultMessage(content=[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_GpimUGED5zUbfxORaGo2JD6F', is_error=False)], type='FunctionExecutionResultMessage')]
We see above that the weather is returned in Centigrade as stated in the user preferences.
Similarly, assuming we ask a separate question about generating a meal plan, the agent is able to retrieve relevant information from the memory store and provide a personalized (vegan) response.
stream = assistant_agent.run_stream(task="Write brief meal recipe with broth")
await Console(stream)
---------- user ----------
Write brief meal recipe with broth
---------- assistant_agent ----------
[MemoryContent(content='The weather should be in metric units', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None)]
---------- assistant_agent ----------
Here's a brief vegan broth recipe:
**Vegan Vegetable Broth**
**Ingredients:**
- 2 tablespoons olive oil
- 1 large onion, chopped
- 2 carrots, sliced
- 2 celery stalks, sliced
- 4 cloves garlic, minced
- 1 teaspoon salt
- 1/2 teaspoon pepper
- 1 bay leaf
- 1 teaspoon thyme
- 1 teaspoon rosemary
- 8 cups water
- 1 cup mushrooms, sliced
- 1 cup chopped leafy greens (e.g., kale, spinach)
- 1 tablespoon soy sauce (optional)
**Instructions:**
1. **Sauté Vegetables:** In a large pot, heat olive oil over medium heat. Add the onion, carrots, and celery, and sauté until they begin to soften, about 5-7 minutes.
2. **Add Garlic & Seasonings:** Stir in the garlic, salt, pepper, bay leaf, thyme, and rosemary. Cook for another 2 minutes until fragrant.
3. **Simmer Broth:** Pour in the water, add mushrooms and soy sauce (if using). Increase heat and bring to a boil. Reduce heat and let it simmer for 30-45 minutes.
4. **Add Greens:** In the last 5 minutes of cooking, add the chopped leafy greens.
5. **Strain & Serve:** Remove from heat, strain out the vegetables (or leave them in for a chunkier texture), and adjust seasoning if needed. Serve hot as a base for soups or enjoy as is!
Enjoy your flavorful, nourishing vegan broth! TERMINATE
TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Write brief meal recipe with broth', type='TextMessage'), MemoryQueryEvent(source='assistant_agent', models_usage=None, metadata={}, content=[MemoryContent(content='The weather should be in metric units', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None)], type='MemoryQueryEvent'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=208, completion_tokens=331), metadata={}, content="Here's a brief vegan broth recipe:\n\n**Vegan Vegetable Broth**\n\n**Ingredients:**\n- 2 tablespoons olive oil\n- 1 large onion, chopped\n- 2 carrots, sliced\n- 2 celery stalks, sliced\n- 4 cloves garlic, minced\n- 1 teaspoon salt\n- 1/2 teaspoon pepper\n- 1 bay leaf\n- 1 teaspoon thyme\n- 1 teaspoon rosemary\n- 8 cups water\n- 1 cup mushrooms, sliced\n- 1 cup chopped leafy greens (e.g., kale, spinach)\n- 1 tablespoon soy sauce (optional)\n\n**Instructions:**\n\n1. **Sauté Vegetables:** In a large pot, heat olive oil over medium heat. Add the onion, carrots, and celery, and sauté until they begin to soften, about 5-7 minutes.\n\n2. **Add Garlic & Seasonings:** Stir in the garlic, salt, pepper, bay leaf, thyme, and rosemary. Cook for another 2 minutes until fragrant.\n\n3. **Simmer Broth:** Pour in the water, add mushrooms and soy sauce (if using). Increase heat and bring to a boil. Reduce heat and let it simmer for 30-45 minutes.\n\n4. **Add Greens:** In the last 5 minutes of cooking, add the chopped leafy greens.\n\n5. **Strain & Serve:** Remove from heat, strain out the vegetables (or leave them in for a chunkier texture), and adjust seasoning if needed. Serve hot as a base for soups or enjoy as is!\n\nEnjoy your flavorful, nourishing vegan broth! TERMINATE", type='TextMessage')], stop_reason=None)
Custom Memory Stores (Vector DBs, etc.)#
You can build on the Memory protocol to implement more complex memory stores. For example, you could implement a custom memory store that uses a vector database to store and retrieve information, or a memory store that uses a machine learning model to generate personalized responses based on the user’s preferences etc.
Specifically, you will need to overload the add, query and update_context methods to implement the desired functionality and pass the memory store to your agent.
Currently the following example memory stores are available as part of the autogen_ext extensions package.
autogen_ext.memory.chromadb.ChromaDBVectorMemory: A memory store that uses a vector database to store and retrieve information.
import os
from pathlib import Path
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_core.memory import MemoryContent, MemoryMimeType
from autogen_ext.memory.chromadb import ChromaDBVectorMemory, PersistentChromaDBVectorMemoryConfig
from autogen_ext.models.openai import OpenAIChatCompletionClient
# Initialize ChromaDB memory with custom config
chroma_user_memory = ChromaDBVectorMemory(
config=PersistentChromaDBVectorMemoryConfig(
collection_name="preferences",
persistence_path=os.path.join(str(Path.home()), ".chromadb_autogen"),
k=2, # Return top k results
score_threshold=0.4, # Minimum similarity score
)
)
# a HttpChromaDBVectorMemoryConfig is also supported for connecting to a remote ChromaDB server
# Add user preferences to memory
await chroma_user_memory.add(
MemoryContent(
content="The weather should be in metric units",
mime_type=MemoryMimeType.TEXT,
metadata={"category": "preferences", "type": "units"},
)
)
await chroma_user_memory.add(
MemoryContent(
content="Meal recipe must be vegan",
mime_type=MemoryMimeType.TEXT,
metadata={"category": "preferences", "type": "dietary"},
)
)
# Create assistant agent with ChromaDB memory
assistant_agent = AssistantAgent(
name="assistant_agent",
model_client=OpenAIChatCompletionClient(
model="gpt-4o",
),
tools=[get_weather],
memory=[user_memory],
)
stream = assistant_agent.run_stream(task="What is the weather in New York?")
await Console(stream)
await user_memory.close()
---------- user ----------
What is the weather in New York?
---------- assistant_agent ----------
[MemoryContent(content='The weather should be in metric units', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None), MemoryContent(content='Meal recipe must be vegan', mime_type=<MemoryMimeType.TEXT: 'text/plain'>, metadata=None)]
---------- assistant_agent ----------
[FunctionCall(id='call_PKcfeEHXimGG2QOhJwXzCLuZ', arguments='{"city":"New York","units":"metric"}', name='get_weather')]
---------- assistant_agent ----------
[FunctionExecutionResult(content='The weather in New York is 23 °C and Sunny.', call_id='call_PKcfeEHXimGG2QOhJwXzCLuZ', is_error=False)]
---------- assistant_agent ----------
The weather in New York is 23 °C and Sunny.
Note that you can also serialize the ChromaDBVectorMemory and save it to disk.
chroma_user_memory.dump_component().model_dump_json()
'{"provider":"autogen_ext.memory.chromadb.ChromaDBVectorMemory","component_type":"memory","version":1,"component_version":1,"description":"ChromaDB-based vector memory implementation with similarity search.","label":"ChromaDBVectorMemory","config":{"client_type":"persistent","collection_name":"preferences","distance_metric":"cosine","k":2,"score_threshold":0.4,"allow_reset":false,"tenant":"default_tenant","database":"default_database","persistence_path":"/Users/victordibia/.chromadb_autogen"}}'