Update main.py

main.py

@@ -1,117 +1,124 @@
+# Import necessary libraries (ensure all required imports are at the top)
 import os
 import pandas as pd
 from fastapi import FastAPI, HTTPException, Body
 from pydantic import BaseModel, Field
-from typing import List, Dict, Any
+from typing import List, Dict, Any #<-- Make sure Any is imported
 from datasets import load_dataset, Dataset, DatasetDict
 from huggingface_hub import HfApi, hf_hub_download
 from datetime import datetime, timezone
 import logging
-import uvicorn # To run the app
-import random # <-- Added import for random choice
+import uvicorn
+import random
 
+# --- Constants and Config ---
 tool_threshold = 3
 step_threshold = 5
-
-# --- Configuration ---
 HF_DATASET_ID = "agents-course/unit4-students-scores"
-# Ensure you have write access to this dataset repository on Hugging Face
-# and are logged in via `huggingface-cli login` or have HF_TOKEN env var set.
-    # Prepare data structures for the API
-questions_for_api: List[Dict[str, str]] = []
+
+# --- Data Structures ---
+# questions_for_api will now hold richer dictionaries
+questions_for_api: List[Dict[str, Any]] = []
 ground_truth_answers: Dict[str, str] = {}
+
 # --- Logging Setup ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-filtered_dataset=None
+
+# --- Filtered Dataset Placeholder ---
+# Note: Making filtered_dataset global might not be ideal in larger apps,
+# but keeping it as is based on the original code.
+filtered_dataset = None
+
+# --- Modified load_questions Function ---
 def load_questions():
     global filtered_dataset
     global questions_for_api
     global ground_truth_answers
     tempo_filtered=[]
-    # Clear existing data to prevent duplication if called multiple times
+    # Clear existing data
     questions_for_api.clear()
     ground_truth_answers.clear()
 
     logger.info("Starting to load and filter GAIA dataset...")
     try:
-        dataset=load_dataset("gaia-benchmark/GAIA","2023_level1",trust_remote_code=True)
-        logger.info("GAIA dataset loaded.")
+        # Load the 'validation' split specifically if that's intended
+        dataset = load_dataset("gaia-benchmark/GAIA", "2023_level1", split='validation', trust_remote_code=True)
+        logger.info("GAIA dataset validation split loaded.")
     except Exception as e:
-        logger.error(f"Failed to load GAIA dataset: {e}", exc_info=True)
-        # Decide how to handle this: maybe raise the error or exit
+        logger.error(f"Failed to load GAIA dataset validation split: {e}", exc_info=True)
         raise RuntimeError("Could not load the primary GAIA dataset.") from e
 
-    for question in dataset['validation']:
-        metadata = question.get('Annotator Metadata') # Use .get() for safety
+    # --- Filtering Logic (remains the same) ---
+    for question in dataset: # Iterate directly over the loaded split
+        metadata = question.get('Annotator Metadata')
 
-        if metadata: # Check if 'Annotator Metadata' exists
+        if metadata:
             num_tools_str = metadata.get('Number of tools')
             num_steps_str = metadata.get('Number of steps')
 
-            # Check if both numbers exist before trying to convert
             if num_tools_str is not None and num_steps_str is not None:
                 try:
-                    # Convert values to integers for comparison
                     num_tools = int(num_tools_str)
                     num_steps = int(num_steps_str)
 
-                    # Apply the filter conditions
                     if num_tools < tool_threshold and num_steps < step_threshold:
-                        # logger.debug(f"MATCH FOUND (Task ID: {question.get('task_id', 'N/A')}) - Tools: {num_tools}, Steps: {num_steps}")
-                        # logger.debug(question) # Print the matching question dictionary
-                        # logger.debug("------------------------------------------------------------------")
-                        tempo_filtered.append(question) # Add to the filtered list
-                    # else: # Optional: Handle items that don't match the filter
-                        # logger.debug(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Tools: {num_tools}, Steps: {num_steps}")
+                        tempo_filtered.append(question)
                 except ValueError:
-                    # Handle cases where 'Number of tools' or 'Number of steps' is not a valid integer
-                    logger.warning(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Could not convert tool/step count to integer: tools='{num_tools_str}', steps='{num_steps_str}'.")
-                    # logger.debug("------------------------------------------------------------------")
+                     logger.warning(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Could not convert tool/step count: tools='{num_tools_str}', steps='{num_steps_str}'.")
+            # else: # Optional: Log if numbers are missing
+                # logger.debug(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Missing tool/step count in metadata.")
         else:
              logger.warning(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Missing 'Annotator Metadata'.")
-             # logger.debug("------------------------------------------------------------------")
 
-    filtered_dataset=tempo_filtered
+    # Store the filtered list (optional, could process directly)
+    filtered_dataset = tempo_filtered
     logger.info(f"Found {len(filtered_dataset)} questions matching the criteria (tools < {tool_threshold}, steps < {step_threshold}).")
-    # print(filtered_dataset) # Keep this commented unless debugging
 
+    # --- Processing Logic (Modified) ---
     processed_count = 0
     for item in filtered_dataset:
         task_id = item.get('task_id')
-        question_text = item.get('Question')
+        question_text = item.get('Question') # Keep original key for now
         final_answer = item.get('Final answer')
-        
-        # Validate required fields
+
+        # Validate required fields needed for processing/scoring
         if task_id and question_text and final_answer is not None:
-            # Create a copy of the item and remove fields we don't want
-            processed_item = item.copy()
-            processed_item.pop('Final answer', None)  # Remove Final answer
-            processed_item.pop('Annotator Annotation', None)  # Remove Annotator Annotation
-            
-            # Ensure the field name matches what's expected by Pydantic models
-            if 'Question' in processed_item and 'question' not in processed_item:
-                processed_item['question'] = processed_item.pop('Question')
-            
-            # Store in questions_for_api
+            # Create a copy to avoid modifying the original item in filtered_dataset
+            processed_item: Dict[str, Any] = item.copy()
+
+            # Remove the fields we explicitly want to exclude
+            processed_item.pop('Final answer', None)
+            processed_item.pop('Annotator Annotation', None)
+            # You could add more fields to pop here if needed later
+            # processed_item.pop('Another field to remove', None)
+
+            # Store the dictionary containing all remaining fields
             questions_for_api.append(processed_item)
-            
-            # Still store the ground truth answers separately
+
+            # Store the ground truth answer separately for scoring
             ground_truth_answers[str(task_id)] = str(final_answer)
             processed_count += 1
         else:
-            logger.warning(f"Skipping item due to missing fields (task_id, Question, or Final answer): {item}")
+            # Log which required field was missing if possible
+            missing = [k for k, v in {'task_id': task_id, 'Question': question_text, 'Final answer': final_answer}.items() if not v and v is not None]
+            logger.warning(f"Skipping item due to missing required fields ({', '.join(missing)}): task_id={task_id}")
+
+    logger.info(f"Successfully processed {processed_count} questions into API format.")
 
     if not questions_for_api:
          logger.error("CRITICAL: No valid questions loaded after filtering. API endpoints needing questions will fail.")
-         # Depending on requirements, you might want to exit or raise an error here
          # raise RuntimeError("Failed to load mandatory question data after filtering.")
 
-# --- Pydantic Models for Data Validation ---
+# --- Pydantic Models ---
+# Keep Question simple for potential internal use or basic validation,
+# but the API will return Dict[str, Any]
 class Question(BaseModel):
     task_id: str
-    question: str
+    Question: str # Keep original casing if that's what in the data
 
+# Keep other models as they are (AnswerItem, Submission, ScoreResponse, ErrorResponse)
+# ... (rest of the Pydantic models remain the same) ...
 class AnswerItem(BaseModel):
     task_id: str
     submitted_answer: str = Field(..., description="The agent's answer for the task_id")
@@ -132,35 +139,30 @@ class ScoreResponse(BaseModel):
 class ErrorResponse(BaseModel):
     detail: str
 
+
 # --- FastAPI Application ---
 app = FastAPI(
     title="Agent Evaluation API",
     description="API to fetch questions and submit agent answers for scoring.",
 )
 
-# --- Startup Event Handler ---
+# --- Startup Event ---
 @app.on_event("startup")
 async def startup_event():
-    """
-    Loads the questions when the FastAPI application starts.
-    """
     logger.info("Application startup: Loading questions...")
     try:
-        load_questions() # Call your loading function here
+        load_questions()
         if not questions_for_api:
-            logger.error("CRITICAL: No questions were loaded during startup. The /questions and /random-question endpoints might fail.")
-            # Depending on requirements, you might want the app to fail startup
-            # raise RuntimeError("Failed to load mandatory question data.")
+            logger.error("CRITICAL: No questions were loaded during startup.")
         else:
             logger.info(f"Successfully loaded {len(questions_for_api)} questions.")
     except Exception as e:
         logger.error(f"CRITICAL ERROR DURING STARTUP while loading questions: {e}", exc_info=True)
-        # Decide if the app should exit if loading fails
         # import sys
-        # sys.exit(1)
+        # sys.exit(1) # Consider exiting if questions are critical
 
-
-# --- Helper Function to interact with HF Dataset ---
+# --- Helper Function (update_huggingface_dataset remains the same) ---
+# ... (update_huggingface_dataset function code) ...
 def update_huggingface_dataset(username: str, score: float):
     """Loads the dataset, updates the score if higher, and pushes back."""
     try:
@@ -242,7 +244,7 @@ def update_huggingface_dataset(username: str, score: float):
 
             updated_ds = DatasetDict({'train': Dataset.from_pandas(df)})
             logger.info(f"Dataset to push: {updated_ds}") # Log the dataset structure
-            updated_ds.push_to_hub(HF_DATASET_ID) # Token should be picked up from env or login
+            # updated_ds.push_to_hub(HF_DATASET_ID) # Uncomment this line to enable leaderboard updates
             logger.warning("Dataset push to hub is currently commented out. Uncomment the line above to enable leaderboard updates.") # REMINDER
             logger.info("Dataset push simulated/attempted.")
             return True
@@ -254,35 +256,35 @@ def update_huggingface_dataset(username: str, score: float):
         # Re-raise the exception to be caught by the endpoint handler
         raise HTTPException(status_code=500, detail=f"Failed to update Hugging Face dataset: {e}")
 
-
-# --- API Endpoints ---
+# --- API Endpoints (Modified response_model) ---
 
 @app.get("/questions",
-         response_model=List[Question],
-         summary="Get All Filtered Questions",
-         description="Returns the complete list of questions (task_id and question text only) filtered based on criteria.")
+         # Return a list of dictionaries with arbitrary keys/values
+         response_model=List[Dict[str, Any]],
+         summary="Get All Filtered Questions (Full Data)",
+         description="Returns the complete list of questions with all associated data (excluding answer/annotation) filtered based on criteria.")
 async def get_questions():
     """
-    Provides the list of questions that agents should answer.
+    Provides the list of questions (with extended data) that agents should answer.
     """
-    # print(f"Returning {len(questions_for_api)} questions.") # Debug log
     if not questions_for_api:
          logger.error("GET /questions requested but no questions are loaded.")
          raise HTTPException(status_code=404, detail="No questions available.")
+    # questions_for_api now contains the richer dictionaries
     return questions_for_api
 
-# --- NEW ENDPOINT ---
 @app.get("/random-question",
-         response_model=Question,
-         summary="Get One Random Question",
-         description="Returns a single random question from the available filtered set.",
+         # Return a single dictionary with arbitrary keys/values
+         response_model=Dict[str, Any],
+         summary="Get One Random Question (Full Data)",
+         description="Returns a single random question with all associated data (excluding answer/annotation) from the available filtered set.",
          responses={
-             200: {"description": "A random question."},
+             200: {"description": "A random question with its full data."},
              404: {"model": ErrorResponse, "description": "No questions available to choose from."}
          })
 async def get_random_question():
     """
-    Provides a single, randomly selected question from the loaded list.
+    Provides a single, randomly selected question with its extended data.
     """
     if not questions_for_api:
         logger.warning("GET /random-question requested but no questions are loaded.")
@@ -290,11 +292,11 @@ async def get_random_question():
 
     # Select and return a random question dictionary
     random_question = random.choice(questions_for_api)
-    logger.info(f"Returning random question with task_id: {random_question['task_id']}")
+    logger.info(f"Returning random question with task_id: {random_question.get('task_id', 'N/A')}")
+    # random_question is already the richer dictionary
     return random_question
-# --- END NEW ENDPOINT ---
-
 
+# --- Submit Endpoint (remains the same, uses ground_truth_answers) ---
 @app.post("/submit",
           response_model=ScoreResponse,
           summary="Submit Agent Answers",
@@ -358,17 +360,22 @@ async def submit_answers(submission: Submission = Body(...)):
              logger.debug(f"Incorrect answer for {task_id} from {submission.username}. Submitted: '{submitted}', Expected: '{ground_truth}'")
 
 
-    # Calculate score based on valid attempts
+    # Calculate score based on valid attempts AND total number of questions available
     if valid_attempted_count == 0:
         score = 0.0
         message = f"Submission received, but no valid/matching task IDs were found in the {total_attempted_in_payload} answers provided."
         logger.warning(f"No valid answers processed for {submission.username} out of {total_attempted_in_payload} submitted.")
+    elif not ground_truth_answers: # Prevent division by zero if no questions loaded
+         score = 0.0
+         message = "Score cannot be calculated because no ground truth answers are loaded."
+         logger.error(f"Cannot calculate score for {submission.username}: ground_truth_answers is empty.")
     else:
+        # Score is based on correct answers divided by the TOTAL number of questions in the filtered set
         score = round((correct_count / len(ground_truth_answers)) * 100, 2)
-        message = f"Score calculated successfully: {correct_count}/{valid_attempted_count} correct answers for valid tasks."
+        message = f"Score calculated successfully: {correct_count}/{len(ground_truth_answers)} total questions answered correctly ({valid_attempted_count} valid tasks attempted)."
         if valid_attempted_count < total_attempted_in_payload:
              message += f" ({total_attempted_in_payload - valid_attempted_count} submitted answers had invalid or duplicate task IDs)."
-        logger.info(f"Score for {submission.username}: {score}% ({correct_count}/{valid_attempted_count})")
+        logger.info(f"Score for {submission.username}: {score}% ({correct_count}/{len(ground_truth_answers)} correct, based on {valid_attempted_count} valid attempts)")
 
 
     # Update Hugging Face dataset
@@ -401,22 +408,18 @@ async def submit_answers(submission: Submission = Body(...)):
     )
 
 # --- Run the application ---
-# This part is mainly for local development without Docker.
-# Docker uses the CMD instruction in the Dockerfile.
 if __name__ == "__main__":
     logger.info("Starting FastAPI server for local development...")
-    # Explicitly call load_questions here for local run,
-    # as the @app.on_event("startup") might not trigger reliably
-    # depending on how uvicorn is invoked directly.
     try:
-        load_questions()
+        load_questions() # Load questions before starting server
         if not questions_for_api:
              logger.error("EXITING: Cannot start server without loaded questions.")
+             # Optional: exit if questions are essential
+             # import sys
+             # sys.exit(1)
         else:
-            # Read port from environment variable for consistency, default to 8000 for local if not set
             local_port = int(os.getenv("PORT", "8000"))
             logger.info(f"Running Uvicorn locally on http://127.0.0.1:{local_port}")
-            # Note: host='127.0.0.1' is usually fine for local runs outside docker
             uvicorn.run(app, host="127.0.0.1", port=local_port, log_level="info")
     except Exception as e:
         logger.error(f"Failed to start server: {e}", exc_info=True)