Autogen memory integration

1
import os
2
import uuid
3
import asyncio
4
from autogen_agentchat.agents import AssistantAgent
5
from autogen_ext.models.openai import OpenAIChatCompletionClient
6
from autogen_core.memory import MemoryContent, MemoryMimeType
7
from zep_cloud.client import AsyncZep
8
from zep_autogen import ZepUserMemory

1
# Initialize Zep client
2
zep_client = AsyncZep(api_key=os.environ.get("ZEP_API_KEY"))
3
user_id = f"user_{uuid.uuid4().hex[:16]}"
4
thread_id = f"thread_{uuid.uuid4().hex[:16]}"
5
6
# Create user (required before using memory)
7
try:
8
    await zep_client.user.add(
9
        user_id=user_id,
10
        email="[email protected]",
11
        first_name="Alice"
12
    )
13
except Exception as e:
14
    print(f"User might already exist: {e}")
15
16
# Create thread (required for conversation memory)
17
try:
18
    await zep_client.thread.create(thread_id=thread_id, user_id=user_id)
19
except Exception as e:
20
    print(f"Thread creation failed: {e}")

1
# Create user memory with configuration
2
memory = ZepUserMemory(
3
    client=zep_client,
4
    user_id=user_id,
5
    thread_id=thread_id
6
)

1
# Create agent with Zep memory
2
agent = AssistantAgent(
3
    name="MemoryAwareAssistant",
4
    model_client=OpenAIChatCompletionClient(
5
        model="gpt-4.1-mini",
6
        api_key=os.environ.get("OPENAI_API_KEY")
7
    ),
8
    memory=[memory],
9
    system_message="You are a helpful assistant with persistent memory."
10
)

1
# Helper function to store messages with proper metadata
2
async def add_message(message: str, role: str, name: str = None):
3
    """Store a message in Zep memory following AutoGen standards."""
4
    metadata = {"type": "message", "role": role}
5
    if name:
6
        metadata["name"] = name
7
    
8
    await memory.add(MemoryContent(
9
        content=message,
10
        mime_type=MemoryMimeType.TEXT,
11
        metadata=metadata
12
    ))
13
14
# Example conversation with memory persistence
15
user_message = "My name is Alice and I love hiking in the mountains."
16
print(f"User: {user_message}")
17
18
# Store user message
19
await add_message(user_message, "user", "Alice")
20
21
# Run agent - it will automatically retrieve context via update_context()
22
response = await agent.run(task=user_message)
23
agent_response = response.messages[-1].content
24
print(f"Agent: {agent_response}")
25
26
# Store agent response
27
await add_message(agent_response, "assistant")

1
from zep_autogen.graph_memory import ZepGraphMemory
2
from zep_cloud.external_clients.ontology import EntityModel, EntityText
3
from pydantic import Field
4
5
# Define entity models using Pydantic
6
class ProgrammingLanguage(EntityModel):
7
    """A programming language entity."""
8
    paradigm: EntityText = Field(
9
        description="programming paradigm (e.g., object-oriented, functional)",
10
        default=None
11
    )
12
    use_case: EntityText = Field(
13
        description="primary use cases for this language",
14
        default=None
15
    )
16
17
class Framework(EntityModel):
18
    """A software framework or library."""
19
    language: EntityText = Field(
20
        description="the programming language this framework is built for",
21
        default=None
22
    )
23
    purpose: EntityText = Field(
24
        description="primary purpose of this framework",
25
        default=None
26
    )

1
from zep_cloud import SearchFilters
2
3
# Set ontology first
4
await zep_client.graph.set_ontology(
5
    entities={
6
        "ProgrammingLanguage": ProgrammingLanguage,
7
        "Framework": Framework,
8
    }
9
)
10
11
# Create graph
12
graph_id = f"graph_{uuid.uuid4().hex[:16]}"
13
try:
14
    await zep_client.graph.create(
15
        graph_id=graph_id,
16
        name="Programming Knowledge Graph"
17
    )
18
    print(f"Created graph: {graph_id}")
19
except Exception as e:
20
    print(f"Graph creation failed: {e}")

1
# Create graph memory with search configuration
2
graph_memory = ZepGraphMemory(
3
    client=zep_client,
4
    graph_id=graph_id,
5
    search_filters=SearchFilters(
6
        node_labels=["ProgrammingLanguage", "Framework"]
7
    ),
8
    facts_limit=20,  # Max facts in context injection (default: 20)
9
    entity_limit=5   # Max entities in context injection (default: 5)
10
)

1
# Add structured knowledge
2
await graph_memory.add(MemoryContent(
3
    content="Python is excellent for data science and AI development",
4
    mime_type=MemoryMimeType.TEXT,
5
    metadata={"type": "data"}  # "data" stores in graph, "message" stores as episode
6
))
7
8
# Wait for graph processing (required)
9
print("Waiting for graph indexing...")
10
await asyncio.sleep(30)  # Allow time for knowledge extraction
11
12
<Callout intent="info">
13
**Graph Memory Context Injection**: ZepGraphMemory automatically retrieves the last 2 episodes from the graph and uses their content to query for relevant facts (up to `facts_limit`) and entities (up to `entity_limit`). This context is injected as a system message during agent interactions.
14
</Callout>

1
# Create agent with graph memory
2
agent = AssistantAgent(
3
    name="GraphMemoryAssistant",
4
    model_client=OpenAIChatCompletionClient(model="gpt-4.1-mini"),
5
    memory=[graph_memory],
6
    system_message="You are a technical assistant with programming knowledge."
7
)