Início rápido 2: pesquisa vetorial com MongoDB e OpenAI

1
{
2
  "fields": [
3
    {
4
      "type": "vector",
5
      "path": "plot_embedding",
6
      "numDimensions": 1536,
7
      "similarity": "dotProduct"
8
    }
9
  ]
10
}

1
# We will keep all global variables in an object to not pollute the global namespace.
2
class MyConfig(object):
3
   pass
4
5
MY_CONFIG = MyConfig()
6
7
MY_CONFIG.ATLAS_URI = "Enter your Atlas URI value here"  ## TODO
8
MY_CONFIG.OPENAI_API_KEY = "Enter your OpenAI API Key here"  ## TODO

1
!pip install   openai==1.13.3     pymongo==4.6.2

1
from pymongo import MongoClient
2
3
4
class AtlasClient ():
5
6
7
   def __init__ (self, altas_uri, dbname):
8
       self.mongodb_client = MongoClient(altas_uri)
9
       self.database = self.mongodb_client[dbname]
10
11
12
   ## A quick way to test if we can connect to Atlas instance
13
   def ping (self):
14
       self.mongodb_client.admin.command('ping')
15
16
17
   def get_collection (self, collection_name):
18
       collection = self.database[collection_name]
19
       return collection
20
21
22
   def find (self, collection_name, filter = {}, limit=10):
23
       collection = self.database[collection_name]
24
       items = list(collection.find(filter=filter, limit=limit))
25
       return items
26
27
28
   # https://www.mongodb.com/pt-br/docs/atlas/atlas-vector-search/vector-search-stage/
29
   def vector_search(self, collection_name, index_name, attr_name, embedding_vector, limit=5):
30
       collection = self.database[collection_name]
31
       results = collection.aggregate([
32
           {
33
               '$vectorSearch': {
34
                   "index": index_name,
35
                   "path": attr_name,
36
                   "queryVector": embedding_vector,
37
                   "numCandidates": 50,
38
                   "limit": limit,
39
               }
40
           },
41
           ## We are extracting 'vectorSearchScore' here
42
           ## columns with 1 are included, columns with 0 are excluded
43
           {
44
               "$project": {
45
                   '_id' : 1,
46
                   'title' : 1,
47
                   'plot' : 1,
48
                   'year' : 1,
49
                   "search_score": { "$meta": "vectorSearchScore" }
50
           }
51
           }
52
           ])
53
       return list(results)
54
55
56
   def close_connection(self):
57
       self.mongodb_client.close()

1
 results = collection.aggregate([
2
           {
3
               '$vectorSearch': {
4
                   "index": index_name,
5
                   "path": attr_name,
6
                   "queryVector": embedding_vector,
7
                   "numCandidates": 50,
8
                   "limit": limit,
9
               }
10
           },
11
           ## We are extracting 'vectorSearchScore' here
12
           ## columns with 1 are included, columns with 0 are excluded
13
           {
14
               "$project": {
15
                   '_id' : 1,
16
                   'title' : 1,
17
                   'plot' : 1,
18
                   'year' : 1,
19
                   "search_score": { "$meta": "vectorSearchScore" }
20
           }
21
           }
22
           ])

1
 "search_score": { "$meta": "vectorSearchScore" }

1
from openai import OpenAI
2
3
4
class OpenAIClient():
5
   def __init__(self, api_key) -> None:
6
       self.client = OpenAI(
7
           api_key= api_key,  # defaults to os.environ.get("OPENAI_API_KEY")
8
       )
9
       # print ("OpenAI Client initialized!")
10
11
12
13
14
   def get_embedding(self, text: str,  model="text-embedding-ada-002") -> list[float]:
15
       text = text.replace("\n", " ")
16
       resp = self.client.embeddings.create (
17
           input=[text],
18
           model=model  )
19
20
21
       return resp.data[0].embedding

1
MY_CONFIG.DB_NAME = 'sample_mflix'
2
MY_CONFIG.COLLECTION_NAME = 'embedded_movies'
3
MY_CONFIG.INDEX_NAME = 'idx_plot_embedding'
4
5
6
atlas_client = AtlasClient (MY_CONFIG.ATLAS_URI, MY_CONFIG.DB_NAME)
7
atlas_client.ping()
8
print ('Connected to Atlas instance! We are good to go!')

1
openAI_client = OpenAIClient (api_key=MY_CONFIG.OPENAI_API_KEY)
2
print ("OpenAI client initialized")

1
import time
2
3
# Handy function
4
def do_vector_search (query:str) -> None:
5
   query = query.lower().strip()  # cleanup query string
6
   print ('query: ', query)
7
8
9
   # call openAI API to convert text into embedding
10
   t1a = time.perf_counter()
11
   embedding = openAI_client.get_embedding(query)
12
   t1b = time.perf_counter()
13
   print (f"Getting embeddings from OpenAI took {(t1b-t1a)*1000:,.0f} ms")
14
15
16
   # perform a vector search on Atlas
17
   # using embeddings (returned from OpenAI above) 
18
   t2a = time.perf_counter()
19
   movies = atlas_client.vector_search(collection_name=MY_CONFIG.COLLECTION_NAME, index_name=MY_CONFIG.INDEX_NAME, attr_name='plot_embedding', embedding_vector=embedding,limit=10 )
20
   t2b = time.perf_counter()
21
22
23
   # and printing out the results
24
   print (f"Altas query returned {len (movies)} movies in {(t2b-t2a)*1000:,.0f} ms")
25
   print()
26
27
28
   for idx, movie in enumerate (movies):
29
       print(f'{idx+1}\nid: {movie["_id"]}\ntitle: {movie["title"]},\nyear: {movie["year"]}' +
30
           f'\nsearch_score(meta):{movie["search_score"]}\nplot: {movie["plot"]}\n')

1
query="humans fighting aliens"
2
do_vector_search (query=query)

1
query:  humans fighting aliens
2
using cached embeddings
3
Altas query returned 10 movies in 138 ms
4
5
6
1
7
id: 573a1398f29313caabce8f83
8
title: V: The Final Battle,
9
year: 1984
10
search_score(meta):0.9573556184768677
11
plot: A small group of human resistance fighters fight a desperate guerilla war against the genocidal extra-terrestrials who dominate Earth.
12
13
14
2
15
id: 573a13c7f29313caabd75324
16
title: Falling Skies,
17
year: 2011è
18
search_score(meta):0.9550596475601196
19
plot: Survivors of an alien attack on earth gather together to fight for their lives and fight back.
20
21
22
3
23
id: 573a139af29313caabcf0cff
24
title: Starship Troopers,
25
year: 1997
26
search_score(meta):0.9523435831069946
27
plot: Humans in a fascistic, militaristic future do battle with giant alien bugs in a fight for survival.
28
29
30
...
31
year: 2002
32
search_score(meta):0.9372057914733887
33
plot: A young woman from the future forces a local gunman to help her stop an impending alien invasion which will wipe out the human race.

1
query="relationship drama between two good friends"
2
do_vector_search (query=query)

1
query:  relationship drama between two good friends
2
using cached embeddings
3
Altas query returned 10 movies in 71 ms
4
5
6
1
7
id: 573a13a3f29313caabd0dfe2
8
title: Dark Blue World,
9
year: 2001
10
search_score(meta):0.9380425214767456
11
plot: The friendship of two men becomes tested when they both fall for the same woman.
12
13
14
2
15
id: 573a13a3f29313caabd0e14b
16
title: Dark Blue World,
17
year: 2001
18
search_score(meta):0.9380425214767456
19
plot: The friendship of two men becomes tested when they both fall for the same woman.
20
21
22
3
23
id: 573a1399f29313caabcec488
24
title: Once a Thief,
25
year: 1991
26
search_score(meta):0.9260045289993286
27
plot: A romantic and action packed story of three best friends, a group of high end art thieves, who come into trouble when a love-triangle forms between them.
28
29
30
...
31
year: 1987
32
search_score(meta):0.9181452989578247
33
plot: A modern day Romeo & Juliet story is told in New York when an Italian boy and a Chinese girl become lovers, causing a tragic conflict between ethnic gangs.