Skip to content

Commit 3efb73a

Browse files
manjunathshivaeavanvalkenburg
manjunathshiva
and
eavanvalkenburg
authored
Python: add a MongoDB agent RAG sample (#8725)
MongoDB has had a vector store since #8184 and core has had the generic VectorCollectionContextProvider since #8421, so grounding an agent on MongoDB already works. Nothing in the package showed it. The only MongoDB sample stores three hand-written vectors and never mentions an agent, both context-provider samples use InMemoryCollection, and the package README did not mention agents at all, so the combination is easy to miss and easy to conclude is unsupported. Two things bite anyone who tries it from the existing material: The documented context-provider sample declares its model with is_full_text_indexed=True. MongoDB rejects that at collection construction with NotImplementedError, because it has no Search text index, so copying that sample and swapping in MongoDBCollection fails on the first call. The new sample uses is_indexed for filterable metadata and says why. MongoDB indexes new documents asynchronously, so an agent that searches straight after upsert can find an empty index and answer that it has no information. The sample waits for the whole batch to become retrievable before running the agent, which is what an ingestion path has to do anyway. Measured against a live Atlas cluster this took five probes, roughly ten seconds, so the wait is load-bearing rather than defensive. Verified end to end against MongoDB Atlas 8.0.32 with Azure OpenAI standing in for the OpenAI endpoint: index creation, embedding generation, index propagation, retrieval and grounded answers. The script dependency floor is core 1.19.0, the first release containing VectorCollectionContextProvider; 1.18.0 was cut a week before it merged and would fail on import. Also records in the package README that a model must not declare is_full_text_indexed, and that newly written records are not immediately searchable. Co-authored-by: Eduard van Valkenburg <[email protected]>
1 parent 3f1f3b5 commit 3efb73a

3 files changed

Lines changed: 185 additions & 3 deletions

File tree

‎python/packages/mongodb/README.md‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,33 @@ async def main() -> None:
6666
asyncio.run(main())
6767
```
6868

69+
## Retrieval-augmented generation
70+
71+
`VectorCollectionContextProvider` exposes a vector collection to an agent as tools,
72+
so grounding an agent on MongoDB needs nothing beyond Agent Framework. Turning off
73+
the write tools keeps retrieval read-only.
74+
75+
```python
76+
from agent_framework import Agent, VectorCollectionContextProvider
77+
78+
knowledge_base = VectorCollectionContextProvider(
79+
collection,
80+
scope_filter=None,
81+
include_upsert_tool=False,
82+
include_get_tool=False,
83+
include_delete_tool=False,
84+
)
85+
86+
async with Agent(client=chat_client, context_providers=[knowledge_base]) as agent:
87+
print((await agent.run("How long does a refund take?")).text)
88+
```
89+
90+
As with any MongoDB record type, the model must not declare `is_full_text_indexed`; mark
91+
metadata you want to filter on with `is_indexed` instead. Newly written records become
92+
searchable asynchronously, so an ingestion path should confirm they are retrievable before
93+
querying them. See [`samples/mongodb_agent_rag.py`](samples/mongodb_agent_rag.py) for a
94+
runnable example that does both.
95+
6996
## Capabilities and limits
7097

7198
- String, signed 64-bit integer, and BSON `ObjectId` `_id` values retain native
@@ -91,7 +118,9 @@ asyncio.run(main())
91118
membership, AND, and OR only. Nested paths, NOT, null/missing, list membership,
92119
and literal/analyzed text are rejected for vector search.
93120
- Keyword-hybrid search, sparse/binary vectors, provider-side embedding
94-
generation, and automatic schema migration are not supported.
121+
generation, and automatic schema migration are not supported. A model that
122+
declares `is_full_text_indexed` is rejected when the collection is constructed,
123+
because there is no Search text index to back it.
95124
- ANN defaults `numCandidates` to the MongoDB recommendation of 20 times
96125
`skip + top`, capped at MongoDB's maximum of 10,000. Explicit values must be
97126
1-10,000 and at least `skip + top`; larger result windows require exact search.

‎python/packages/mongodb/samples/README.md‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,16 @@
44
prefiltering, multiple vector fields, and vector exclusion during ordinary
55
retrieval. It creates and deletes a unique collection.
66

7+
`mongodb_agent_rag.py` grounds an agent on a MongoDB collection through
8+
`VectorCollectionContextProvider`, which exposes the collection to the model as a
9+
search tool. It seeds a small knowledge base, waits for the new documents to become
10+
searchable, answers two questions, and deletes the collection it created.
11+
712
Run MongoDB Atlas or
813
[Atlas Local](https://www.mongodb.com/docs/atlas/cli/current/atlas-cli-deploy-local/),
914
set `MONGODB_URI` and `MONGODB_DATABASE_NAME`, and follow the
10-
[package connection guidance](../README.md#connection). No embedding service is
11-
required.
15+
[package connection guidance](../README.md#connection).
16+
17+
`mongodb_vectors.py` supplies its own vectors and needs no embedding service.
18+
`mongodb_agent_rag.py` additionally needs `OPENAI_API_KEY` for the embeddings and
19+
the agent.

‎python/packages/mongodb/samples/mongodb_agent_rag.py‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
# /// script
2+
# dependencies = [
3+
# "agent-framework-core>=1.19.0,<2",
4+
# "agent-framework-mongodb>=1.0.0a260909,<2",
5+
# "agent-framework-openai>=1.14.4,<2",
6+
# ]
7+
# ///
8+
9+
# Copyright (c) Microsoft. All rights reserved.
10+
11+
from __future__ import annotations
12+
13+
import asyncio
14+
from dataclasses import dataclass
15+
from typing import Annotated
16+
from uuid import uuid4
17+
18+
from agent_framework import (
19+
Agent,
20+
VectorCollectionContextProvider,
21+
VectorStoreField,
22+
vectorstoremodel,
23+
)
24+
from agent_framework.openai import OpenAIChatClient, OpenAIEmbeddingClient
25+
26+
from agent_framework_mongodb import MongoDBCollection, MongoDBStore
27+
28+
"""Answer questions from a MongoDB Atlas collection (retrieval-augmented generation).
29+
30+
``VectorCollectionContextProvider`` turns any vector collection into agent tools, so a
31+
MongoDB collection is grounded through the same path as any other store, with no
32+
retrieval library beyond Agent Framework itself.
33+
34+
Set ``MONGODB_URI`` and ``MONGODB_DATABASE_NAME`` for the deployment, and ``OPENAI_API_KEY``
35+
for the embeddings and the agent. Vector search requires Atlas or Atlas Local; a standalone
36+
``mongod`` cannot serve the query.
37+
38+
The collection name is unique per run and is deleted during cleanup, so existing
39+
collections are untouched.
40+
"""
41+
42+
43+
@vectorstoremodel
44+
@dataclass
45+
class Article:
46+
"""A support article the agent can retrieve but not modify."""
47+
48+
id: Annotated[str, VectorStoreField("key")]
49+
# MongoDB rejects `is_full_text_indexed`; it has no Search text index or
50+
# keyword-hybrid support. Use `is_indexed` for metadata you want to filter on.
51+
text: Annotated[str, VectorStoreField("data")]
52+
category: Annotated[str, VectorStoreField("data", is_indexed=True)]
53+
text_vector: Annotated[
54+
list[float] | str | None,
55+
VectorStoreField("vector", dimensions=1536, distance_function="cosine_similarity"),
56+
] = None
57+
58+
def __post_init__(self) -> None:
59+
"""Embed the article body when no vector is supplied."""
60+
if self.text_vector is None:
61+
self.text_vector = self.text
62+
63+
64+
async def wait_until_searchable(collection: MongoDBCollection[str, Article], probe: str, expected: int) -> None:
65+
"""Block until every newly written article is returned by vector search.
66+
67+
MongoDB indexes new documents asynchronously, so a search issued straight after
68+
`upsert` can legitimately return nothing even though the index is queryable. A
69+
production ingestion path needs the same allowance.
70+
71+
Vector search returns nearest neighbours rather than matches above a threshold, so
72+
the probe text only has to be non-empty; counting the results is what establishes
73+
that the whole batch is visible.
74+
"""
75+
for _ in range(30):
76+
results = await collection.search(values=probe, top=expected)
77+
if len([record async for record in results]) >= expected:
78+
return
79+
await asyncio.sleep(2)
80+
raise TimeoutError(f"Only part of the batch became searchable; expected {expected} articles.")
81+
82+
83+
async def main() -> None:
84+
articles = [
85+
Article("refunds", "Refunds are issued to the original payment method within five business days.", "billing"),
86+
Article("seats", "Team plans may add seats at any time; the next invoice is prorated.", "billing"),
87+
Article("sso", "Single sign-on is configured from the admin console under Security.", "account"),
88+
]
89+
90+
# 1. MongoDBStore resolves the URI and database through MongoDBSettings.
91+
async with MongoDBStore() as store:
92+
collection = store.get_collection(
93+
Article,
94+
collection_name=f"af_rag_{uuid4().hex}",
95+
embedding_generator=OpenAIEmbeddingClient(model="text-embedding-3-small"),
96+
)
97+
try:
98+
# 2. Create the collection and its Atlas vectorSearch index, then seed the
99+
# knowledge base. Passing strings in the vector field lets upsert
100+
# generate the embeddings.
101+
await collection.ensure_collection_exists()
102+
await collection.upsert(articles)
103+
await wait_until_searchable(collection, probe="support article", expected=len(articles))
104+
105+
# 3. Expose retrieval only. The agent grounds its answers on the collection
106+
# and cannot write to it, which is the usual shape for RAG.
107+
knowledge_base = VectorCollectionContextProvider(
108+
collection,
109+
scope_filter=None,
110+
include_upsert_tool=False,
111+
include_get_tool=False,
112+
include_delete_tool=False,
113+
)
114+
115+
async with Agent(
116+
client=OpenAIChatClient(model="gpt-5.4-nano"),
117+
name="SupportAssistant",
118+
instructions=(
119+
"Answer using the support articles returned by the search tool. "
120+
"If the articles do not cover the question, say so instead of guessing."
121+
),
122+
context_providers=[knowledge_base],
123+
) as agent:
124+
session = agent.create_session()
125+
for question in (
126+
"How long does a refund take?",
127+
"Can we add more seats mid-cycle?",
128+
):
129+
response = await agent.run(question, session=session)
130+
print(f"User: {question}")
131+
print(f"Assistant: {response.text}\n")
132+
finally:
133+
await collection.ensure_collection_deleted()
134+
135+
136+
if __name__ == "__main__":
137+
asyncio.run(main())
138+
139+
140+
# Expected output resembles:
141+
# User: How long does a refund take?
142+
# Assistant: Refunds are issued to the original payment method within five business days.
143+
#
144+
# User: Can we add more seats mid-cycle?
145+
# Assistant: Yes. Team plans may add seats at any time; the next invoice is prorated.

0 commit comments

Comments
 (0)