Update custom_utils.py

custom_utils.py

@@ -28,22 +28,23 @@ def rag_retrieval_naive(openai_api_key,
                         db, 
                         collection, 
                         vector_index="vector_index"):
-    
-    get_knowledge = vector_search_naive(
+    # Naive RAG: Semantic search
+    retrieval_result = vector_search_naive(
         openai_api_key, 
         prompt, 
         db, 
         collection, 
-        vector_index)
+        vector_index
+    )
 
-    if not get_knowledge:
-        return "No results found.", "No source information available."
+    if not retrieval_result:
+        return "No results found."
 
-    print("###")
-    print(get_knowledge)
-    print("###")
+    #print("###")
+    #print(retrieval_result)
+    #print("###")
     
-    return get_knowledge
+    return retrieval_result
 
 def rag_retrieval_advanced(openai_api_key, 
                            prompt, 
@@ -52,43 +53,23 @@ def rag_retrieval_advanced(openai_api_key,
                            db, 
                            collection, 
                            vector_index="vector_index"):
-    ###
-    ### Pre-retrieval processing: index filter
-    ### Post-retrieval processing: result filter
+    # Advanced RAG: Semantic search plus...
+    
+    # 1a) Pre-retrieval processing: index filter (accomodates, bedrooms) plus...
+    
+    # 1b) Post-retrieval processing: result filter (accomodates, bedrooms) plus...
+    
     #match_stage = {
     #    "$match": {
     #        "accommodates": { "$eq": 2},
     #        "bedrooms": { "$eq": 1}
     #    }
     #}
-
-    #additional_stages = [match_stage]
-    ###
-    """
-    projection_stage = {
-        "$project": {
-            "_id": 0,  
-            "name": 1,
-            "accommodates": 1,
-            "address.street": 1,
-            "address.government_area": 1, 
-            "address.market": 1,
-            "address.country": 1, 
-            "address.country_code": 1, 
-            "address.location.type": 1, 
-            "address.location.coordinates": 1,  
-            "address.location.is_location_exact": 1,
-            "summary": 1,
-            "space": 1,  
-            "neighborhood_overview": 1, 
-            "notes": 1, 
-            "score": {"$meta": "vectorSearchScore"} 
-        }
-    }
     
-    additional_stages = [projection_stage]
-    """
-    ###
+    #additional_stages = [match_stage]
+
+    # 2) Average review score and review count boost, sorted in descending order
+
     review_average_stage = {
         "$addFields": {
             "averageReviewScore": {
@@ -104,10 +85,9 @@ def rag_retrieval_advanced(openai_api_key,
                             "$review_scores_value",
                         ]
                     },
-                    6  # Divide by the number of review score types to get the average
+                    7
                 ]
             },
-            # Calculate a score boost factor based on the number of reviews
             "reviewCountBoost": "$number_of_reviews"
         }
     }
@@ -115,27 +95,21 @@ def rag_retrieval_advanced(openai_api_key,
     weighting_stage = {
         "$addFields": {
             "combinedScore": {
-                # Example formula that combines average review score and review count boost
                 "$add": [
-                    {"$multiply": ["$averageReviewScore", 0.9]},  # Weighted average review score
-                    {"$multiply": ["$reviewCountBoost", 0.1]}   # Weighted review count boost
+                    {"$multiply": ["$averageReviewScore", 0.9]},
+                    {"$multiply": ["$reviewCountBoost", 0.1]},
                 ]
             }
         }
     }
 
-    # Apply the combinedScore for sorting
     sorting_stage_sort = {
-        "$sort": {"combinedScore": -1}  # Descending order to boost higher combined scores
+        "$sort": {"combinedScore": -1}
     }
 
     additional_stages = [review_average_stage, weighting_stage, sorting_stage_sort]
-    ###
-    #additional_stages = []
-    ###    
-    ###
     
-    get_knowledge = vector_search_advanced(
+    retrieval_result = vector_search_advanced(
         openai_api_key, 
         prompt, 
         accomodates, 
@@ -143,45 +117,29 @@ def rag_retrieval_advanced(openai_api_key,
         db, 
         collection, 
         additional_stages, 
-        vector_index)
+        vector_index
+    )
 
-    if not get_knowledge:
-        return "No results found.", "No source information available."
+    if not retrieval_result:
+        return "No results found."
 
-    print("###")
-    print(get_knowledge)
-    print("###")
+    #print("###")
+    #print(retrieval_result)
+    #print("###")
     
-    return get_knowledge
+    return retrieval_result
 
-def rag_inference(openai_api_key, 
-                  prompt, 
-                  search_results):
-    openai.api_key = openai_api_key
-  
-    content = f"Answer this user question: {prompt} with the following context:\n{search_results}"
-    
-    completion = openai.chat.completions.create(
-        model="gpt-4o",
-        messages=[
-            {
-                "role": "system", 
-                "content": "You are an AirBnB listing recommendation system."},
-            {
-                "role": "user", 
-                "content": content
-            }
-        ]
-    )
+def inference(openai_api_key, prompt):
+    content = f"Answer this user question: {prompt}"
+    return invoke_llm(openai_api_key, content)
 
-    return completion.choices[0].message.content
+def rag_inference(openai_api_key, prompt, retrieval_result):
+    content = f"Answer this user question: {prompt} with the following context:\n{retrieval_result}"
+    return invoke_llm(openai_api_key, content)
 
-def inference(openai_api_key, 
-              prompt):
+def invoke_llm(openai_api_key, content):
     openai.api_key = openai_api_key
-  
-    content = f"Answer this user question: {prompt}"
-    
+
     completion = openai.chat.completions.create(
         model="gpt-4o",
         messages=[
@@ -196,7 +154,7 @@ def inference(openai_api_key,
     )
 
     return completion.choices[0].message.content
-
+    
 def vector_search_naive(openai_api_key, 
                         user_query, 
                         db, 
@@ -223,21 +181,7 @@ def vector_search_naive(openai_api_key,
     
     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)
+    return invoke_search(collection, pipeline)
 
 def vector_search_advanced(openai_api_key, 
                            user_query, 
@@ -252,16 +196,6 @@ def vector_search_advanced(openai_api_key,
     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,
-    #    }
-    #}
-
     vector_search_stage = {
         "$vectorSearch": {
             "index": vector_index,
@@ -284,20 +218,10 @@ def vector_search_advanced(openai_api_key,
     
     pipeline = [vector_search_stage, remove_embedding_stage] + additional_stages
 
-    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 invoke_search(collection, pipeline)
 
+def invoke_search(collection, pipeline):
+    results = collection.aggregate(pipeline)
     return list(results)
 
 def get_text_embedding(openai_api_key, text):
@@ -307,10 +231,10 @@ def get_text_embedding(openai_api_key, text):
     openai.api_key = openai_api_key
         
     try:
-        embedding = openai.embeddings.create(
+        return openai.embeddings.create(
             input=text,
-            model="text-embedding-3-small", dimensions=1536).data[0].embedding
-        return embedding
+            model="text-embedding-3-small", dimensions=1536
+        ).data[0].embedding
     except Exception as e:
         print(f"Error in get_embedding: {e}")
         return None