Parte #3: pesquise seus dados semanticamente com o Atlas Vector Search do MongoDB

1
{
2
  "fields": [
3
    {
4
      "numDimensions": 384,
5
      "path": "embedding",
6
      "similarity": "cosine",
7
      "type": "vector"
8
    }
9
  ]
10
}

1
MONGODB_CONNECTION_STRING="mongodb+srv://<username>:<password>@<cluster_id>.mongodb.net/?retryWrites=true&w=majority"
2
EMBEDDING_SERVICE="https://<your_endpoint>.amazonaws.com/TEST/sageMakerResource"

1
import os
2
from typing import Optional
3
from urllib.parse import quote
4
5
import requests
6
from dotenv import load_dotenv
7
8
load_dotenv()
9
10
EMBEDDING_SERVICE = os.environ.get("EMBEDDING_SERVICE")

1
def create_embedding(plot: str) -> Optional[float]:
2
    encoded_plot = quote(plot)
3
    embedding_url = f"{EMBEDDING_SERVICE}?query={encoded_plot}"
4
5
    embedding_response = requests.get(embedding_url)
6
    embedding_vector = embedding_response.json()["embedding"]
7
8
    return embedding_vector

1
import os
2
3
from dotenv import load_dotenv
4
from pymongo import MongoClient, UpdateOne
5
6
from sagemaker import create_embedding
7
8
load_dotenv()
9
10
MONGODB_CONNECTION_STRING = os.environ.get("MONGODB_CONNECTION_STRING")
11
DATABASE_NAME = "sample_mflix"
12
COLLECTION_NAME = "embedded_movies"
13
VECTOR_SEARCH_INDEX_NAME = "VectorSearchIndex"
14
EMBEDDING_PATH = "embedding"
15
mongo_client = MongoClient(MONGODB_CONNECTION_STRING)
16
database = mongo_client[DATABASE_NAME]
17
movies_collection = database[COLLECTION_NAME]

1
def add_missing_embeddings():
2
    movies_with_a_plot_without_embedding_filter = {
3
        "$and": [
4
            {"plot": {"$exists": True, "$ne": ""}},
5
            {"embedding": {"$exists": False}},
6
        ]
7
    }
8
    only_show_plot_projection = {"plot": 1}
9
10
    movies = movies_collection.find(
11
        movies_with_a_plot_without_embedding_filter,
12
        only_show_plot_projection,
13
    )
14
15
    movies_to_update = []
16
17
    for movie in movies:
18
        embedding = create_embedding(movie["plot"])
19
        update_operation = UpdateOne(
20
            {"_id": movie["_id"]},
21
            {"$set": {"embedding": embedding}},
22
        )
23
        movies_to_update.append(update_operation)
24
25
    if movies_to_update:
26
        result = movies_collection.bulk_write(movies_to_update)
27
        print(f"Updated {result.modified_count} movies")
28
29
    else:
30
        print("No movies to update")

1
def update_plot(title: str, plot: str) -> dict:
2
    embedding = create_embedding(plot)
3
4
    result = movies_collection.find_one_and_update(
5
        {"title": title},
6
        {"$set": {"plot": plot, "embedding": embedding}},
7
        return_document=True,
8
    )
9
10
    return result

1
def execute_vector_search(vector: [float]) -> list[dict]:
2
    vector_search_query = {
3
        "$vectorSearch": {
4
            "index": VECTOR_SEARCH_INDEX_NAME,
5
            "path": EMBEDDING_PATH,
6
            "queryVector": vector,
7
            "numCandidates": 10,
8
            "limit": 5,
9
        }
10
    }
11
    projection = {"$project": {"_id": 0, "title": 1}}
12
    results = movies_collection.aggregate([vector_search_query, projection])
13
    results_list = list(results)
14
15
    return results_list

1
from flask import Flask, request, jsonify
2
3
from atlas import execute_vector_search, update_plot
4
from sagemaker import create_embedding
5
6
app = Flask(__name__)

1
@app.route("/movies/<title>", methods=["PUT"])
2
def update_movie(title: str):
3
    try:
4
        request_json = request.get_json()
5
        plot = request_json["plot"]
6
        updated_movie = update_plot(title, plot)
7
8
        if updated_movie:
9
            return jsonify(
10
                {
11
                    "message": "Movie updated successfully",
12
                    "updated_movie": updated_movie,
13
                }
14
            )
15
        else:
16
            return jsonify({"error": f"Movie with title {title} not found"}), 404
17
18
    except Exception as e:
19
        return jsonify({"error": str(e)}), 500

1
@app.route("/movies/search", methods=["POST"])
2
def search_movies():
3
    try:
4
        request_json = request.get_json()
5
        query = request_json["query"]
6
        embedding = create_embedding(query)
7
8
        results = execute_vector_search(embedding)
9
10
        jsonified_results = jsonify(
11
            {
12
                "message": "Movies searched successfully",
13
                "results": results,
14
            }
15
        )
16
17
        return jsonified_results
18
19
    except Exception as e:
20
        return jsonify({"error": str(e)}), 500
21
22
23
if __name__ == "__main__":
24
    app.run(debug=True)

1
curl -X POST -H "Content-Type: application/json" -d '{"query": "A movie about the Earth, Mars and an invasion."}' http://127.0.0.1:5000/movies/search

1
{
2
  "message": "Movies searched successfully",
3
  "results": [
4
    {
5
      "title": "The War of the Worlds"
6
    },
7
    {
8
      "title": "The 6th Day"
9
    },
10
    {
11
      "title": "Pixels"
12
    },
13
    {
14
      "title": "Journey to Saturn"
15
    },
16
    {
17
      "title": "Moonraker"
18
    }
19
  ]
20
}