app.py

import time
import traceback
import gradio as gr
import os
import asyncio
from pymongo import MongoClient
from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings
from langchain_community.llms import OpenAI
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

output_parser = StrOutputParser()

import json

## uri removed
## 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 = 'embedded_movies'
collection = client[db_name][collection_name]

try:
    vector_store = MongoDBAtlasVectorSearch(embedding=OpenAIEmbeddings(), collection=collection, index_name='vector_index', text_key='plot', embedding_key='plot_embedding')
    llm = ChatOpenAI(temperature=0)
    prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a movie recommendation engine which post a concise and short summary on relevant movies."),
    ("user", "List of movies: {input}")
     ])
    chain = prompt | llm | output_parser

except:
    #If open ai key is wrong
    print ('Open AI key is wrong')
    vector_store = None
    print("An error occurred: \n" + error_message)

def get_movies(message, history):

    try:
        movies =  vector_store.similarity_search(query=message, k=3, embedding_key='plot_embedding')
        return_text = ''
        for movie in movies:
            return_text = return_text + 'Title : ' +  movie.metadata['title'] + '\n------------\n' + 'Plot: ' + movie.page_content + '\n\n'
    
        print_llm_text = chain.invoke({"input": return_text})
    
        for i in range(len(print_llm_text)):
            time.sleep(0.05)
            yield "Found: " + "\n\n" + print_llm_text[: i+1]
    except Exception as e:
        error_message = traceback.format_exc()
        print("An error occurred: \n" + error_message)
        yield "Please clone the repo and add your open ai key as well as your MongoDB Atlas URI in the Secret Section of you Space\n OPENAI_API_KEY (your Open AI key) and MONGODB_ATLAS_CLUSTER_URI (0.0.0.0/0 whitelisted instance with Vector index created) \n\n For more information : https://mongodb.com/products/platform/atlas-vector-search"
        
    

demo = gr.ChatInterface(get_movies, examples=["What movies are scary?", "Find me a comedy", "Movies for kids"], title="Movies Atlas Vector Search",description="This small chat uses a similarity search to find relevant movies, it uses MongoDB Atlas Vector Search read more here: https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-tutorial",submit_btn="Search").queue()

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