Update app.py

app.py

@@ -1,11 +1,15 @@
 import gradio as gr
+import pandas as pd
 import logging, os, sys, threading
 
+from datasets import load_dataset
 from dotenv import load_dotenv, find_dotenv
-from document_model import Listing
-from hugging_face_ import get_listings
-from mongodb_ import get_db_collection, create_vector_search_index
-from openai_ import handle_user_prompt
+from utils import process_records, connect_to_database, setup_vector_search_index
+
+from pydantic import BaseModel
+from typing import Optional
+
+from IPython.display import display, HTML
 
 lock = threading.Lock()
 
@@ -20,6 +24,108 @@ RAG_ADVANCED = "Advanced RAG"
 logging.basicConfig(stream = sys.stdout, level = logging.INFO)
 logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))
 
+def vector_search(user_query, db, collection, additional_stages=[], vector_index="vector_index_text"):
+    """
+    Perform a vector search in the MongoDB collection based on the user query.
+
+    Args:
+    user_query (str): The user's query string.
+    db (MongoClient.database): The database object.
+    collection (MongoCollection): The MongoDB collection to search.
+    additional_stages (list): Additional aggregation stages to include in the pipeline.
+
+    Returns:
+    list: A list of matching documents.
+    """
+
+    # Generate embedding for the user query
+    query_embedding = custom_utils.get_embedding(user_query)
+
+    if query_embedding is None:
+        return "Invalid query or embedding generation failed."
+
+    # Define the vector search stage
+    vector_search_stage = {
+        "$vectorSearch": {
+            "index": vector_index, # specifies the index to use for the search
+            "queryVector": query_embedding, # the vector representing the query
+            "path": "text_embeddings", # field in the documents containing the vectors to search against
+            "numCandidates": 150, # number of candidate matches to consider
+            "limit": 20, # return top 20 matches
+        }
+    }
+
+    # Define the aggregate pipeline with the vector search stage and additional stages
+    pipeline = [vector_search_stage] + additional_stages
+
+    # Execute the search
+    results = collection.aggregate(pipeline)
+
+    explain_query_execution = db.command( # sends a database command directly to the MongoDB server
+        'explain', { # return information about how MongoDB executes a query or command without actually running it
+            'aggregate': collection.name, # specifies the name of the collection on which the aggregation is performed
+            'pipeline': pipeline, # the aggregation pipeline to analyze
+            'cursor': {} # indicates that default cursor behavior should be used
+        }, 
+        verbosity='executionStats') # detailed statistics about the execution of each stage of the aggregation pipeline
+
+    vector_search_explain = explain_query_execution['stages'][0]['$vectorSearch']
+    millis_elapsed = vector_search_explain['explain']['collectStats']['millisElapsed']
+
+    print(f"Total time for the execution to complete on the database server: {millis_elapsed} milliseconds")
+
+    return list(results)
+
+class SearchResultItem(BaseModel):
+    name: str
+    accommodates: Optional[int] = None
+    bedrooms: Optional[int] = None
+    address: custom_utils.Address
+    space: str = None
+
+def handle_user_query(query, db, collection, stages=[], vector_index="vector_index_text"):
+    # Assuming vector_search returns a list of dictionaries with keys 'title' and 'plot'
+    get_knowledge = vector_search(query, db, collection, stages, vector_index)
+
+    # Check if there are any results
+    if not get_knowledge:
+        return "No results found.", "No source information available."
+
+    # Convert search results into a list of SearchResultItem models
+    search_results_models = [
+        SearchResultItem(**result)
+        for result in get_knowledge
+    ]
+
+    # Convert search results into a DataFrame for better rendering in Jupyter
+    search_results_df = pd.DataFrame([item.dict() for item in search_results_models])
+
+    # Generate system response using OpenAI's completion
+    completion = custom_utils.openai.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {
+                "role": "system", 
+                "content": "You are a airbnb listing recommendation system."},
+            {
+                "role": "user", 
+                "content": f"Answer this user query: {query} with the following context:\n{search_results_df}"
+            }
+        ]
+    )
+
+    system_response = completion.choices[0].message.content
+
+    # Print User Question, System Response, and Source Information
+    print(f"- User Question:\n{query}\n")
+    print(f"- System Response:\n{system_response}\n")
+
+    # Display the DataFrame as an HTML table
+    display(HTML(search_results_df.to_html()))
+
+    # Return structured response and source info as a string
+    return system_response
+    
 def invoke(openai_api_key, prompt, rag_option):
     if not openai_api_key:
         raise gr.Error("OpenAI API Key is required.")
@@ -37,41 +143,41 @@ def invoke(openai_api_key, prompt, rag_option):
         and not too far from resturants, can you recommend a place? 
         Include a reason as to why you've chosen your selection.
         """
-        listings = get_listings()
-        db, collection = get_db_collection(listings)
-        create_vector_search_index(collection)
-        result = handle_user_prompt(openai_api_key, prompt, db, collection)
-        ###
+        dataset = load_dataset("MongoDB/airbnb_embeddings", streaming=True, split="train")
+        dataset = dataset.take(100)
+        # Convert the dataset to a pandas dataframe
+        dataset_df = pd.DataFrame(dataset)
+        dataset_df.head(5)
+        print("Columns:", dataset_df.columns)
+        
+        listings = process_records(dataset_df)
+        
+        db, collection = connect_to_database()
+        collection.delete_many({})
+        collection.insert_many(listings)
+        print("Data ingestion into MongoDB completed")
 
-        #del os.environ["OPENAI_API_KEY"]
+        setup_vector_search_index(collection=collection)
+
+        search_path = "address.country"
+
+        # Create a match stage
+        match_stage = {
+            "$match": {
+               search_path: re.compile(r"United States"),
+               "accommodates": { "$gt": 1, "$lt": 5}
+            }
+        }
         
+        additional_stages = [match_stage]
+
+        query = """
+        I want to stay in a place that's warm and friendly, 
+        and not too far from resturants, can you recommend a place? 
+        Include a reason as to why you've chosen your selection"
         """
-        if (RAG_INGESTION):
-            if (rag_option == RAG_LANGCHAIN):
-                #rag = LangChainRAG()
-                #rag.ingestion(config)
-            elif (rag_option == RAG_LLAMAINDEX):
-                #rag = LlamaIndexRAG()
-                #rag.ingestion(config)
-    
-        try:
-                #rag = LangChainRAG()
-                #completion, callback = rag.rag_chain(config, prompt)
-                #result = completion["result"]
-            elif (rag_option == RAG_LLAMAINDEX):
-                #rag = LlamaIndexRAG()
-                #result, callback = rag.retrieval(config, prompt)
-            else:
-                #rag = LangChainRAG()
-                #completion, callback = rag.llm_chain(config, prompt)
-                #result = completion.generations[0][0].text
-        except Exception as e:
-            err_msg = e
-    
-            raise gr.Error(e)
-        finally:
-            del os.environ["OPENAI_API_KEY"]
-        """
+        result = handle_user_query(query, db, collection, additional_stages)
+        ###
 
         return result