Update custom_utils.py

custom_utils.py

@@ -1,8 +1,6 @@
 import openai, os, time
-#import pandas as pd
 
 from datasets import load_dataset
-#from pydantic import ValidationError
 from pymongo.collection import Collection
 from pymongo.errors import OperationFailure
 from pymongo.mongo_client import MongoClient
@@ -25,13 +23,35 @@ def rag_ingestion(collection):
     collection.insert_many(dataset)
     return "Manually create a vector search index (in free tier, this feature is not available via SDK)"
 
-def rag_retrieval(openai_api_key, 
-                  prompt, 
-                  accomodates, 
-                  bedrooms, 
-                  db, 
-                  collection, 
-                  vector_index="vector_index"):
+def rag_retrieval_naive(openai_api_key, 
+                        prompt, 
+                        db, 
+                        collection, 
+                        vector_index="vector_index"):
+    
+    get_knowledge = vector_search_naive(
+        openai_api_key, 
+        prompt, 
+        db, 
+        collection, 
+        vector_index)
+
+    if not get_knowledge:
+        return "No results found.", "No source information available."
+
+    print("###")
+    print(get_knowledge)
+    print("###")
+    
+    return get_knowledge
+
+def rag_retrieval_advanced(openai_api_key, 
+                           prompt, 
+                           accomodates, 
+                           bedrooms, 
+                           db, 
+                           collection, 
+                           vector_index="vector_index"):
     ###
     ### Pre-retrieval processing: index filter
     ### Post-retrieval processing: result filter
@@ -115,7 +135,7 @@ def rag_retrieval(openai_api_key,
     ###    
     ###
     
-    get_knowledge = vector_search(
+    get_knowledge = vector_search_advanced(
         openai_api_key, 
         prompt, 
         accomodates, 
@@ -156,14 +176,77 @@ def rag_inference(openai_api_key,
 
     return completion.choices[0].message.content
 
-def vector_search(openai_api_key, 
-                  user_query, 
-                  accommodates, 
-                  bedrooms, 
-                  db, 
-                  collection, 
-                  additional_stages=[], 
-                  vector_index="vector_index"):
+def inference(openai_api_key, 
+              prompt):
+    openai.api_key = openai_api_key
+  
+    content = f"Answer this user question: {prompt}"
+    
+    completion = openai.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "system", 
+                "content": "You are an AirBnB listing recommendation system."},
+            {
+                "role": "user", 
+                "content": content
+            }
+        ]
+    )
+
+    return completion.choices[0].message.content
+
+def vector_search_naive(openai_api_key, 
+                        user_query, 
+                        db, 
+                        collection, 
+                        vector_index="vector_index"):
+    query_embedding = get_text_embedding(openai_api_key, user_query)
+
+    if query_embedding is None:
+        return "Invalid query or embedding generation failed."
+
+    vector_search_stage = {
+        "$vectorSearch": {
+            "index": vector_index,
+            "queryVector": query_embedding,
+            "path": "description_embedding",
+            "numCandidates": 150,
+            "limit": 25,
+        }
+    }
+
+    remove_embedding_stage =     {
+        "$unset": "description_embedding"
+    }
+    
+    pipeline = [vector_search_stage, remove_embedding_stage]
+
+    results = collection.aggregate(pipeline)
+
+    #explain_query_execution = db.command(
+    #    "explain", {
+    #        "aggregate": collection.name,
+    #        "pipeline": pipeline,
+    #        "cursor": {}
+    #    }, 
+    #    verbosity='executionStats')
+
+    #vector_search_explain = explain_query_execution["stages"][0]["$vectorSearch"]
+    #millis_elapsed = vector_search_explain["explain"]["collectStats"]["millisElapsed"]
+    #print(f"Query execution time: {millis_elapsed} milliseconds")
+
+    return list(results)
+
+def vector_search_advanced(openai_api_key, 
+                           user_query, 
+                           accommodates, 
+                           bedrooms, 
+                           db, 
+                           collection, 
+                           additional_stages=[], 
+                           vector_index="vector_index"):
     query_embedding = get_text_embedding(openai_api_key, user_query)
 
     if query_embedding is None: