diff --git a/ai-agent-memory/README.md b/ai-agent-memory/README.md new file mode 100644 index 0000000000..ea50bbfdff --- /dev/null +++ b/ai-agent-memory/README.md @@ -0,0 +1,60 @@ +# AI Agent Memory in Python: A Primer + +This folder contains the sample code for the Real Python tutorial [AI Agent Memory in Python: A Primer](https://realpython.com/ai-agent-memory/). + +## Setup + +Both scripts call the OpenAI API, so they need an `OPENAI_API_KEY` and make paid API calls: + +```console +$ export OPENAI_API_KEY="your-api-key" +``` + +Each script declares its dependency, `openai==3.35.1`, in an inline script metadata block, so [uv](https://docs.astral.sh/uv/) installs it for you. If you don't use uv, install the pinned package with `python -m pip install openai==3.35.1` and run the scripts with `python`. + +## Short-Term Memory + +`short_term.py` asks the same question with and without the conversation history: + +```console +$ uv run short_term.py +``` + +## Agent-Managed Memory + +`managed_memory.py` gives the model four memory tools and stores what it saves in `memories.json`. Each memory has text and topic tags, which the search tool uses to find matching entries. Each command is a fresh session, so the agent remembers only what's in the file. Delete `memories.json` to start over. + +Here's the four-message run from the tutorial. The model's wording, its memories, and even its choice of tools vary from run to run: + +```console +$ uv run managed_memory.py "Hi! I'm Ana, I'm learning Python, and I prefer short answers with code examples." + tool: save_memory({'text': 'User name is Ana. She is learning Python and prefers short answers with code examples.', 'tags': 'name, learning, preferences'}) +Hi Ana! I’ll keep answers short and include code examples. +$ uv run managed_memory.py "I moved from Lisbon to Berlin last month." + tool: save_memory({'text': 'User moved from Lisbon to Berlin last month.', 'tags': 'location'}) +Got it. +$ uv run managed_memory.py "Forget that I prefer short answers." + tool: search_memory({'query': 'preferences'}) + tool: update_memory({'id': '1', 'text': 'User name is Ana. She is learning Python and no longer prefers short answers.', 'tags': 'name, learning, preferences'}) +Got it — I’ve removed the short-answer preference. +$ uv run managed_memory.py "Where do I live, and what do you know about me?" + tool: search_memory({'query': 'name learning location'}) +I know that you’re Ana, you’re learning Python, and you moved from Lisbon to Berlin last month. + +So, you likely live in Berlin now. +``` + +After this run, `memories.json` holds two entries: + +```json +{ + "1": { + "text": "User name is Ana. She is learning Python and no longer prefers short answers.", + "tags": ["learning", "name", "preferences"] + }, + "2": { + "text": "User moved from Lisbon to Berlin last month.", + "tags": ["location"] + } +} +``` diff --git a/ai-agent-memory/managed_memory.py b/ai-agent-memory/managed_memory.py new file mode 100644 index 0000000000..3b23a7e973 --- /dev/null +++ b/ai-agent-memory/managed_memory.py @@ -0,0 +1,155 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["openai==3.35.1"] +# /// +import json +import re +import sys +from pathlib import Path + +from openai import OpenAI + +client = OpenAI() +MODEL = "gpt-5.4-mini" +STORE = Path("memories.json") +memories = json.loads(STORE.read_text()) if STORE.exists() else {} + + +def save_memory(text, tags): + new_id = str(max(map(int, memories), default=0) + 1) + memories[new_id] = {"text": text, "tags": _parse_tags(tags)} + return "saved" + + +def _parse_tags(tags): + return sorted( + {tag.strip().lower() for tag in tags.split(",") if tag.strip()} + ) + + +def search_memory(query): + query_words = _words(query) + matches = { + id: memory + for id, memory in memories.items() + if query_words + & (_words(memory["text"]) | _words(" ".join(memory["tags"]))) + } + return json.dumps(matches) + + +STOP_WORDS = { + "a", + "an", + "and", + "are", + "do", + "i", + "is", + "me", + "my", + "of", + "the", + "to", + "user", + "what", + "where", +} + + +def _words(text): + return set(re.findall(r"[a-z0-9]+", text.lower())) - STOP_WORDS + + +def update_memory(id, text, tags): + if id not in memories: + return "not found" + memories[id] = {"text": text, "tags": _parse_tags(tags)} + return "updated" + + +def delete_memory(id): + if id not in memories: + return "not found" + del memories[id] + return "deleted" + + +def spec(fn, description, *params): + props = {name: {"type": "string"} for name in params} + schema = { + "type": "object", + "properties": props, + "required": list(params), + "additionalProperties": False, + } + return { + "type": "function", + "function": { + "name": fn.__name__, + "description": description, + "parameters": schema, + "strict": True, + }, + } + + +TOOLS = { + f.__name__: f + for f in (save_memory, search_memory, update_memory, delete_memory) +} +SPECS = [ + spec( + save_memory, + "Store a fact with comma-separated topic tags.", + "text", + "tags", + ), + spec( + search_memory, + "Search memories by topic or keywords, such as location or preferences.", + "query", + ), + spec( + update_memory, + "Replace a memory and its comma-separated tags.", + "id", + "text", + "tags", + ), + spec(delete_memory, "Delete a wrong or unneeded memory.", "id"), +] + +if len(sys.argv) != 2: + sys.exit('Usage: uv run managed_memory.py "Your message"') + +messages = [ + { + "role": "system", + "content": "You have tools to manage your own long-term memory. Decide " + "when to save, search, update, or delete memories. Give each memory " + "topic tags, such as location or preferences. Search by topic first " + "to get ids before you update or delete. Answer briefly.", + }, + {"role": "user", "content": sys.argv[1]}, +] +for _ in range(10): # The agent loop, capped at ten rounds + reply = ( + client.chat.completions.create( + model=MODEL, messages=messages, tools=SPECS + ) + .choices[0] + .message + ) + messages.append(reply) + if not reply.tool_calls: + break + for call in reply.tool_calls: + args = json.loads(call.function.arguments) + print(f" tool: {call.function.name}({args})") + result = TOOLS[call.function.name](**args) + STORE.write_text(json.dumps(memories, indent=2)) # Save every change + messages.append( + {"role": "tool", "tool_call_id": call.id, "content": result} + ) +print(reply.content) diff --git a/ai-agent-memory/short_term.py b/ai-agent-memory/short_term.py new file mode 100644 index 0000000000..fefa4e8ac4 --- /dev/null +++ b/ai-agent-memory/short_term.py @@ -0,0 +1,31 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = ["openai==3.35.1"] +# /// +from openai import OpenAI + +client = OpenAI() +MODEL = "gpt-5.4-mini" + + +SYSTEM = {"role": "system", "content": "Answer in one short sentence."} + + +def chat(messages): + response = client.chat.completions.create( + model=MODEL, messages=[SYSTEM, *messages] + ) + return response.choices[0].message.content + + +history = [] +for turn in ["Hi, I'm Ana and I'm learning Python.", "What's my name?"]: + history.append({"role": "user", "content": turn}) + reply = chat(history) # The whole history goes in every call + history.append({"role": "assistant", "content": reply}) + print(f"With history: {reply}") + +# Same question, but without the earlier turns +print( + "Without history:", chat([{"role": "user", "content": "What's my name?"}]) +)