Autogen Memory Example

1
import os
2
from dotenv import load_dotenv
3
import uuid
4
from typing import Union, Dict
5
from autogen import ConversableAgent, Agent
6
7
load_dotenv()
8
9
config_list = [
10
    {
11
        "model": "gpt-4o-mini",
12
        "api_key": os.environ.get("OPENAI_API_KEY"),
13
        "max_tokens": 1024,
14
    }
15
]

1
from zep_cloud.client import AsyncZep
2
from zep_cloud import Message, FactRatingExamples, FactRatingInstruction
3
4
MIN_FACT_RATING = 0.3
5
6
# Configure Zep
7
zep = AsyncZep(api_key=os.environ.get("ZEP_API_KEY"))

1
class ZepConversableAgent(ConversableAgent):
2
    """
3
    A custom ConversableAgent that integrates with Zep for long-term memory.
4
    """
5
    def __init__(
6
        self,
7
        name: str,
8
        system_message: str,
9
        llm_config: dict,
10
        function_map: dict,
11
        human_input_mode: str,
12
        zep_session_id: str,
13
    ):
14
        super().__init__(
15
            name=name,
16
            system_message=system_message,
17
            llm_config=llm_config,
18
            function_map=function_map,
19
            human_input_mode=human_input_mode,
20
        )
21
        self.zep_session_id = zep_session_id
22
        # store the original system message as we will update it with relevant facts from Zep
23
        self.original_system_message = system_message
24
        self.register_hook(
25
            "a_process_last_received_message", self.persist_user_messages
26
        )
27
        self.register_hook(
28
            "a_process_message_before_send", self.persist_assistant_messages
29
        )
30
31
    async def persist_assistant_messages(
32
        self, sender: Agent, message: Union[Dict, str], recipient: Agent, silent: bool
33
    ):
34
        """Agent sends a message to the user. Add the message to Zep."""
35
36
        # Assume message is a string
37
        zep_messages = convert_to_zep_messages(
38
            [{"role": "assistant", "name": self.name, "content": message}]
39
        )
40
        await zep.memory.add(session_id=self.zep_session_id, messages=zep_messages)
41
42
        return message
43
44
    async def persist_user_messages(self, messages: list[dict[str, str]] | str):
45
        """
46
        User sends a message to the agent. Add the message to Zep and
47
        update the system message with relevant facts from Zep.
48
        """
49
        # Assume messages is a string
50
        zep_messages = convert_to_zep_messages([{"role": "user", "content": messages}])
51
        await zep.memory.add(session_id=self.zep_session_id, messages=zep_messages)
52
53
        memory = await zep.memory.get(self.zep_session_id, min_rating=MIN_FACT_RATING)
54
55
        # Update the system message with the relevant facts retrieved from Zep
56
        self.update_system_message(
57
            self.original_system_message
58
            + f"\n\nRelevant facts about the user and their prior conversation:\n{memory.relevant_facts}"
59
        )
60
61
        return messages

1
bot_name = "CareBot"
2
user_name = "Cathy"
3
4
user_id = user_name + str(uuid.uuid4())[:4]
5
session_id = str(uuid.uuid4())
6
7
fact_rating_instruction = """Rate the facts by poignancy. Highly poignant
8
facts have a significant emotional impact or relevance to the user.
9
Low poignant facts are minimally relevant or of little emotional
10
significance."""
11
fact_rating_examples = FactRatingExamples(
12
    high="The user received news of a family member's serious illness.",
13
    medium="The user completed a challenging marathon.",
14
    low="The user bought a new brand of toothpaste.",
15
)
16
17
await zep.user.add(user_id=user_id,
18
fact_rating_instruction=FactRatingInstruction(
19
        instruction=fact_rating_instruction,
20
        examples=fact_rating_examples,
21
    )
22
)

1
chat_history = [
2
    {
3
        "role": "assistant",
4
        "name": "carebot",
5
        "content": "Hi Cathy, how are you doing today?",
6
    },
7
    {
8
        "role": "user",
9
        "name": "Cathy",
10
        "content": "To be honest, I've been feeling a bit down and demotivated lately. It's been tough.",
11
    },
12
    {
13
        "role": "assistant",
14
        "name": "CareBot",
15
        "content": "I'm sorry to hear that you're feeling down and demotivated, Cathy. It's understandable given the challenges you're facing. Can you tell me more about what's been going on?",
16
    },
17
    {
18
        "role": "user",
19
        "name": "Cathy",
20
        "content": "Well, I'm really struggling to process the passing of my mother.",
21
    },
22
    {
23
        "role": "assistant",
24
        "name": "CareBot",
25
        "content": "I'm deeply sorry for your loss, Cathy. Losing a parent is incredibly difficult. It's normal to struggle with grief, and there's no 'right' way to process it. Would you like to talk about your mother or how you're coping?",
26
    },
27
    {
28
        "role": "user",
29
        "name": "Cathy",
30
        "content": "Yes, I'd like to talk about my mother. She was a kind and loving person.",
31
    },
32
]
33
34
# Convert chat history to Zep messages
35
zep_messages = convert_to_zep_messages(chat_history)
36
37
await zep.memory.add(session_id=session_id, messages=zep_messages)

1
response = await zep.memory.get(session_id=session_id, min_rating=MIN_FACT_RATING)
2
3
for r in response.relevant_facts:
4
    print(r)

1
carebot_system_message = """
2
You are a compassionate mental health bot and caregiver. Review information about the user and their prior conversation below and respond accordingly.
3
Keep responses empathetic and supportive. And remember, always prioritize the user's well-being and mental health. Keep your responses very concise and to the point.
4
"""
5
6
agent = ZepConversableAgent(
7
    bot_name,
8
    system_message=carebot_system_message,
9
    llm_config={"config_list": config_list},
10
    function_map=None,  # No registered functions, by default it is None.
11
    human_input_mode="NEVER",  # Never ask for human input.
12
    zep_session_id=session_id,
13
)

1
cathy = ConversableAgent(
2
    user_name,
3
    system_message="You are a helpful mental health bot. You are seeking counsel from a mental health bot. Ask the bot about your previous conversation.",
4
    llm_config={"config_list": config_list},
5
    human_input_mode="NEVER",  # Never ask for human input.
6
)

1
result = await agent.a_initiate_chat(
2
    cathy,
3
    message="Hi Cathy, nice to see you again. How are you doing today?",
4
    max_turns=3,
5
)

1
response = await zep.memory.get(session_id, min_rating=MIN_FACT_RATING)
2
3
for r in response.relevant_facts:
4
    print(r)

1
response = await zep.graph.search(
2
    query="What do you know about Cathy's family?",
3
    user_id=user_id,
4
)
5
relevant_edges = response.edges
6
formatted_facts = []
7
for edge in relevant_edges:
8
    valid_at = edge.valid_at if edge.valid_at is not None else "date unknown"
9
    invalid_at = edge.invalid_at if edge.invalid_at is not None else "present"
10
    formatted_fact = f"{edge.fact} (Date range: {valid_at} - {invalid_at})"
11
    formatted_facts.append(formatted_fact)
12
13
# Print the results
14
print("\nFound facts:")
15
for fact in formatted_facts:
16
    print(f"- {fact}")