How to Do Semantic Search in MongoDB Using Atlas Vector Search

1
exports = async function(changeEvent) {
2
    // Get the full document from the change event.
3
    const doc = changeEvent.fullDocument;
4
5
    // Define the OpenAI API url and key.
6
    const url = 'https://api.openai.com/v1/embeddings';
7
    // Use the name you gave the value of your API key in the "Values" utility inside of App Services
8
    const openai_key = context.values.get("openAI_value");
9
    try {
10
        console.log(`Processing document with id: ${doc._id}`);
11
12
        // Call OpenAI API to get the embeddings.
13
        let response = await context.http.post({
14
            url: url,
15
             headers: {
16
                'Authorization': [`Bearer ${openai_key}`],
17
                'Content-Type': ['application/json']
18
            },
19
            body: JSON.stringify({
20
                // The field inside your document that contains the data to embed, here it is the "plot" field from the sample movie data.
21
                input: doc.plot,
22
                model: "text-embedding-ada-002"
23
            })
24
        });
25
26
        // Parse the JSON response
27
        let responseData = EJSON.parse(response.body.text());
28
29
        // Check the response status.
30
        if(response.statusCode === 200) {
31
            console.log("Successfully received embedding.");
32
33
            const embedding = responseData.data[0].embedding;
34
35
            // Use the name of your MongoDB Atlas Cluster
36
            const collection = context.services.get("<CLUSTER_NAME>").db("sample_mflix").collection("movies");
37
38
            // Update the document in MongoDB.
39
            const result = await collection.updateOne(
40
                { _id: doc._id },
41
                // The name of the new field you'd like to contain your embeddings.
42
                { $set: { plot_embedding: embedding }}
43
            );
44
45
            if(result.modifiedCount === 1) {
46
                console.log("Successfully updated the document.");
47
            } else {
48
                console.log("Failed to update the document.");
49
            }
50
        } else {
51
            console.log(`Failed to receive embedding. Status code: ${response.statusCode}`);
52
        }
53
54
    } catch(err) {
55
        console.error(err);
56
    }
57
};

1
{
2
  "type": "vectorSearch",
3
  "fields": [{
4
    "path": "plot_embedding",
5
    "dimensions": 1536,
6
    "similarity": "cosine",
7
    "type": "vector"
8
  }]
9
}

1
const axios = require('axios');
2
const MongoClient = require('mongodb').MongoClient;
3
4
async function getEmbedding(query) {
5
    // Define the OpenAI API url and key.
6
    const url = 'https://api.openai.com/v1/embeddings';
7
    const openai_key = 'your_openai_key'; // Replace with your OpenAI key.
8
    
9
    // Call OpenAI API to get the embeddings.
10
    let response = await axios.post(url, {
11
        input: query,
12
        model: "text-embedding-ada-002"
13
    }, {
14
        headers: {
15
            'Authorization': `Bearer ${openai_key}`,
16
            'Content-Type': 'application/json'
17
        }
18
    });
19
    
20
    if(response.status === 200) {
21
        return response.data.data[0].embedding;
22
    } else {
23
        throw new Error(`Failed to get embedding. Status code: ${response.status}`);
24
    }
25
}
26
27
async function findSimilarDocuments(embedding) {
28
    const url = 'your_mongodb_url'; // Replace with your MongoDB url.
29
    const client = new MongoClient(url);
30
    
31
    try {
32
        await client.connect();
33
        
34
        const db = client.db('<DB_NAME>'); // Replace with your database name.
35
        const collection = db.collection('<COLLECTION_NAME>'); // Replace with your collection name.
36
        
37
        // Query for similar documents.
38
        const documents = await collection.aggregate([
39
  {"$vectorSearch": {
40
    "queryVector": embedding,
41
    "path": "plot_embedding",
42
    "numCandidates": 100,
43
    "limit": 5,
44
    "index": "moviesPlotIndex",
45
      }}
46
]).toArray();
47
        
48
        return documents;
49
    } finally {
50
        await client.close();
51
    }
52
}
53
54
async function main() {
55
    const query = 'your_query'; // Replace with your query.
56
    
57
    try {
58
        const embedding = await getEmbedding(query);
59
        const documents = await findSimilarDocuments(embedding);
60
        
61
        console.log(documents);
62
    } catch(err) {
63
        console.error(err);
64
    }
65
}
66
67
main();