gaiaX/api.py

#!/usr/bin/env python3
"""
API interaction module for GAIA Benchmark Agent.

This module handles all interactions with the GAIA benchmark API,
including fetching questions, downloading files, and submitting answers.
"""

import json
import requests
from typing import Dict, List, Any, Optional
from pathlib import Path

from gaiaX.config import logger, API_BASE_URL

def get_all_questions(api_base_url: str = API_BASE_URL) -> List[Dict[str, Any]]:
    """
    Retrieve all available questions from the GAIA benchmark.
    
    Args:
        api_base_url: Base URL for the GAIA API
        
    Returns:
        List of question dictionaries
        
    Raises:
        requests.RequestException: If the API request fails
        ValueError: If the response is not valid JSON or doesn't contain expected data
    """
    try:
        response = requests.get(f"{api_base_url}/questions")
        response.raise_for_status()  # Raise exception for 4XX/5XX responses
        
        questions = response.json()
        
        if not isinstance(questions, list):
            raise ValueError("Expected a list of questions but received a different format")
            
        return questions
    
    except requests.RequestException as e:
        logger.error(f"Error fetching questions: {e}")
        raise
    
    except json.JSONDecodeError:
        logger.error("Error decoding response as JSON")
        raise ValueError("Invalid JSON response from the API")


def get_random_question(api_base_url: str = API_BASE_URL) -> Dict[str, Any]:
    """
    Retrieve a random question from the GAIA benchmark.
    
    Args:
        api_base_url: Base URL for the GAIA API
        
    Returns:
        A single question dictionary
        
    Raises:
        requests.RequestException: If the API request fails
        ValueError: If the response is not valid JSON or doesn't contain expected data
    """
    try:
        response = requests.get(f"{api_base_url}/questions/random")
        response.raise_for_status()
        
        question = response.json()
        
        if not isinstance(question, dict):
            raise ValueError("Expected a question dictionary but received a different format")
            
        return question
    
    except requests.RequestException as e:
        logger.error(f"Error fetching random question: {e}")
        raise
    
    except json.JSONDecodeError:
        logger.error("Error decoding response as JSON")
        raise ValueError("Invalid JSON response from the API")


def download_file_for_task(api_base_url: str, task_id: str, download_path: str) -> str:
    """
    Download a file associated with a specific task.
    
    Args:
        api_base_url: Base URL for the GAIA API
        task_id: ID of the task to download files for
        download_path: Directory path where the file should be saved
        
    Returns:
        Path to the downloaded file
        
    Raises:
        requests.RequestException: If the API request fails
        IOError: If there's an error writing the file
        ValueError: If the task_id is invalid or the response is unexpected
    """
    if not task_id:
        raise ValueError("Task ID cannot be empty")
    
    # Ensure download directory exists
    download_dir = Path(download_path)
    download_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        response = requests.get(
            f"{api_base_url}/tasks/{task_id}/file",
            stream=True  # Stream the response for large files
        )
        response.raise_for_status()
        
        # Get filename from Content-Disposition header or use task_id as fallback
        content_disposition = response.headers.get('Content-Disposition', '')
        filename = None
        
        if 'filename=' in content_disposition:
            filename = content_disposition.split('filename=')[1].strip('"\'')
        
        if not filename:
            filename = f"{task_id}_file.txt"
        
        file_path = download_dir / filename
        
        # Write the file
        with open(file_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        return str(file_path)
    
    except requests.RequestException as e:
        logger.error(f"Error downloading file for task {task_id}: {e}")
        raise
    
    except IOError as e:
        logger.error(f"Error writing file to {download_path}: {e}")
        raise


def submit_answers(
    api_base_url: str, 
    username: str, 
    agent_code_link: str, 
    answers: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Submit answers to the GAIA benchmark.
    
    Args:
        api_base_url: Base URL for the GAIA API
        username: Hugging Face username
        agent_code_link: Link to the agent code (e.g., GitHub repository)
        answers: Dictionary of answers to submit
        
    Returns:
        Response from the API containing submission results
        
    Raises:
        requests.RequestException: If the API request fails
        ValueError: If the response is not valid JSON or contains an error message
    """
    if not username:
        raise ValueError("Username cannot be empty")
    
    if not agent_code_link:
        raise ValueError("Agent code link cannot be empty")
    
    if not answers or not isinstance(answers, dict):
        raise ValueError("Answers must be a non-empty dictionary")
    
    payload = {
        "username": username,
        "agent_code_link": agent_code_link,
        "answers": answers
    }
    
    try:
        response = requests.post(
            f"{api_base_url}/submit",
            json=payload,
            headers={"Content-Type": "application/json"}
        )
        response.raise_for_status()
        
        result = response.json()
        
        # Check if the response contains an error message
        if isinstance(result, dict) and result.get("error"):
            raise ValueError(f"API returned an error: {result['error']}")
            
        return result
    
    except requests.RequestException as e:
        logger.error(f"Error submitting answers: {e}")
        raise
    
    except json.JSONDecodeError:
        logger.error("Error decoding response as JSON")
        raise ValueError("Invalid JSON response from the API")


def get_question_details(task_id: str, api_base_url: str = API_BASE_URL) -> Dict[str, Any]:
    """
    Get detailed information about a specific question/task.
    
    Args:
        task_id: The ID of the task to get details for
        api_base_url: Base URL for the GAIA API
        
    Returns:
        Dictionary containing question details
    """
    try:
        response = requests.get(f"{api_base_url}/questions/{task_id}")
        response.raise_for_status()
        return response.json()
    except requests.RequestException as e:
        logger.error(f"Failed to get question details: {str(e)}")
        return {"error": f"Failed to get question details: {str(e)}"}
    except json.JSONDecodeError:
        logger.error("Invalid JSON response from the API")
        return {"error": "Invalid JSON response from the API"}