DEV Community

AIRabbit
AIRabbit

Posted on

Building Semantic Search with MongoDB Atlas Vector Search

MongoDB Atlas now supports vector search capabilities, enabling powerful semantic search applications. In this post, we'll explore how to implement neural search using MongoDB Atlas and OpenAI embeddings, with a practical example of a movie search application.

What is Neural Search?

Neural search, also known as semantic search, goes beyond traditional keyword matching by understanding the meaning and context of search queries. It uses vector embeddings to represent text as high-dimensional vectors, allowing for similarity-based matching that captures semantic relationships.

Try it out

Want to experiment with MongoDB Atlas vector search before diving into the implementation? Try our live demo on Hugging Face Spaces: MongoDB Vector Search Util

To get started:

  1. Clone the repository
  2. Set up your environment variables:
  3. ATLAS_URI: Your MongoDB Atlas connection string
  4. OPENAI_API_KEY: Your OpenAI API key

The demo provides two tabs:

  1. Embedding Generation: Start here to generate embeddings for your documents
  2. Search: Once embeddings are created, use this tab to perform semantic searches

First we do the embedding. For testing purposes, you can just embed a limited number of focal points, say 10. Then you can check the generated embedding fields with a query like this { "embedding": { "$exists": true }}. If all goes well, you can set the number to 0 to embed all entries in the DB.

Bug: If the UI does not show the collection for the selected DB, select another DB (such as admin) and then switch back to the DB you want to index (e.g. sample*mflix)*

Then we have to create the index in atlas.

Finally, do a neural search using this index (make sure the new index is fully indexed, otherwise you will not be able to find any matches).

Now, let's dive into how it all works!

Prerequisites

  • MongoDB Atlas account with vector search enabled
  • OpenAI API key
  • Python 3.8+
  • Sample movie data (we'll use the sample_mflix database)

Implementation Overview

Our implementation consists of two main components:

  1. Embedding Generation: Converting text data into vector embeddings
  2. Vector Search: Performing similarity-based searches using these embeddings

Step 1: Setting Up Vector Search in MongoDB Atlas

First, we need to create a vector search index in MongoDB Atlas. This index will store our embeddings and enable efficient similarity searches:

{
  "fields": [
    {
      "type": "vector",
      "path": "plot_embedding",
      "numDimensions": 1536,
      "similarity": "dotProduct"
    }
  ]
}
Enter fullscreen mode Exit fullscreen mode

Step 2: Generating Embeddings

We use OpenAI's text-embedding-ada-002 model to generate embeddings. Here's how we handle the embedding generation:

def get_embedding(text: str, openai_client, model="text-embedding-ada-002", max_retries=3) -> list[float]:
    """Get embeddings for given text using OpenAI API with retry logic"""
    text = text.replace("\n", " ")

    for attempt in range(max_retries):
        try:
            resp = openai_client.embeddings.create(
                input=[text],
                model=model
            )
            return resp.data[0].embedding
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            time.sleep(2 ** attempt)  # Exponential backoff
Enter fullscreen mode Exit fullscreen mode

For efficiency, we process documents in parallel batches:

def parallel_generate_embeddings(
    collection: Collection,
    cursor,
    field_name: str,
    embedding_field: str,
    openai_client,
    total_docs: int,
    batch_size: int = 10
) -> int:
    """Generate embeddings in parallel using ThreadPoolExecutor"""
    with ThreadPoolExecutor(max_workers=10) as executor:
        # Process documents in batches
        batch = []
        for doc in cursor:
            batch.append(doc)
            if len(batch) >= batch_size:
                future = executor.submit(process_batch, batch.copy(), 
                                      field_name, embedding_field, openai_client)
                # ... process futures and update progress
Enter fullscreen mode Exit fullscreen mode

Step 3: Implementing Vector Search

Once we have our embeddings, we can perform semantic searches using MongoDB's $vectorSearch operator:

def vector_search(query_text: str, db_name: str, collection_name: str, 
                 embedding_field: str, index_name: str) -> str:
    """Perform vector search using embeddings"""
    # Generate embedding for search query
    embedding = get_embedding(query_text, openai_client)

    # Perform vector search
    results = collection.aggregate([
        {
            '$vectorSearch': {
                "index": index_name,
                "path": embedding_field,
                "queryVector": embedding,
                "numCandidates": 50,
                "limit": 5
            }
        },
        {
            "$project": {
                "search_score": { "$meta": "vectorSearchScore" },
                "document": "$$ROOT"
            }
        }
    ])
Enter fullscreen mode Exit fullscreen mode

Example: Movie Search Application

Let's look at how this works in practice with a movie search application. The application allows users to search through movie plots using natural language queries.

AI Rabbit News

AI News & Tutorials

favicon airabbit.blog

Sample Queries and Results

  1. "humans fighting aliens"
  2. This might return sci-fi movies like "Independence Day" or "War of the Worlds"
  3. "relationship drama between two good friends"
  4. Could match movies like "When Harry Met Sally" or "Good Will Hunting"
  5. "comedy about family vacation"
  6. Might return movies like "National Lampoon's Vacation" or "Little Miss Sunshine"

The search results are ranked by similarity score, with higher scores indicating better semantic matches to the query.

Performance Optimizations

Our implementation includes several optimizations:

  1. Parallel Processing Uses ThreadPoolExecutor for concurrent embedding generation Dynamically adjusts batch size and worker count based on performance
  2. Error Handling Implements retry logic with exponential backoff Gracefully handles API rate limits and failures
  3. Memory Efficiency Processes documents in batches to manage memory usage Uses cursor-based iteration for large collections
  4. Progress Tracking Real-time progress updates during embedding generation Supports processing of large collections (100,000+ documents)

Conclusion

MongoDB Atlas's vector search capabilities, combined with OpenAI's embeddings, provide a powerful foundation for building semantic search applications. This approach enables natural language queries that understand context and meaning, going beyond simple keyword matching.

AWS Security LIVE!

Tune in for AWS Security LIVE!

Join AWS Security LIVE! for expert insights and actionable tips to protect your organization and keep security teams prepared.

Learn More

Top comments (0)

AWS Security LIVE!

Tune in for AWS Security LIVE!

Join AWS Security LIVE! for expert insights and actionable tips to protect your organization and keep security teams prepared.

Learn More

👋 Kindness is contagious

Discover a treasure trove of wisdom within this insightful piece, highly respected in the nurturing DEV Community enviroment. Developers, whether novice or expert, are encouraged to participate and add to our shared knowledge basin.

A simple "thank you" can illuminate someone's day. Express your appreciation in the comments section!

On DEV, sharing ideas smoothens our journey and strengthens our community ties. Learn something useful? Offering a quick thanks to the author is deeply appreciated.

Okay