from typing import TypedDict, List, Dict, Optional, Any
from langchain_core.messages import BaseMessage

class JARVISState(TypedDict):
    """
    State dictionary for the JARVIS GAIA Agent, used with LangGraph to manage task processing.

    Attributes:
        task_id: Unique identifier for the GAIA task.
        question: The question text to be answered.
        tools_needed: List of tool names to be used for the task.
        web_results: List of web search results (e.g., from SERPAPI, DuckDuckGo).
        file_results: Parsed content from text, CSV, or Excel files.
        image_results: OCR or description results from image files.
        calculation_results: Results from mathematical calculations.
        document_results: Extracted content from PDF documents.
        multi_hop_results: Results from iterative multi-hop searches.
        messages: List of messages for LLM context (e.g., user prompts, system instructions).
        answer: Final answer for the task, formatted for GAIA submission.
        results_table: List of task results for Gradio display (Task ID, Question, Answer).
        status_output: Status message for Gradio output (e.g., submission result).
        error: Optional error message if task processing fails.
        metadata: Optional metadata (e.g., timestamps, tool execution status).
    """
    task_id: str
    question: str
    tools_needed: List[str]
    web_results: List[str]
    file_results: str
    image_results: str
    calculation_results: str
    document_results: str
    multi_hop_results: List[str]
    messages: List[BaseMessage]
    answer: str
    results_table: List[Dict[str, str]]
    status_output: str
    error: Optional[str]
    metadata: Optional[Dict[str, Any]]