Update custom_utils.py

custom_utils.py

@@ -6,8 +6,8 @@ from document_model import Listing, SearchResultItem
 from pydantic import ValidationError
 from pymongo.collection import Collection
 from pymongo.errors import OperationFailure
-from pymongo.operations import SearchIndexModel
 from pymongo.mongo_client import MongoClient
+from pymongo.operations import SearchIndexModel
 
 DB_NAME = "airbnb_dataset"
 COLLECTION_NAME = "listings_reviews"
@@ -28,6 +28,28 @@ def rag_ingestion(collection):
     collection.insert_many(listings)
     return "Manually create a vector search index (in free tier, this feature is not available via SDK)"
 
+def process_records(data_frame):
+    records = data_frame.to_dict(orient="records")
+
+    # Handle potential NaT values
+    for record in records:
+        for key, value in record.items():
+            # List values
+            if isinstance(value, list):
+                processed_list = [None if pd.isnull(v) else v for v in value]
+                record[key] = processed_list
+            # Scalar values
+            else:
+                if pd.isnull(value):
+                    record[key] = None
+
+    try:
+        # Convert each dictionary to a Listing instance
+        return [Listing(**record).dict() for record in records]
+    except ValidationError as e:
+        print("Validation error:", e)
+        return []
+
 def rag_retrieval(openai_api_key, prompt, db, collection, stages=[], vector_index="vector_index"):
     get_knowledge = vector_search(openai_api_key, prompt, db, collection, stages, vector_index)
 
@@ -67,30 +89,6 @@ def rag_inference(openai_api_key, prompt, search_results):
 
     return completion.choices[0].message.content
 
-def process_records(data_frame):
-    records = data_frame.to_dict(orient="records")
-
-    # Handle potential NaT values
-    for record in records:
-        print("###")
-        print(record)
-        #for key, value in record.items():
-        #    # List values
-        #    if isinstance(value, list):
-        #        processed_list = [None if pd.isnull(v) else v for v in value]
-        #        record[key] = processed_list
-        #    # Scalar values
-        #    else:
-        #        if pd.isnull(value):
-        #            record[key] = None
-
-    try:
-        # Convert each dictionary to a Listing instance
-        return [Listing(**record).dict() for record in records]
-    except ValidationError as e:
-        print("Validation error:", e)
-        return []
-
 def vector_search(openai_api_key, user_query, db, collection, additional_stages=[], vector_index="vector_index"):
     query_embedding = get_text_embedding(openai_api_key, user_query)