# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
# https://github.com/gepa-ai/gepa

from typing import Any, TypedDict


class MCPDataInst(TypedDict):
    """
    Dataset item for MCP tool optimization.

    Attributes:
        user_query: The user's question or request that may require tool usage
        tool_arguments: Expected arguments for the tool call (used for validation/guidance)
        reference_answer: Optional reference answer for scoring
        additional_context: Optional additional context for the task
    """

    user_query: str
    tool_arguments: dict[str, Any]
    reference_answer: str | None
    additional_context: dict[str, str]


class MCPTrajectory(TypedDict):
    """
    Trajectory capturing MCP tool invocation and model behavior.

    This captures the full execution trace including:
    - What tool was called and with what arguments
    - The tool's response
    - The tool description that was used
    - Model outputs at different stages
    - Tool selection process (for multi-tool scenarios)
    """

    user_query: str
    tool_names: list[str]  # Available tools
    selected_tool: str | None  # Which tool was selected (if any)
    tool_called: bool
    tool_arguments: dict[str, Any] | None
    tool_response: str | None
    tool_description_used: str
    system_prompt_used: str
    model_first_pass_output: str
    model_final_output: str
    score: float


class MCPOutput(TypedDict):
    """
    Output from MCP-enabled system evaluation.

    Attributes:
        final_answer: The final answer generated by the model
        tool_called: Whether a tool was called during execution
        selected_tool: Which tool was selected (if any)
        tool_response: The raw tool response (if tool was called)
    """

    final_answer: str
    tool_called: bool
    selected_tool: str | None
    tool_response: str | None
