Build Persistent Agent Working Memory with MongoDB and Tavily

This guide is adapted from the MongoDB GenAI Showcase notebook implementing_working_memory_with_tavily_and_mongodb.ipynb, available at https://github.com/mongodb-developer/GenAI-Showcase under the MIT licence.

Most AI agent tutorials demonstrate memory as an in-process Python dict that evaporates the moment the process exits. That works for demos, but it breaks the instant your agent needs to span multiple sessions, recover from a crash mid-task, or share state with a second agent running in parallel. Persistent working memory — backed by a real database — solves all three problems at once. The pattern shown here pairs MongoDB as the persistent store with Tavily's hybrid search client, so the agent can combine a private product catalog (local, vector-indexed) with live web results (remote) and write everything back into the same collection. The result is a memory layer that survives process restarts, scales horizontally, and doubles as a retrieval index.

This pattern suits engineers building multi-step research agents, customer-facing sales assistants, or any agentic system where a legal agent model or similar long-horizon task runner must remember what it found two steps ago without reprocessing from scratch. You need a free MongoDB Atlas cluster (M0 tier is sufficient for experiments), a Cohere account (embed-english-v3.0 has a free tier), and a Tavily API key (also free-tier-eligible). Hardware requirements are minimal — a standard Colab CPU runtime is enough.

Prerequisites

• Python 3.9 or later
• A MongoDB Atlas account with a cluster that has Atlas Vector Search enabled (M0 free tier works)
• A Cohere API key — used for embed-english-v3.0 (1024 dimensions)
• A Tavily API key — used for hybrid RAG search
• A Hugging Face account with HF_TOKEN set in your environment, to access the philschmid/amazon-product-descriptions-vlm dataset

Step 1: Install dependencies and configure credentials

Install the four libraries the notebook depends on. The -q flag keeps output tidy in Colab.

%pip install -U -q tavily-python cohere pymongo datasets pandas

Credentials are collected interactively rather than hard-coded, which keeps secrets out of notebooks that get committed to version control. The helper below wraps getpass so you can call it once per secret without repeating the pattern.

import getpass
import os


# Function to securely get and set environment variables
def set_env_securely(var_name, prompt):
    value = getpass.getpass(prompt)
    os.environ[var_name] = value

You will call set_env_securely three more times throughout the guide — once for Cohere, once for MongoDB, and once for Tavily. Keep this cell's output visible so you notice if a prompt is skipped.

Step 2: Build the product dataset and generate embeddings

The long-term memory backing the agent is a MongoDB collection of Amazon product records. The first task is to shape the raw dataset into something semantically searchable, which means concatenating the fields that carry product meaning into a single string before embedding.

import pandas as pd
from datasets import load_dataset

# Make sure you have an HF_TOKEN in your environment varibales to access dataset on hugging face
product_dataset = load_dataset("philschmid/amazon-product-descriptions-vlm")

# Convert product_dataset to pandas dataframe
product_dataframe = pd.DataFrame(product_dataset["train"])

# Display top 5 rows
product_dataframe.head()

The raw dataset includes an image column that adds weight without contributing to text retrieval. Before embedding, drop it and construct product_semantics — the field that the vector index will be built on.

# Create a new coloumn in the dataset that combines existing coloumns that captures a product semantics
product_dataframe["product_semantics"] = product_dataframe.apply(
    lambda row: " ".join(
        str(x)
        for x in [
            row["Product Name"],
            row["Category"],
            row["About Product"],
            row["Technical Details"],
            row["description"],
        ]
        if x
    ),
    axis=1,
)

# Display top 5 data point to see new coloumn
product_dataframe.head()

# We are not using the image data, so let's remove this
product_dataframe = product_dataframe.drop(columns=["image"])

Now attach your Cohere key and generate the embeddings. The get_embedding function wraps the Cohere V2 client and passes input_type="search_document" — Cohere's asymmetric embedding model expects documents and queries to be tagged differently, and using the wrong tag silently degrades retrieval quality.

# Add Cohere API Key to Environment Variable
set_env_securely("COHERE_API_KEY", "Enter your Cohere API Key: ")

import cohere

co = cohere.ClientV2()


def get_embedding(texts, model="embed-english-v3.0", input_type="search_document"):
    """Gets embeddings for a list of texts using the Cohere API.

    Args:
      texts: A list of texts to embed.
      model: The Cohere embedding model to use.
      input_type: The input type for the embedding model.

    Returns:
      A list of embeddings, where each embedding is a list of floats.
    """
    try:
        response = co.embed(
            texts=[texts],
            model=model,
            input_type=input_type,
            embedding_types=["float"],
        )
        # Extract and return the embeddings
        return response.embeddings.float[0]
    except Exception as e:
        print(f"Error generating embeddings: {e}")
        print("Couldn't generate emebedding for text: ")
        print(texts)
        return None

# Generate an embedding coloum for each datapoint in the dataset
# Embedding is generated from the new product semantics attribute
try:
    product_dataframe["embedding"] = product_dataframe["product_semantics"].apply(
        get_embedding
    )
    print("Embeddings generated successfully")
except Exception as e:
    print(f"Error generating embeddings: {e}")

product_dataframe.head()

This loop hits the Cohere API once per row. Cohere's free tier rate-limits at 100 calls per minute, so the loop may throttle on a large dataset; if you see repeated errors, add a time.sleep(0.7) inside get_embedding. The embeddings are 1024-dimensional floats. Any row where get_embedding returns None will be silently excluded from vector search — validate that the embedding column has no null values before proceeding to ingestion.

Step 3: Persist embeddings to MongoDB and create a vector search index

With embeddings in memory, push everything to MongoDB. The connection helper validates the link with a ping before returning the client, which surfaces misconfigured URIs immediately rather than at first query time.

# Set MongoDB URI
set_env_securely("MONGO_URI", "Enter your MONGO URI: ")

import pymongo


def get_mongo_client(mongo_uri):
    """Establish and validate connection to the MongoDB."""

    client = pymongo.MongoClient(mongo_uri, appname="devrel.showcase.tavily_mongodb")

    # Validate the connection
    ping_result = client.admin.command("ping")
    if ping_result.get("ok") == 1.0:
        # Connection successful
        print("Connection to MongoDB successful")
        return client
    print("Connection to MongoDB failed")
    return None


MONGO_URI = os.environ["MONGO_URI"]
if not MONGO_URI:
    print("MONGO_URI not set in environment variables")

mongo_client = get_mongo_client(MONGO_URI)

DB_NAME = "amazon_products"
COLLECTION_NAME = "products"

# Create or get the database
db = mongo_client[DB_NAME]

# Create or get the collections
product_collection = db[COLLECTION_NAME]

Clear the collection before re-ingesting — this makes the notebook idempotent so re-runs do not accumulate duplicate documents. Note that delete_many({}) is destructive and instant; in production, use targeted upserts instead.

product_collection.delete_many({})

try:
    documents = product_dataframe.to_dict("records")
    product_collection.insert_many(documents)

    print("Data ingestion into MongoDB completed")
except Exception as e:
    print(f"Error during data ingestion into MongoDB: {e}")

Next, create the Atlas Vector Search index. The setup_vector_search_index function includes a mandatory 30-second sleep because Atlas builds the index asynchronously — queries issued before the index is ready silently fall back to a collection scan, returning wrong results without raising an error.

# The field containing the text embeddings on each document
embedding_field_name = "embedding"
# MongoDB Vector Search index name
vector_search_index_name = "vector_index"

import time

from pymongo.operations import SearchIndexModel


def setup_vector_search_index(collection, index_definition, index_name="vector_index"):
    """
    Setup a vector search index for a MongoDB collection and wait for 30 seconds.

    Args:
    collection: MongoDB collection object
    index_definition: Dictionary containing the index definition
    index_name: Name of the index (default: "vector_index")
    """
    new_vector_search_index_model = SearchIndexModel(
        definition=index_definition, name=index_name, type="vectorSearch"
    )

    # Create the new index
    try:
        result = collection.create_search_index(model=new_vector_search_index_model)
        print(f"Creating index '{index_name}'...")

        # Sleep for 30 seconds
        print(f"Waiting for 30 seconds to allow index '{index_name}' to be created...")
        time.sleep(30)

        print(f"30-second wait completed for index '{index_name}'.")
        return result

    except Exception as e:
        print(f"Error creating new vector search index '{index_name}': {e!s}")
        return None

def create_vector_index_definition(dimensions):
    return {
        "fields": [
            {
                "type": "vector",
                "path": "embedding",
                "numDimensions": dimensions,
                "similarity": "cosine",
            }
        ]
    }

DIMENSIONS = 1024
vector_index_definition = create_vector_index_definition(dimensions=DIMENSIONS)

setup_vector_search_index(product_collection, vector_index_definition, "vector_index")

The numDimensions value must match the model exactly. embed-english-v3.0 always produces 1024-dimensional vectors; if you swap in a different model and forget to update this constant, index creation will fail with a dimension-mismatch error.

Step 4: Query with Tavily hybrid search and persist foreign results

This is where working memory comes alive. TavilyHybridClient wraps both your MongoDB collection (local results, retrieved via vector search) and the Tavily web search API (foreign results, retrieved from the live internet) behind a single .search() call. The content_field parameter tells Tavily which field to display as the result snippet.

# Set up Tavily API Key
set_env_securely("TAVILY_API_KEY", "Enter your Tavily API Key: ")

from tavily import TavilyHybridClient

hybrid_rag = TavilyHybridClient(
    api_key=os.environ.get("TAVILY_API_KEY"),
    db_provider="mongodb",
    collection=product_collection,
    index=vector_search_index_name,
    embeddings_field="embedding",
    content_field="product_semantics",
)

A basic hybrid search returns local and foreign results without writing anything back.

results = hybrid_rag.search(
    "Get me a black laptop to use in a office", max_local=5, max_foreign=2
)

# Create dataframe from the result and view as table
pd.DataFrame(results)

max_local caps how many MongoDB vector-search results are returned; max_foreign caps the live web results. Tuning these numbers is the main lever for balancing freshness against latency — each foreign result requires an outbound HTTP call.

The critical parameter for persistent working memory is save_foreign=True. When set, Tavily writes the web results back into your MongoDB collection, embedding them alongside your existing product records. On the next query, those saved documents are eligible to appear as local results — the collection has genuinely learned from the previous search.

results = hybrid_rag.search(
    "Get me a black laptop to use in a office",
    max_local=5,
    max_foreign=2,
    save_foreign=True,
)

pd.DataFrame(results)

Run the same query a second time with save_foreign=True and watch what happens: the web results saved in the previous round now appear as local results (because they are in MongoDB and indexed), while the foreign slot fetches fresh live data. The collection is acting as a growing, persistent memory that accumulates context across agent steps.

results = hybrid_rag.search(
    "Get me a black laptop to use in a office",
    max_local=5,
    max_foreign=2,
    save_foreign=True,
)

pd.DataFrame(results)

Comparing search parameter strategies

What to watch out for

The 30-second sleep is a lower bound, not a guarantee. Under load, Atlas can take longer to bring a new index online. If your first query returns zero results, poll the index status via the Atlas UI or the list_search_indexes API and wait for the READY state — do not simply increase the sleep duration.

Dimension mismatches fail at query time, not index-creation time. If you create the index with numDimensions=1024 and later embed a query with a different model that outputs 768 dimensions, Atlas Vector Search will return an error at query time. Pin the embedding model name in a constant and reference it from both get_embedding and create_vector_index_definition.

save_foreign=True has no built-in deduplication. Re-running the same query repeatedly inserts the same web documents multiple times, bloating the collection and skewing retrieval toward duplicated content. Implement an upsert strategy — check for a URL or content hash before inserting — if your agent is likely to revisit the same query.

delete_many({}) is destructive and instant. The notebook runs it unconditionally at the top of the ingestion step to keep reruns clean. Dropping and re-inserting documents also forces a full index rebuild. Use this pattern only in development; in production, use targeted upserts.

Cohere free-tier rate limits will stall the embedding loop. The get_embedding function catches exceptions and returns None, but the loop continues to the next row and inserts None into the embedding column. A document with a null embedding is silently ignored by vector search. Add retry logic with backoff if you are working near the rate limit.

Collection growth is unbounded. Every save_foreign=True call adds records. Over many agent sessions this will eventually degrade vector search performance or exceed Atlas tier storage limits. Build a periodic cleanup job — evict documents older than N days or trim to a maximum collection size — before deploying this pattern in production. This is a broader challenge in agent memory calibration that the research community is actively working on.

Where to go next

The pattern established here — MongoDB as a dual-purpose operational and vector store, augmented by live web retrieval — extends naturally in several directions. You can replace the Cohere embedder with a locally hosted model to eliminate per-call cost and latency. You can introduce a reranker between the hybrid results and the agent's LLM context window to improve result ordering. You can layer guardrails on top of the hybrid client to prevent the agent from saving unsafe or off-topic web content back to the collection. For teams already running workloads on managed infrastructure, serverless async pipelines offer a complementary approach to handling embedding and ingestion at scale without blocking the agent's main loop.

The full source notebook, including additional context on memory taxonomies and the sales-assistant framing, is at the MongoDB GenAI Showcase (MIT licence).