Update main.py

main.py

@@ -114,76 +114,7 @@ def load_questions():
     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/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}")
-
-    # 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):
@@ -259,8 +190,76 @@ async def startup_event():
         # import sys
         # sys.exit(1) # Consider exiting if questions are critical
 
-# --- Helper Function (update_huggingface_dataset remains the same) ---
-# ... (update_huggingface_dataset function code) ...
+# --- 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}")
+
+    # Use FileResponse to efficiently stream the file
+    return FileResponse(path=abs_file_path, media_type=media_type, filename=file_name_for_download)
 def update_huggingface_dataset(username: str, score: float):
     """Loads the dataset, updates the score if higher, and pushes back."""
     try: