Update custom_utils.py

custom_utils.py

@@ -25,7 +25,13 @@ 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, db, collection, vector_index="vector_index"):
+def rag_retrieval(openai_api_key, 
+                  prompt, 
+                  accomodates, 
+                  bedrooms, 
+                  db, 
+                  collection, 
+                  vector_index="vector_index"):
     ###
     ### Pre-retrieval processing: index filter
     ### Post-retrieval processing: result filter
@@ -108,7 +114,15 @@ def rag_retrieval(openai_api_key, prompt, db, collection, vector_index="vector_i
     ###    
     ###
     
-    get_knowledge = vector_search(openai_api_key, prompt, db, collection, additional_stages, vector_index)
+    get_knowledge = vector_search(
+        openai_api_key, 
+        prompt, 
+        accomodates, 
+        bedrooms, 
+        db, 
+        collection, 
+        additional_stages, 
+        vector_index)
 
     if not get_knowledge:
         return "No results found.", "No source information available."
@@ -119,7 +133,9 @@ def rag_retrieval(openai_api_key, prompt, db, collection, vector_index="vector_i
     
     return get_knowledge
 
-def rag_inference(openai_api_key, prompt, search_results):
+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}"
@@ -139,39 +155,44 @@ def rag_inference(openai_api_key, prompt, search_results):
 
     return completion.choices[0].message.content
 
-def vector_search(openai_api_key, user_query, db, collection, additional_stages=[], vector_index="vector_index"):
+def vector_search(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:
         return "Invalid query or embedding generation failed."
 
-    vector_search_stage = {
-        "$vectorSearch": {
-            "index": vector_index,
-            "queryVector": query_embedding,
-            "path": "description_embedding",
-            "numCandidates": 150,
-            "limit": 3,
-        }
-    }
+    #vector_search_stage = {
+    #    "$vectorSearch": {
+    #        "index": vector_index,
+    #        "queryVector": query_embedding,
+    #        "path": "description_embedding",
+    #        "numCandidates": 150,
+    #        "limit": 3,
+    #    }
+    #}
 
-    """ filter
     vector_search_stage = {
         "$vectorSearch": {
             "index": vector_index,
             "queryVector": query_embedding,
             "path": "description_embedding",
             "numCandidates": 150,
-            "limit": 3,
+            "limit": 10,
             "filter": {
                 "$and": [
-                    {"accommodates": {"$eq": 2}}, 
-                    {"bedrooms": {"$eq": 1}}
+                    {"accommodates": {"$eq": accommodates}}, 
+                    {"bedrooms": {"$eq": bedrooms}}
                 ]
             },
         }
     }
-    """
 
     remove_embedding_stage =     {
         "$unset": "description_embedding"