Update custom_utils.py

custom_utils.py

@@ -55,53 +55,9 @@ def rag_retrieval_advanced(openai_api_key,
     
     # 1b) Post-retrieval processing: result filter (accomodates, bedrooms) plus...
     
-    #match_stage = {
-    #    "$match": {
-    #        "accommodates": { "$eq": 2},
-    #        "bedrooms": { "$eq": 1}
-    #    }
-    #}
+    # 2) Weighted average review, sorted in descending order
 
-    # 2) Average review score and review count boost, sorted in descending order
-
-    review_average_stage = {
-        "$addFields": {
-            "averageReviewScore": {
-                "$divide": [
-                    {
-                        "$add": [
-                            "$review_scores_rating",
-                            "$review_scores_accuracy",
-                            "$review_scores_cleanliness",
-                            "$review_scores_checkin",
-                            "$review_scores_communication",
-                            "$review_scores_location",
-                            "$review_scores_value",
-                        ]
-                    },
-                    7
-                ]
-            },
-            "reviewCountBoost": "$number_of_reviews"
-        }
-    }
-
-    weighting_stage = {
-        "$addFields": {
-            "combinedScore": {
-                "$add": [
-                    {"$multiply": ["$averageReviewScore", 0.9]},
-                    {"$multiply": ["$reviewCountBoost", 0.1]},
-                ]
-            }
-        }
-    }
-
-    sorting_stage_sort = {
-        "$sort": {"combinedScore": -1}
-    }
-
-    additional_stages = [review_average_stage, weighting_stage, sorting_stage_sort]
+    additional_stages = [get_average_review_stage(), get_weighting_stage(), get_sorting_stage()]
     
     retrieval_result = vector_search_advanced(
         openai_api_key, 
@@ -201,7 +157,7 @@ def vector_search_advanced(openai_api_key,
     if query_embedding is None:
         return "Invalid query or embedding generation failed."
 
-    vector_search_stage = {
+    vector_search_filter_stage = {
         "$vectorSearch": {
             "index": vector_index,
             "queryVector": query_embedding,
@@ -217,7 +173,7 @@ def vector_search_advanced(openai_api_key,
         }
     }
   
-    pipeline = [vector_search_stage, get_remove_embedding_stage()] + additional_stages
+    pipeline = [vector_search_filter_stage, get_remove_embedding_stage()] + additional_stages
 
     return invoke_search(db, collection, pipeline)
 
@@ -226,6 +182,54 @@ def get_remove_embedding_stage():
         "$unset": "description_embedding"
     }
 
+def get_result_filter_stage():
+    return {
+        "$match": {
+            "accommodates": { "$eq": 2},
+            "bedrooms": { "$eq": 1}
+        }
+    }
+
+def get_average_review_stage():
+    return {
+        "$addFields": {
+            "averageReview": {
+                "$divide": [
+                    {
+                        "$add": [
+                            "$review_scores_rating",
+                            "$review_scores_accuracy",
+                            "$review_scores_cleanliness",
+                            "$review_scores_checkin",
+                            "$review_scores_communication",
+                            "$review_scores_location",
+                            "$review_scores_value",
+                        ]
+                    },
+                    7
+                ]
+            },
+            "reviewCount": "$number_of_reviews"
+        }
+    }
+
+def get_weighting_stage():
+    return {
+        "$addFields": {
+            "weightedAverageReview": {
+                "$add": [
+                    {"$multiply": ["$averageReview", 0.9]},
+                    {"$multiply": ["$reviewCount", 0.1]},
+                ]
+            }
+        }
+    }
+
+def get_sorting_stage():
+    return {
+        "$sort": {"weightedAverageReview": -1}
+    }
+
 def invoke_search(db, collection, pipeline):
     results = collection.aggregate(pipeline)