feat: mongodb vector io #3772

seyeong-han · 2025-10-10T17:49:32Z

What does this PR do?

This pull request adds support for MongoDB Atlas as a remote vector database provider to Llama Stack, enabling cloud-native vector search functionality. The changes introduce configuration, documentation, and integration for MongoDB Atlas Vector Search across the codebase, including distribution templates, provider registry, and sample configurations.

MongoDB Atlas Vector Search Integration

Added remote::mongodb as a supported vector database provider in distribution build/run YAML files for starter, starter-gpu, and CI-tests, allowing MongoDB Atlas to be used for vector storage and search. [1] [2] [3] [4] [5] [6]
Integrated MongoDB Atlas configuration into the starter distribution Python template, including environment variables, provider registration, and sample config generation using MongoDBVectorIOConfig. [1] [2] [3] [4]

Provider Implementation

Added the MongoDB Atlas provider implementation in llama_stack/providers/remote/vector_io/mongodb, including adapter loading logic and a pydantic-based configuration class for connection and search options. [1] [2]

Documentation

Created comprehensive documentation for the MongoDB Atlas provider, detailing features, search modes (vector, keyword, hybrid), configuration, installation, and usage instructions.

Provider Registry

Registered the MongoDB Atlas provider in llama_stack/providers/registry/vector_io.py, including metadata, pip dependencies, and a detailed description for discoverability.

Miscellaneous

Added copyright header to the test module for providers.

Test Plan

Example RAG with response API

Get MongoDB connection key and Fireworks API key
Run llama-stack server with Fireworks API

MONGODB_CONNECTION_STRING="mongodb+srv://{GET_THIS_FROM_MONGODB}" \
FIREWORKS_API_KEY={YOUR_API_KEY} uv run --with llama-stack llama stack build --distro starter --image-type venv --run

Run RAG

from io import BytesIO
from llama_stack_client import LlamaStackClient

# Initialize client (use port 8321 which is the default for starter distribution)
client = LlamaStackClient(base_url="http://localhost:8321")

model_id = "accounts/fireworks/models/llama4-maverick-instruct-basic"

knowledge_base = [
    (
        "Python is a high-level programming language known for its readable syntax and versatility.",
        {"category": "Programming"},
    ),
]
file_ids = []
for content, metadata in knowledge_base:
    with BytesIO(content.encode()) as file_buffer:
        file_buffer.name = f"{metadata['category'].lower()}_{len(file_ids)}.txt"
        response = client.files.create(file=file_buffer, purpose="assistants")
        file_ids.append(response.id)

print(f"✅ Uploaded {len(file_ids)} documents")

print("Creating vector store with sentence-transformers/all-MiniLM-L6-v2...")

vector_store = client.vector_stores.create(
    name="mongodb_tech_knowledge_base",
    file_ids=file_ids,
    embedding_model="sentence-transformers/all-MiniLM-L6-v2",
    embedding_dimension=384,
    provider_id="mongodb_atlas",  # MongoDB Atlas provider
)

print(f"✅ Created MongoDB Atlas vector store: {vector_store.name}")
print(f"Vector store id: {vector_store.id}")

query = "What is Python?"
try:
    response = client.responses.create(
        model=model_id,
        input=query,
        tools=[
            {
                "type": "file_search",
                "vector_store_ids": [vector_store.id],
            },
        ],
    )
    # print(f"Response: {response}\n")
    print(f"File search results: {response.output[-1].content[0].text}\n")
except Exception as e:
    print(f"❌ Query failed: {str(e)[:150]}...")

--> Result

source .venv/bin/activate
python examples/mongodb_response_api_simple.py
INFO:httpx:HTTP Request: POST http://localhost:8321/v1/openai/v1/files "HTTP/1.1 200 OK"
✅ Uploaded 1 documents
Creating vector store with sentence-transformers/all-MiniLM-L6-v2...
INFO:httpx:HTTP Request: POST http://localhost:8321/v1/openai/v1/vector_stores "HTTP/1.1 200 OK"
✅ Created MongoDB Atlas vector store: mongodb_tech_knowledge_base
Vector store id: vs_f09d3c4b-0163-43b2-8f49-d6701643cf7a
INFO:httpx:HTTP Request: POST http://localhost:8321/v1/openai/v1/responses "HTTP/1.1 200 OK"
File search results: Python is a high-level, interpreted programming language known for its simplicity and readability. It is widely used for various purposes such as web development, data analysis, artificial intelligence, and more. Python's syntax is designed to be easy to understand, making it a great language for beginners and experienced developers alike.

MongoDB Atlas llama_stack.vs_f09d3c4b_0163_43b2_8f49_d6701643cf7a collection

…provider

franciscojavierarceo · 2025-10-15T03:39:52Z

docs/docs/providers/vector_io/remote_mongodb.mdx

+)
+```
+
+## Usage


i think part of this readme is duplicated.

I believe the # Usage and # Configuration sections are duplicated.

franciscojavierarceo · 2025-10-15T03:41:26Z

llama_stack/distributions/ci-tests/build.yaml

    - provider_type: inline::milvus
    - provider_type: remote::chromadb
    - provider_type: remote::pgvector
+    - provider_type: remote::mongodb


you need to also add an image that'll run MongoDB and expose the vector search functionality.

you'll need to add this to the ./github/workflows/integratino-vector-io.tests.yaml file like here: https://github.com/llamastack/llama-stack/blob/main/.github/workflows/integration-vector-io-tests.yml#L43

franciscojavierarceo · 2025-10-15T03:44:29Z

llama_stack/distributions/starter/starter.py

+                            provider_id="${env.MONGODB_CONNECTION_STRING:+mongodb_atlas}",
+                            provider_type="remote::mongodb",
+                            config=MongoDBVectorIOConfig.sample_run_config(
+                                f"~/.llama/distributions/{name}",


you probably want to add the other configuration parameters.

franciscojavierarceo · 2025-10-15T03:45:25Z

llama_stack/providers/registry/vector_io.py

+)
+```
+
+## Usage


BTW I believe the docs are now compiled from this description, so this may be why you have duplication.

franciscojavierarceo · 2025-10-15T03:45:43Z

llama_stack/providers/remote/vector_io/mongodb/__init__.py

+
+    # Handle the deps resolution - if files API exists, pass it, otherwise None
+    files_api = deps.get(Api.files)
+    impl = MongoDBVectorIOAdapter(config, deps[Api.inference], files_api)


need to add models_api

franciscojavierarceo · 2025-10-15T03:46:46Z

llama_stack/providers/remote/vector_io/mongodb/mongodb.py

+                logger.info(f"Creating collection '{self.collection.name}'")
+                # Create collection by inserting a dummy document
+                dummy_doc = {"_id": "__dummy__", "dummy": True}
+                self.collection.insert_one(dummy_doc)


why insert a dummy record and delete it? was this for testing?

franciscojavierarceo · 2025-10-15T03:47:30Z

llama_stack/providers/remote/vector_io/mongodb/mongodb.py

+                f"Collection name: {self.collection.name}. Error: {str(e)}"
+            )
+            # Don't fail completely - just log the error and continue
+            logger.warning(


we should raise if we can't start the index because then vector stores doesn't work.

franciscojavierarceo · 2025-10-15T03:55:30Z

llama_stack/providers/remote/vector_io/mongodb/mongodb.py

+                        "index": self.config.index_name,
+                        "queryVector": embedding.tolist(),
+                        "path": self.config.path_field,
+                        "numCandidates": k * 10,  # Get more candidates for better results


is this intentional? multiplicative increase seems aggressive, i could understand a constant offset but this feels like it could end up with bad consequences.

franciscojavierarceo · 2025-10-15T03:59:42Z

llama_stack/providers/remote/vector_io/mongodb/mongodb.py

+        self,
+        config: MongoDBVectorIOConfig,
+        inference_api,
+        files_api=None,


needs models_api

franciscojavierarceo

tried to add some comments throughout, you will need some adjustments now that #3803 has landed.

You also need to update the vector-io action to enable the integration tests. along with that you need to update the skip_if_provider_doesnt_support_openai_vector_stores test in ./tests/integration/vector_io/test_openai_vector_stores.py.

This is definitely great to see so thanks for contributing this! BTW I"M not sure if you've seen MongoDB does has an external provided implementation here: https://github.com/mongodb-partners/mongodb-llama-stack.

They actually created a PR quite some time ago here: https://github.com/llamastack/llama-stack/pull/1603/files but the stack has changed some since then. Think we'll be able to get this in though!

seyeong-han added 5 commits October 10, 2025 09:43

init: add mongodb in vector_io

dd1136a

fix: add missing __init__.py files for vector_io test directories

aa21996

fix: add missing __init__.py files for vector_io test directories

2b6d57a

Merge remote-tracking branch 'upstream/main' into add-mongodb-vector_io

d8c82c1

feat: add mongodb provider

885631d

seyeong-han requested review from ashwinb, bbrowning, ehhuang, franciscojavierarceo, hardikjshah, leseb, mattf, raghotham, reluctantfuturist, slekkala1, terrytangyuan and yanxi0830 as code owners October 10, 2025 17:49

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 10, 2025

seyeong-han changed the title ~~Add: mongodb vector io~~ feat: mongodb vector io Oct 10, 2025

fix: resolve unit test failures in starter distribution with MongoDB …

572062d

…provider

leseb added the new-in-tree-provider label Oct 13, 2025

franciscojavierarceo reviewed Oct 15, 2025

View reviewed changes

franciscojavierarceo requested changes Oct 15, 2025

View reviewed changes

feat: mongodb vector io #3772

Are you sure you want to change the base?

feat: mongodb vector io #3772

Uh oh!

Conversation

seyeong-han commented Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Test Plan

Example RAG with response API

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

franciscojavierarceo left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

seyeong-han commented Oct 10, 2025 •

edited

Loading

franciscojavierarceo left a comment •

edited

Loading