Update main.py

main.py

@@ -3,7 +3,7 @@ import os
 import pandas as pd
 from fastapi import FastAPI, HTTPException, Body
 from pydantic import BaseModel, Field
-from typing import List, Dict, Any #<-- Make sure Any is imported
+from typing import List, Dict, Any, Optional
 from datasets import load_dataset, Dataset, DatasetDict
 from huggingface_hub import HfApi, hf_hub_download
 from datetime import datetime, timezone
@@ -12,8 +12,6 @@ import uvicorn
 import random
 
 # --- Constants and Config ---
-tool_threshold = 3
-step_threshold = 5
 HF_DATASET_ID = "agents-course/unit4-students-scores"
 
 # --- Data Structures ---
@@ -25,33 +23,37 @@ ground_truth_answers: Dict[str, str] = {}
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- 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 ---
+logger = logging.getLogger(__name__) # Make sure logger is initialized
+tool_threshold = 3
+step_threshold = 5
+questions_for_api: List[Dict[str, Any]] = [] # Use Dict[str, Any] for flexibility before validation
+ground_truth_answers: Dict[str, str] = {}
+filtered_dataset = None # Or initialize as empty list: []
+
 def load_questions():
     global filtered_dataset
     global questions_for_api
     global ground_truth_answers
-    tempo_filtered=[]
+    tempo_filtered = []
     # Clear existing data
     questions_for_api.clear()
     ground_truth_answers.clear()
 
-    logger.info("Starting to load and filter GAIA dataset...")
+    logger.info("Starting to load and filter GAIA dataset (validation split)...")
     try:
-        # 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.")
+        # Load the specified split and features
+        dataset = load_dataset("gaia-benchmark/GAIA", "2023_level1", split="validation", trust_remote_code=True)
+        logger.info(f"GAIA dataset validation split loaded. Features: {dataset.features}")
+        # You can uncomment below to see the first item's structure if needed
+        # logger.debug(f"First item structure: {dataset[0]}")
     except Exception as e:
-        logger.error(f"Failed to load GAIA dataset validation split: {e}", exc_info=True)
+        logger.error(f"Failed to load GAIA dataset: {e}", exc_info=True)
         raise RuntimeError("Could not load the primary GAIA dataset.") from e
 
     # --- Filtering Logic (remains the same) ---
-    for question in dataset: # Iterate directly over the loaded split
-        metadata = question.get('Annotator Metadata')
+    for item in dataset:
+        metadata = item.get('Annotator Metadata')
 
         if metadata:
             num_tools_str = metadata.get('Number of tools')
@@ -63,59 +65,88 @@ def load_questions():
                     num_steps = int(num_steps_str)
 
                     if num_tools < tool_threshold and num_steps < step_threshold:
-                        tempo_filtered.append(question)
+                        tempo_filtered.append(item) # Add the original item if it matches filter
                 except ValueError:
-                     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.")
+                    logger.warning(f"Skipping Task ID: {item.get('task_id', 'N/A')} - Could not convert tool/step count: tools='{num_tools_str}', steps='{num_steps_str}'.")
+            # else: # If needed: log missing numbers in metadata
+                # logger.warning(f"Skipping Task ID: {item.get('task_id', 'N/A')} - 'Number of tools' or 'Number of steps' missing in Metadata.")
         else:
-             logger.warning(f"Skipping Task ID: {question.get('task_id', 'N/A')} - Missing 'Annotator Metadata'.")
+             logger.warning(f"Skipping Task ID: {item.get('task_id', 'N/A')} - Missing 'Annotator Metadata'.")
 
-    # Store the filtered list (optional, could process directly)
-    filtered_dataset = tempo_filtered
+    filtered_dataset = tempo_filtered # Store the list of filtered original items
     logger.info(f"Found {len(filtered_dataset)} questions matching the criteria (tools < {tool_threshold}, steps < {step_threshold}).")
 
-    # --- Processing Logic (Modified) ---
     processed_count = 0
+    # --- REVISED Processing Logic to match the new Pydantic model ---
     for item in filtered_dataset:
         task_id = item.get('task_id')
-        question_text = item.get('Question') # Keep original key for now
+        original_question_text = item.get('Question') # Get original text
         final_answer = item.get('Final answer')
 
-        # Validate required fields needed for processing/scoring
-        if task_id and question_text and final_answer is not None:
-            # 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
+        # Validate essential fields needed for processing & ground truth
+        if task_id and original_question_text and final_answer is not None:
+
+            # Create the dictionary for the API, selecting only the desired fields
+            processed_item = {
+                "task_id": str(task_id), # Ensure string type
+                "question": str(original_question_text), # Rename and ensure string type
+                # Include optional fields *if they exist* in the source item
+                "Level": item.get("Level"), # Use .get() for safety, Pydantic handles None
+                "file_name": item.get("file_name"),
+                "file_path": item.get("file_path"),
+            }
+            # Optional: Clean up None values if Pydantic model doesn't handle them as desired
+            # processed_item = {k: v for k, v in processed_item.items() if v is not None}
+            # However, the Optional[...] fields in Pydantic should handle None correctly.
+
+            # Append the structured dictionary matching the Pydantic model
             questions_for_api.append(processed_item)
 
-            # Store the ground truth answer separately for scoring
+            # Store the ground truth answer separately (as before)
             ground_truth_answers[str(task_id)] = str(final_answer)
             processed_count += 1
         else:
-            # 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.")
+            logger.warning(f"Skipping item due to missing essential fields (task_id, Question, or Final answer): task_id={task_id}")
 
+    logger.info(f"Successfully processed {processed_count} questions for the API matching the Pydantic model.")
     if not questions_for_api:
-         logger.error("CRITICAL: No valid questions loaded after filtering. API endpoints needing questions will fail.")
+         logger.error("CRITICAL: No valid questions loaded after filtering and processing. API endpoints needing questions will fail.")
          # raise RuntimeError("Failed to load mandatory question data after filtering.")
+    # --- END REVISED Processing Logic ---
+
+
 
 # --- 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 # Keep original casing if that's what in the data
+    question: str
+    Level: Optional[str] = None
+    file_name: Optional[str] = None
+    file_path: Optional[str] = None
+
+   
+# --- The rest of your Pydantic models remain the same ---
+class AnswerItem(BaseModel):
+    task_id: str
+    submitted_answer: str = Field(..., description="The agent's answer for the task_id")
+
+class Submission(BaseModel):
+    username: str = Field(..., description="Hugging Face username", min_length=1)
+    agent_code: str = Field(..., description="The Python class code for the agent", min_length=10) # Basic check
+    answers: List[AnswerItem] = Field(..., description="List of answers submitted by the agent")
+
+class ScoreResponse(BaseModel):
+    username: str
+    score: float
+    correct_count: int
+    total_attempted: int
+    message: str
+    timestamp: str
+
+class ErrorResponse(BaseModel):
+    detail: str
 
 # Keep other models as they are (AnswerItem, Submission, ScoreResponse, ErrorResponse)
 # ... (rest of the Pydantic models remain the same) ...