Create app.py

app.py

@@ -0,0 +1,118 @@
+import streamlit as st
+from time import sleep
+import json
+from pymongo import MongoClient
+from bson import ObjectId
+from openai import OpenAI
+openai_client = OpenAI()
+import os
+
+# Get the restaurants based on the search and location
+def get_restaurants(search, location, meters):
+    try:
+        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']
+    except:
+        st.error("Error Connecting to the MongoDB Atlas Cluster")
+        return None, None, None, None
+
+    try:
+        newTrip, pre_agg = pre_aggregate_meters(restaurants_collection, 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-4o",
+            messages=[
+                {"role": "system", "content": "You are a helpful restaurant assistant. Answer shortly and quickly. You will get a context if the context is not relevant to the user query please address that and not provide by default the restaurants as is."},
+                {"role": "user", "content": f"Find me the 2 best restaurant and why based on {search} and {restaurant_docs}. Shortly explain trades offs and why I should go to each one. You can mention the third option as a possible alternative in one sentence."}
+            ]
+        )
+
+        trips_collection.delete_many({"searchTrip": newTrip})
+
+        if len(restaurant_docs) == 0:
+            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>', str(pre_agg), str(vectorQuery)
+
+        first_restaurant = restaurant_docs[0]['restaurant_id']
+        second_restaurant = restaurant_docs[1]['restaurant_id']
+        third_restaurant = restaurant_docs[2]['restaurant_id']
+        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>'
+        client.close()
+        return chat_response.choices[0].message.content, iframe, str(pre_agg), str(vectorQuery)
+    except Exception as e:
+        st.error(f"Your query caused an error: {e}")
+        return "Your query caused an error, please retry with allowed input only ...", '<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>', str(pre_agg), str(vectorQuery)
+
+def pre_aggregate_meters(restaurants_collection, 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)
+    sleep(3)
+    return tripId, pre_aggregate_pipeline
+
+st.markdown(
+    """
+    # MongoDB's Vector Restaurant Planner 
+    Start typing below to see the results. You can search a specific cuisine for you and choose 3 predefined locations.
+
+    The radius specifies the distance from the start search location. This space uses the dataset called [whatscooking.restaurants](https://huggingface.co/datasets/AIatMongoDB/whatscooking.restaurants)
+    """
+)
+
+search = st.text_input("What type of dinner are you looking for?")
+location = st.radio("Location", options=[
+    {"label": "Timesquare Manhattan", "value": {"type": "Point", "coordinates": [-73.98527039999999, 40.7589099]}},
+    {"label": "Westside Manhattan", "value": {"type": "Point", "coordinates": [-74.013686, 40.701975]}},
+    {"label": "Downtown Manhattan", "value": {"type": "Point", "coordinates": [-74.000468, 40.720777]}}
+], format_func=lambda x: x['label'])
+meters = st.slider("Radius in meters", min_value=500, max_value=10000, step=5)
+
+if st.button("Get Restaurants"):
+    location_value = location['value']
+    result, iframe, pre_agg, vectorQuery = get_restaurants(search, location_value, meters)
+    if result:
+        st.markdown(result)
+        st.markdown(iframe, unsafe_allow_html=True)
+        st.json(pre_agg)
+        st.json(vectorQuery)