Enable Generative AI and Semantic Search Capabilities on Your Database With MongoDB Atlas and OpenAI

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.description,
22
                model: "text-embedding-3-small"
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("AtlasSearch").db("sample_airbnb").collection("listingsAndReviews");
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: { description_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
import pymongo
2
import requests
3
import pprint
4
5
def get_vector_embeddings_from_openai(query):
6
    openai_api_url = "https://api.openai.com/v1/embeddings"
7
    openai_api_key = "<your-open-ai-api-key>"
8
9
    data = {
10
        'input': query,
11
        'model': "text-embedding-3-small"
12
    }
13
14
    headers = {
15
        'Authorization': 'Bearer {0}'.format(openai_api_key),
16
        'Content-Type': 'application/json'
17
    }
18
19
    response = requests.post(openai_api_url, json=data, headers=headers)
20
    embedding = []
21
    if response.status_code == 200:
22
        embedding = response.json()['data'][0]['embedding']
23
    return embedding
24
25
def find_similar_documents(embedding):
26
    mongo_url = 'mongodb+srv://<username>:<password>@<cluster-url.mongodb.net>/?retryWrites=true&w=majority'
27
    client = pymongo.MongoClient(mongo_url)
28
    db = client.sample_airbnb
29
    collection = db["listingsAndReviews"]
30
31
    pipeline = [
32
        {
33
            "$vectorSearch": {
34
                "index": "default",
35
                "path": "descriptions_embedding",
36
                “queryVector”: “embedding”,
37
                “numCandidates”: 150, 
38
                “limit”: 10
39
             }
40
        },
41
        {
42
            "$project": {
43
                "_id": 0,
44
                "description": 1
45
            }
46
        }
47
    ]
48
    documents = collection.aggregate(pipeline)
49
    return documents
50
51
def main():
52
    query = "Best for couples, nearby beach area with cool weather"
53
    try:
54
        embedding = get_vector_embeddings_from_openai(query)
55
        documents = find_similar_documents(embedding)
56
        print("Documents")
57
        pprint.pprint(list(documents))
58
    except Exception as e:
59
        print("Error occured: {0}".format(e))
60
61
main()