Pavel Duchovny
init
c63f1bb
raw
history blame
5.1 kB
import gradio as gr
from time import sleep
import json
from pymongo import MongoClient
from bson import ObjectId
from openai import OpenAI
openai_client = OpenAI()
import os
uri = os.environ.get('MONGODB_ATLAS_URI')
client = MongoClient(uri)
db_name = 'whatscooking'
collection_name = 'restaurants'
restaurants_collection = client[db_name][collection_name]
trips_collection = client[db_name]['smart_trips']
def get_restaurants(search, location, meters):
newTrip, pre_agg = pre_aggregate_meters(location, meters)
response = openai_client.embeddings.create(
input=search,
model="text-embedding-3-small",
dimensions=256
)
vectorQuery = {
"$vectorSearch": {
"index" : "vector_index",
"queryVector": response.data[0].embedding,
"path" : "embedding",
"numCandidates": 10,
"limit": 3,
"filter": {"searchTrip": newTrip}
}}
restaurant_docs = list(trips_collection.aggregate([vectorQuery,
{"$project": {"_id" : 0, "embedding": 0}}]))
chat_response = openai_client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a helpful restaurant assistant."},
{ "role": "user", "content": f"Find me the 2 best restaurant and why based on {search} and {restaurant_docs}. explain trades offs and why I should go to each one. You can mention the third option as a possible alternative."}
]
)
trips_collection.delete_many({"searchTrip": newTrip})
first_restaurant = restaurant_docs[0]['restaurant_id']
second_restaurant = restaurant_docs[1]['restaurant_id']
third_restaurant = restaurant_docs[2]['restaurant_id']
if (first_restaurant or second_restaurant or third_restaurant) is None:
return "No restaurants found", '<iframe style="background: #FFFFFF;border: none;border-radius: 2px;box-shadow: 0 2px 10px 0 rgba(70, 76, 79, .2);" width="640" height="480" src="https://charts.mongodb.com/charts-paveldev-wiumf/embed/charts?id=65c24b0c-2215-4e6f-829c-f484dfd8a90c&filter={\'restaurant_id\':\'\'}&maxDataAge=3600&theme=light&autoRefresh=true"></iframe>'
else:
restaurant_string = f"\'{first_restaurant}\', \'{second_restaurant}\', \'{third_restaurant}\'"
iframe = '<iframe style="background: #FFFFFF;border: none;border-radius: 2px;box-shadow: 0 2px 10px 0 rgba(70, 76, 79, .2);" width="640" height="480" src="https://charts.mongodb.com/charts-paveldev-wiumf/embed/charts?id=65c24b0c-2215-4e6f-829c-f484dfd8a90c&filter={\'restaurant_id\':{$in:[' + restaurant_string + ']}}&maxDataAge=3600&theme=light&autoRefresh=true"></iframe>'
return chat_response.choices[0].message.content, iframe,str(pre_agg), str(vectorQuery)
def pre_aggregate_meters(location, meters):
tripId = ObjectId()
pre_aggregate_pipeline = [{
"$geoNear": {
"near": location,
"distanceField": "distance",
"maxDistance": meters,
"spherical": True,
},
},
{
"$addFields": {
"searchTrip" : tripId,
"date" : tripId.generation_time
}
},
{
"$merge": {
"into": "smart_trips"
}
} ]
result = restaurants_collection.aggregate(pre_aggregate_pipeline);
print(trips_collection.count_documents({"searchTrip": tripId}));
sleep(1)
return tripId, pre_aggregate_pipeline
with gr.Blocks() as demo:
gr.Markdown(
"""
# MongoDB's Vector Restaurant planner
Start typing below to see the results
""")
#gr.HTML(value='<iframe style="background: #FFFFFF;border: none;border-radius: 2px;box-shadow: 0 2px 10px 0 rgba(70, 76, 79, .2);" width="640" height="480" src="https://charts.mongodb.com/charts-paveldev-wiumf/embed/charts?id=65c24b0c-2215-4e6f-829c-f484dfd8a90c&maxDataAge=3600&theme=light&autoRefresh=true"></iframe>')
#
gr.Interface(
get_restaurants,
[
gr.Textbox(placeholder="What type of dinner are you looking for?"),
gr.Radio([("work",{
"type": "Point",
"coordinates": [
-73.98527039999999,
40.7589099
]
}), ("home",{
"type": "Point",
"coordinates": [
-74.013686, 40.701975
]
}), ("park", {
"type": "Point",
"coordinates": [ -74.000468,40.720777
]
})], label="Location", info="What location you need?"),
gr.Slider(minimum=500, maximum=10000, randomize=False, step=5, label="Radius in meters")],
[gr.Textbox(label="MongoDB Vector Recommendations", placeholder="Results will be displayed here"), "html",
gr.Code(label="Pre-aggregate pipeline",language="json" ),
gr.Code(label="Vector Query", language="json")],
)
#radio.change(location_searched, loc, out)
if __name__ == "__main__":
demo.launch()