Update server.py

server.py

@@ -17,7 +17,7 @@ except ImportError:
     HF_DATASETS_AVAILABLE = False
     Features, Value = None, None
 
-STORAGE_BACKEND_CONFIG = os.getenv("STORAGE_BACKEND", "JSON").upper()
+STORAGE_BACKEND_CONFIG = os.getenv("STORAGE_BACKEND", "HF_DATASET").upper()
 HF_DATASET_REPO = os.getenv("HF_DATASET_REPO")
 HF_TOKEN = os.getenv("HF_TOKEN")
 DB_FILE_JSON = "social_data.json"
@@ -74,10 +74,15 @@ def handle_persistence_after_change():
 def load_data():
     global STORAGE_BACKEND_CONFIG
     storage_backend = STORAGE_BACKEND_CONFIG
+    
+    # Define explicit dtypes for empty dataframes to ensure schema consistency
+    posts_schema = {"post_id": "int64", "username": "object", "content": "object", "timestamp": "object"}
+    comments_schema = {"comment_id": "int64", "post_id": "int64", "username": "object", "content": "object", "timestamp": "object", "reply_to_comment_id": "float64"} # Use float for optional int
+
     with db_lock:
         users = {"admin": "password"}
-        posts = pd.DataFrame(columns=["post_id", "username", "content", "timestamp"])
-        comments = pd.DataFrame(columns=["comment_id", "post_id", "username", "content", "timestamp", "reply_to_comment_id"])
+        posts = pd.DataFrame({k: pd.Series(dtype=v) for k, v in posts_schema.items()})
+        comments = pd.DataFrame({k: pd.Series(dtype=v) for k, v in comments_schema.items()})
 
         if storage_backend == "SQLITE":
             try:
@@ -105,12 +110,11 @@ def load_data():
                 try:
                     ds_dict = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
                     if ds_dict and all(k in ds_dict for k in ['users', 'posts', 'comments']):
-                        users = dict(zip(ds_dict['users']['username'], ds_dict['users']['password']))
+                        users = dict(zip(ds_dict['users']['username'], ds_dict['users']['password'])) if ds_dict['users'].num_rows > 0 else {"admin":"password"}
                         posts = ds_dict['posts'].to_pandas()
                         comments = ds_dict['comments'].to_pandas()
                         print("Successfully loaded data from HF Dataset.")
-                    else:
-                        raise ValueError("Dataset dictionary is empty or malformed.")
+                    else: raise ValueError("Dataset dictionary is empty or malformed.")
                 except Exception as e:
                     print(f"Could not load from HF Dataset '{HF_DATASET_REPO}'. Attempting to initialize. Error: {e}")
                     try:
@@ -135,7 +139,7 @@ def load_data():
                  STORAGE_BACKEND_CONFIG = "RAM"
 
     if "reply_to_comment_id" not in comments.columns:
-        comments["reply_to_comment_id"] = None
+        comments["reply_to_comment_id"] = pd.Series(dtype='float64')
         
     post_counter = int(posts['post_id'].max()) if not posts.empty else 0
     comment_counter = int(comments['comment_id'].max()) if not comments.empty else 0