Sample code for: AI Agent Memory in Python: A Primer

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"]
+  }
+}
+```

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)

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?"}])
+)