pydantic-ai-mongodb-memory

A MongoDB Atlas–backed memory capability for Pydantic AI and pydantic-ai-harness.

Drop MongoDBMemoryCapability into an agent's capabilities=[...] array and it gains durable, cross-session, semantically-searchable memory — backed by MongoDB documents + Atlas Vector Search. Anchored to pydantic-ai-harness#255.

Why

In a stateless agent, every new session forgets the user. Re-injecting the full transcript inflates tokens and hits context limits. This capability gives the agent a tiered memory:

• Short-term / episodic — recent turns, semantically searchable via $vectorSearch.
• Long-term structured profile — durable user preferences (always injected).
• Semantic recall — Voyage 3.5 embeddings + Atlas Vector Search, with prefiltering.

Install

pip install -e ".[embeddings]"   # pymongo + pydantic-ai-slim + voyageai

Quickstart

from pydantic_ai import Agent
from pydantic_ai_mongodb_memory import MongoDBMemoryCapability

memory = MongoDBMemoryCapability(
    connection_string="mongodb+srv://...",
    database_name="agent_memory",
    semantic_recall=True,            # uses Atlas Vector Search + voyage-3.5
)

agent = Agent("google:gemini-2.0-flash", capabilities=[memory])

The capability implements two AbstractCapability lifecycle hooks:

• before_model_request → recalls memories for the active scope and injects them.
• after_model_request → persists the latest turn(s) back to MongoDB.

How it maps to MongoDB

Method	MongoDB operation
add_memory(scope, role, content)	insert_one (+ voyage-3.5 embed if semantic)
get_recent(scope, n)	find({scope}).sort([(ts,-1),(_id,-1)])
recall_semantic(scope, query, k, filters=…)	$vectorSearch on embedding + prefilter
set_preference(scope, k, v)	update_one(upsert) on user_profiles
get_profile(scope)	find_one on user_profiles

Conventions: the store always constructs and owns its own MongoClient from your connection string, so connection appName = devrel-integ-pydanticai-python and handshake driver_info name pydantic-ai-mongodb-memory are always present and non-overridable; embeddings voyage-3.5 (1024-dim). Extra MongoClient options (e.g. tls=True, maxPoolSize=…) can be passed through as keyword args.

v0.1.1: the store now always owns its client — pass connection_string (the former client= parameter was removed) so appName + driver-info are guaranteed on every connection.

Demos

Both auto-load demo/.env (ATLAS_URI, VOYAGE_API_KEY, GEMINI_API_KEY).

python demo/main.py        # data-level: seed team dataset, multi-turn, scope isolation, $vectorSearch
python demo/agent_demo.py  # a Gemini agent: cross-session memory + Atlas Vector Search w/ department prefilter

agent_demo.py proves the headline value: an agent told "only staff Engineering" in session 1 answers a fresh session-2 staffing question using $vectorSearch prefiltered to Engineering — recalling the constraint from MongoDB with no history passed in.

Tests

PYTHONPATH=src pytest tests/ -v

Non-search tests run on mongomock (offline). test_vector_recall needs ATLAS_URI + VOYAGE_API_KEY and is skipped otherwise. Suite: round-trip, scope isolation, profile upsert, TTL index, vector recall, appName present, driver-info present.

Layout

src/pydantic_ai_mongodb_memory/
  capability.py   # MongoDBMemoryCapability (AbstractCapability subclass)
  store.py        # MongoDBMemoryStore — connection, indexes, CRUD, $vectorSearch
  embeddings.py   # voyage-3.5 helper
demo/             # main.py (data) + agent_demo.py (Gemini agent)
tests/            # acceptance suite
PLAN.md           # the 7-phase build plan