Update src/mongo_search.py

src/mongo_search.py

@@ -1,6 +1,7 @@
 from utils.mongo_utils import generate_mongodb_query, get_prompt
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_core.output_parsers import JsonOutputParser
+import random
 
 from utils.utils import timing_decorator
 
@@ -16,7 +17,9 @@ class MongoSearch:
         self.index_variable = index_variable
 
     @timing_decorator
-    def __call__(self, query, k=3, use_filter=True):
+    def __call__(self, query, k=4, use_filter=True):
+
+
 
         query_filter = {}
         if use_filter:
@@ -52,6 +55,8 @@ class MongoSearch:
         # Extract unique makeModel values for the next step
         make_model_list = [doc['makeModel'] for doc in first_search_results]
 
+        k = k * 3
+
         # Define the second pipeline
         second_pipeline = [
             {
@@ -79,6 +84,16 @@ class MongoSearch:
         # run pipeline
         result = self.collection.aggregate(second_pipeline)
 
+        # Convert the result cursor to a list
+        result_list = list(result)
+
+        # Randomly select k/3 objects from the results
+        k_third = k // 3
+        selected_cars = random.sample(result_list, min(k_third, len(result_list)))
+
+        # Update the result with the randomly selected cars
+        result = selected_cars
+
         cars = []
         for i in result:
             cars.append(i)