1from kern.db.postgres import PostgresDb
2from kern.memory import MemoryManager, UserMemory
3from rich.pretty import pprint
4
5db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
6
7memory = MemoryManager(db=PostgresDb(db_url=db_url))
8
9# Add a memory for the default user
10memory.add_user_memory(
11 memory=UserMemory(memory="The user's name is John Doe", topics=["name"]),
12)
13print("Memories:")
14pprint(memory.get_user_memories())
15
16# Add memories for Jane Doe
17jane_doe_id = "jane_doe@example.com"
18print(f"\nUser: {jane_doe_id}")
19memory_id_1 = memory.add_user_memory(
20 memory=UserMemory(memory="The user's name is Jane Doe", topics=["name"]),
21 user_id=jane_doe_id,
22)
23memory_id_2 = memory.add_user_memory(
24 memory=UserMemory(memory="She likes to play tennis", topics=["hobbies"]),
25 user_id=jane_doe_id,
26)
27memories = memory.get_user_memories(user_id=jane_doe_id)
28print("Memories:")
29pprint(memories)
30
31# Delete a memory
32print("\nDeleting memory")
33assert memory_id_2 is not None
34memory.delete_user_memory(user_id=jane_doe_id, memory_id=memory_id_2)
35print("Memory deleted\n")
36memories = memory.get_user_memories(user_id=jane_doe_id)
37print("Memories:")
38pprint(memories)
39
40# Replace a memory
41print("\nReplacing memory")
42assert memory_id_1 is not None
43memory.replace_user_memory(
44 memory_id=memory_id_1,
45 memory=UserMemory(memory="The user's name is Jane Mary Doe", topics=["name"]),
46 user_id=jane_doe_id,
47)
48print("Memory replaced")
49memories = memory.get_user_memories(user_id=jane_doe_id)
50print("Memories:")
51pprint(memories)
