Test

app.py

@@ -6,6 +6,7 @@ import datasets
 from mini_agents import master_agent
 from utils import get_full_file_path
 from smolagents.memory import ActionStep, PlanningStep, TaskStep, SystemPromptStep, FinalAnswerStep
+from typing import Optional
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -72,6 +73,60 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer, df_agent_steps
 
+def check_required_env_vars() -> tuple[bool, Optional[str]]:
+    """Check if required environment variables are set"""
+    missing_vars = []
+    
+    # Check HF_TOKEN
+    if not os.getenv("HF_TOKEN"):
+        missing_vars.append("HF_TOKEN")
+    
+    # Check SPACE_ID (only warn, not required)
+    if not os.getenv("SPACE_ID"):
+        print("⚠️  SPACE_ID not set - this is normal when running locally")
+    
+    if missing_vars:
+        return False, f"Missing required environment variables: {', '.join(missing_vars)}"
+    return True, None
+
+def save_dataset_to_hub(df: pd.DataFrame, dataset_name: str) -> tuple[bool, str]:
+    """Save DataFrame to Hugging Face dataset with proper error handling"""
+    # Check environment variables
+    env_ok, env_error = check_required_env_vars()
+    if not env_ok:
+        return False, f"Cannot save dataset: {env_error}"
+    
+    try:
+        if len(df) == 0:
+            return False, "Cannot save empty dataset"
+            
+        print(f"Saving {len(df)} steps to {dataset_name}...")
+        
+        # Convert to dataset
+        dataset = datasets.Dataset.from_pandas(df)
+        
+        # Add metadata
+        dataset.info.description = "Agent steps data from evaluation run"
+        dataset.info.features = {
+            'task_id': datasets.Value('string'),
+            'step_class': datasets.Value('string'),
+            # Add other feature definitions as needed
+        }
+        
+        # Save to hub with token
+        dataset.push_to_hub(
+            dataset_name,
+            private=True,
+            token=os.getenv("HF_TOKEN")
+        )
+        
+        return True, f"Successfully saved {len(df)} steps to {dataset_name}"
+        
+    except Exception as e:
+        error_msg = f"Error saving dataset: {str(e)}"
+        print(error_msg)
+        return False, error_msg
+
 def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool = False):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
@@ -156,10 +211,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool =
     print(status_update)
 
     # 5. Save steps data to huggingface dataset
-    print("Commiting steps data to huggingface dataset...")
-    dataset = datasets.Dataset.from_pandas(df_agent_steps)
-    dataset.push_to_hub("huytofu92/agent_steps_huggingface_course_unit4")
-    print("Agent steps data committed to huggingface dataset.")
+    print("\nSaving agent steps to Hugging Face dataset...")
+    success, message = save_dataset_to_hub(df_agent_steps, "huytofu92/agent_steps_huggingface_course_unit4")
+    if success:
+        print(message)
+    else:
+        print(f"⚠️  {message}")
+        print("Continuing with submission despite dataset save failure...")
+    
     # 6. Submit
     print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
     if mock_submission:
@@ -180,7 +239,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool =
             )
             print("Submission successful.")
             results_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])
-            print(results_df["Task ID", "Submitted Answer"].head(20))
+            print(results_df[["Task ID", "Submitted Answer"]].head(20))
             return final_status, results_df
         except requests.exceptions.HTTPError as e:
             error_detail = f"Server responded with status {e.response.status_code}."
@@ -192,25 +251,25 @@ def run_and_submit_all( profile: gr.OAuthProfile | None, mock_submission: bool =
             status_message = f"Submission Failed: {error_detail}"
             print(status_message)
             results_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])
-            print(results_df["Task ID", "Submitted Answer"].head(20))
+            print(results_df[["Task ID", "Submitted Answer"]].head(20))
             return status_message, results_df
         except requests.exceptions.Timeout:
             status_message = "Submission Failed: The request timed out."
             print(status_message)
             results_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])
-            print(results_df["Task ID", "Submitted Answer"].head(20))
+            print(results_df[["Task ID", "Submitted Answer"]].head(20))
             return status_message, results_df
         except requests.exceptions.RequestException as e:
             status_message = f"Submission Failed: Network error - {e}"
             print(status_message)
             results_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])
-            print(results_df["Task ID", "Submitted Answer"].head(20))
+            print(results_df[["Task ID", "Submitted Answer"]].head(20))
             return status_message, results_df
         except Exception as e:
             status_message = f"An unexpected error occurred during submission: {e}"
             print(status_message)
             results_df = pd.DataFrame(results_log, columns=["Task ID", "Question", "Submitted Answer"])
-            print(results_df["Task ID", "Submitted Answer"].head(20))
+            print(results_df[["Task ID", "Submitted Answer"]].head(20))
             return status_message, results_df