File size: 1,844 Bytes
8f835ef
 
90f787c
8f835ef
 
e2275e6
 
 
 
 
8f835ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import time
import gradio as gr
import os
import asyncio
from pymongo import MongoClient
from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langchain_community.embeddings import OpenAIEmbeddings
# from langchain_community.llms import OpenAI
# from langchain_community.prompts import PromptTemplate
# from langchain.chains import LLMChain
import json


## Connect to MongoDB Atlas local cluster
MONGODB_ATLAS_CLUSTER_URI = os.getenv('MONGODB_ATLAS_CLUSTER_URI')
client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
db_name = 'sample_mflix'
collection_name = 'movies'
collection = client[db_name][collection_name]

## Create a collection and insert data
print ('Creating collection movies')
client[db_name].create_collection(collection_name)

## Create a vector search index
print ('Creating vector search index')
collection.create_search_index(model={"definition": {"mappings":{
  "dynamic":True,
  "fields": {
    "plot_embedding": {
      "type": "knnVector",
      "dimensions": 1536,
      "similarity": "euclidean"
    }
  }
}}, "name":'default'})

# sleep for minute
print ('Waiting for vector index on field "embedding" to be created')
time.sleep(60)

vector_store = MongoDBAtlasVectorSearch(embedding=OpenAIEmbeddings(), collection=collection, index_name='default', text_key='plot', embedding_key='plot_embedding')

def get_movies(message, history):
    movies = vector_store.similarity_search(message, 3)
    for movie in movies:
        for i in range(len(movie.metadata['title'])):
            time.sleep(0.05)
            yield "Movie " + i + " :  Title - " + movie.metadata['title'][: i+1]

demo = gr.ChatInterface(get_movies, examples=["What movies are scary?", "Find me a comedy", "Movies for kids"], title="Movies Atlas Vector Search", submit_btn="Search").queue()

if __name__ == "__main__":
    demo.launch()