Update main.py

main.py

@@ -2,6 +2,7 @@
 import os
 import pandas as pd
 from fastapi import FastAPI, HTTPException, Body
+from fastapi.responses import FileResponse
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 from datasets import load_dataset, Dataset, DatasetDict
@@ -14,117 +15,183 @@ import random
 # --- Constants and Config ---
 HF_DATASET_ID = "agents-course/unit4-students-scores"
 
-# --- 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__)
 
-
-logger = logging.getLogger(__name__) # Make sure logger is initialized
+task_file_paths: Dict[str, str] = {}
 tool_threshold = 3
 step_threshold = 5
-questions_for_api: List[Dict[str, Any]] = [] # Use Dict[str, Any] for flexibility before validation
+questions_for_api: List[Dict[str, Any]] = []
 ground_truth_answers: Dict[str, str] = {}
-filtered_dataset = None # Or initialize as empty list: []
+filtered_dataset = None
+# --- Define ErrorResponse if not already defined ---
+class ErrorResponse(BaseModel):
+    detail: str
 
 def load_questions():
     global filtered_dataset
     global questions_for_api
     global ground_truth_answers
+    global task_file_paths # Declare modification of global
     tempo_filtered = []
     # Clear existing data
     questions_for_api.clear()
     ground_truth_answers.clear()
+    task_file_paths.clear() # Clear the mapping too
 
     logger.info("Starting to load and filter GAIA dataset (validation split)...")
     try:
-        # 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: {e}", exc_info=True)
         raise RuntimeError("Could not load the primary GAIA dataset.") from e
 
-    # --- Filtering Logic (remains the same) ---
+    # --- Filtering Logic (remains same) ---
+    # [ ... Same filtering code as before ... ]
     for item in dataset:
         metadata = item.get('Annotator Metadata')
-
-        if metadata:
+        if metadata: # Check if 'Annotator Metadata' exists
             num_tools_str = metadata.get('Number of tools')
             num_steps_str = metadata.get('Number of steps')
-
             if num_tools_str is not None and num_steps_str is not None:
                 try:
                     num_tools = int(num_tools_str)
                     num_steps = int(num_steps_str)
-
                     if num_tools < tool_threshold and num_steps < step_threshold:
-                        tempo_filtered.append(item) # Add the original item if it matches filter
+                        tempo_filtered.append(item)
                 except ValueError:
-                    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: {item.get('task_id', 'N/A')} - Missing 'Annotator Metadata'.")
+                     logger.warning(f"Skipping Task ID: {item.get('task_id', 'N/A')} - Could not convert tool/step count.")
+        # else: # Log missing metadata if needed
+             # logger.warning(f"Skipping Task ID: {item.get('task_id', 'N/A')} - Missing 'Annotator Metadata'.")
+
 
-    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}).")
+    filtered_dataset = tempo_filtered
+    logger.info(f"Found {len(filtered_dataset)} questions matching the criteria.")
 
     processed_count = 0
-    # --- REVISED Processing Logic to match the new Pydantic model ---
+    # --- Processing Logic (includes storing file path mapping) ---
     for item in filtered_dataset:
         task_id = item.get('task_id')
-        original_question_text = item.get('Question') # Get original text
+        original_question_text = item.get('Question')
         final_answer = item.get('Final answer')
+        local_file_path = item.get('file_path') # Get the local path
+        file_name = item.get('file_name') # Get the filename
 
-        # Validate essential fields needed for processing & ground truth
+        # Validate essential fields
         if task_id and original_question_text and final_answer is not None:
-
-            # Create the dictionary for the API, selecting only the desired fields
+            # Create the dictionary for the API (WITHOUT file_path)
             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"),
+                "task_id": str(task_id),
+                "question": str(original_question_text),
+                "Level": item.get("Level"),
+                "file_name": file_name, # Include filename for info
             }
-            # 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.
+            # Clean None values if you prefer not to send nulls for optional fields
+            processed_item = {k: v for k, v in processed_item.items() if v is not None}
 
-            # Append the structured dictionary matching the Pydantic model
             questions_for_api.append(processed_item)
 
-            # Store the ground truth answer separately (as before)
+            # Store ground truth
             ground_truth_answers[str(task_id)] = str(final_answer)
+
+            # --- Store the file path mapping ---
+            if local_file_path and file_name: # Only store if path and name exist
+                 # Basic check if path looks plausible (optional)
+                 if os.path.exists(local_file_path):
+                      task_file_paths[str(task_id)] = local_file_path
+                      logger.debug(f"Stored file path for task_id {task_id}: {local_file_path}")
+                 else:
+                      logger.warning(f"File path '{local_file_path}' for task_id {task_id} does not exist on server. Mapping skipped.")
+
+
             processed_count += 1
         else:
-            logger.warning(f"Skipping item due to missing essential fields (task_id, Question, or Final answer): task_id={task_id}")
+            logger.warning(f"Skipping item due to missing essential fields: task_id={task_id}")
 
-    logger.info(f"Successfully processed {processed_count} questions for the API matching the Pydantic model.")
+    logger.info(f"Successfully processed {processed_count} questions for the API.")
+    logger.info(f"Stored file path mappings for {len(task_file_paths)} tasks.")
     if not questions_for_api:
-         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 ---
+         logger.error("CRITICAL: No valid questions loaded after filtering/processing.")
+# --- Add this endpoint definition to your FastAPI app ---
+
+# Determine a base path for security. This should be the root directory
+# where Hugging Face datasets cache is allowed to serve files from.
+# IMPORTANT: Adjust this path based on your server's environment or use
+# environment variables for configuration.
+# Using expanduser handles '~' correctly.
+ALLOWED_CACHE_BASE = os.path.abspath(os.path.expanduser("~/.cache/huggingface/datasets"))
+logger.info(f"Configured allowed base path for file serving: {ALLOWED_CACHE_BASE}")
+
+@app.get("/files/{task_id}",
+         summary="Get Associated File by Task ID",
+         description="Downloads the file associated with the given task_id, if one exists and is mapped.",
+         responses={
+             200: {
+                 "description": "File content.",
+                 "content": {"*/*": {}} # Indicates response can be any file type
+             },
+             403: {"model": ErrorResponse, "description": "Access denied (e.g., path traversal attempt)."},
+             404: {"model": ErrorResponse, "description": "Task ID not found, no file associated, or file missing on server."},
+             500: {"model": ErrorResponse, "description": "Server error reading file."}
+         })
+async def get_task_file(task_id: str):
+    """
+    Serves the file associated with a specific task ID.
+    Includes security checks to prevent accessing arbitrary files.
+    """
+    logger.info(f"Request received for file associated with task_id: {task_id}")
+
+    if task_id not in task_file_paths:
+        logger.warning(f"File request failed: task_id '{task_id}' not found in file path mapping.")
+        raise HTTPException(status_code=404, detail=f"No file path associated with task_id {task_id}.")
+
+    local_file_path = task_file_paths[task_id]
+    logger.debug(f"Mapped task_id '{task_id}' to local path: {local_file_path}")
+
+    # --- CRUCIAL SECURITY CHECK ---
+    try:
+        # Resolve to absolute paths to prevent '..' tricks
+        abs_file_path = os.path.abspath(local_file_path)
+        abs_base_path = ALLOWED_CACHE_BASE # Already absolute
+
+        # Check if the resolved file path starts with the allowed base directory
+        if not abs_file_path.startswith(abs_base_path):
+            logger.error(f"SECURITY ALERT: Path traversal attempt denied for task_id '{task_id}'. Path '{local_file_path}' resolves outside base '{abs_base_path}'.")
+            raise HTTPException(status_code=403, detail="File access denied.")
+
+        # Check if the file exists at the resolved, validated path
+        if not os.path.exists(abs_file_path) or not os.path.isfile(abs_file_path):
+             logger.error(f"File not found on server for task_id '{task_id}' at expected path: {abs_file_path}")
+             raise HTTPException(status_code=404, detail=f"File associated with task_id {task_id} not found on server disk.")
+
+    except HTTPException as http_exc:
+         raise http_exc # Re-raise our own security/404 exceptions
+    except Exception as path_err:
+         logger.error(f"Error resolving or checking path '{local_file_path}' for task_id '{task_id}': {path_err}", exc_info=True)
+         raise HTTPException(status_code=500, detail="Server error validating file path.")
+    # --- END SECURITY CHECK ---
+
+    # Determine MIME type for the Content-Type header
+    mime_type, _ = mimetypes.guess_type(abs_file_path)
+    media_type = mime_type if mime_type else "application/octet-stream" # Default if unknown
 
+    # Extract filename for the Content-Disposition header (suggests filename to browser/client)
+    file_name_for_download = os.path.basename(abs_file_path)
 
+    logger.info(f"Serving file '{file_name_for_download}' (type: {media_type}) for task_id '{task_id}' from path: {abs_file_path}")
 
-# --- Pydantic Models ---
+    # Use FileResponse to efficiently stream the file
+    return FileResponse(path=abs_file_path, media_type=media_type, filename=file_name_for_download)
 
 
 class Question(BaseModel):
     task_id: str
     question: str
     Level: Optional[str] = None
-    file_name: Optional[str] = None
-    file_path: Optional[str] = None
+    file_name: Optional[str] = None # Keep filename for info
+    # file_path: Optional[str] = None # REMOVE file_path from the response model
 
    
 # --- The rest of your Pydantic models remain the same ---