Specification

JSONType = Union[Dict[str, Any], List[Any], str, int, float, bool, None]

class ChatCompletionContentPartTextParam(BaseModel):
    text: str
    type: Literal["text"] = "text"

class Message(BaseModel):
    role: str  # assistant, user, system, tool
    content: Optional[Union[str, List[ChatCompletionContentPartTextParam]]] = ""
    reasoning_content: Optional[str] = None
    name: Optional[str] = None
    tool_call_id: Optional[str] = None
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
    function_call: Optional[FunctionCall] = None
    control_plane_step: Optional[Dict[str, Any]] = None

CompletionParams = Dict[str, Any]
"""
Provider-agnostic completion parameters.

Required:
- model: str

Common fields:
- temperature: Optional[float]
- max_tokens: Optional[int]
- top_p: Optional[float]

Extra provider-specific fields are allowed and passed through (e.g., max_tool_calls).
"""

class InputMetadata(BaseModel):
    # Accepts additional keys for future extensibility
    # (model_config = ConfigDict(extra="allow") in implementation)

    row_id: Optional[str]  # defaulted to a generated ID
    completion_params: CompletionParams = Field(default_factory=dict)
    dataset_info: Optional[Dict[str, Any]]  # seed, system_prompt, environment_context, etc.
    session_data: Optional[Dict[str, Any]]

class ErrorInfo(BaseModel):
    reason: str
    domain: str
    metadata: Dict[str, Any] = {}

class Status(BaseModel):
    class Code(int, Enum):
        OK = 0
        CANCELLED = 1
        UNKNOWN = 2
        INVALID_ARGUMENT = 3
        DEADLINE_EXCEEDED = 4
        NOT_FOUND = 5
        ALREADY_EXISTS = 6
        PERMISSION_DENIED = 7
        RESOURCE_EXHAUSTED = 8
        FAILED_PRECONDITION = 9
        ABORTED = 10
        OUT_OF_RANGE = 11
        UNIMPLEMENTED = 12
        INTERNAL = 13
        UNAVAILABLE = 14
        DATA_LOSS = 15
        UNAUTHENTICATED = 16

        # Custom codes used by Eval Protocol
        FINISHED = 100
        RUNNING = 101
        SCORE_INVALID = 102

    code: Code
    message: str
    details: List[Dict[str, Any]] = []

class TerminationReason(str, Enum):
    MAX_STEPS = "max_steps"
    CONTROL_PLANE_SIGNAL = "control_plane_signal"
    USER_STOP = "user_stop"
    SKIPPABLE_ERROR = "skippable_error"
    NON_SKIPPABLE_ERROR = "non_skippable_error"
    STOP = "stop"
    LENGTH = "length"
    TOOL_CALLS = "tool_calls"

class MetricResult(BaseModel):
    is_score_valid: bool = True
    score: float  # Between 0.0 and 1.0
    reason: str  # Explanation for the score
    data: Dict[str, Any] = Field(default_factory=dict)  # Optional extra metric data

class StepOutput(BaseModel):
    step_index: Union[int, str]  # User-defined index for the step
    base_reward: float  # Base reward calculated by the user's reward function
    terminated: bool = False  # Whether the environment signaled termination
    control_plane_info: Optional[Dict[str, Any]]  # Structured info from environment
    metrics: Dict[str, Any] = Field(default_factory=dict)  # Optional custom metrics
    reason: Optional[str]  # Optional explanation for the step's base reward

class EvaluationThreshold(BaseModel):
    success: float  # Minimum success rate threshold (0.0 to 1.0)
    standard_error: Optional[float]  # Optional maximum standard error threshold

class EvalMetadata(BaseModel):
    name: str
    description: Optional[str]
    version: str  # PEP 440 version string (auto-populated)
    status: Optional[Status]
    num_runs: int
    aggregation_method: str
    passed_threshold: Optional[EvaluationThreshold]
    passed: Optional[bool]

class CostMetrics(BaseModel):
    input_cost: Optional[float]
    output_cost: Optional[float]
    total_cost_dollar: Optional[float]

class ExecutionMetadata(BaseModel):
    invocation_id: Optional[str]
    experiment_id: Optional[str]
    rollout_id: Optional[str]
    run_id: Optional[str]
    usage: Optional[CompletionUsage]
    cost_metrics: Optional[CostMetrics]
    duration_seconds: Optional[float]
    experiment_duration_seconds: Optional[float]

class EvaluateResult(BaseModel):
    # Core evaluation data
    score: float  # Overall evaluation score (0.0 to 1.0)
    is_score_valid: bool  # Whether the overall score is valid (defaults to True)
    reason: Optional[str]  # Optional explanation for the overall score
    
    # Component metrics
    metrics: Dict[str, MetricResult]  # Dictionary of component metrics
    
    # RL-specific fields
    step_outputs: Optional[List[StepOutput]]  # Per-step base rewards for RL
    
    # Error handling
    error: Optional[str]  # Optional error message if evaluation failed
    
    # Trajectory information
    trajectory_info: Optional[Dict[str, Any]]  # Additional trajectory-level information
    final_control_plane_info: Optional[Dict[str, Any]]  # Final control plane state

    # Aggregation across runs
    agg_score: Optional[float]  # Aggregated score across runs
    standard_error: Optional[float]  # Standard error across runs

class EvaluationRow(BaseModel):
    # Core conversation (trajectory) data
    messages: List[Message]

    # Tool and function call information
    tools: Optional[List[Dict[str, Any]]] = None

    # Input-related metadata
    input_metadata: InputMetadata = Field(default_factory=InputMetadata)

    # Rollout status (AIP-193)
    rollout_status: Status = Field(default_factory=Status.rollout_running)

    # Optional ground truth reference
    ground_truth: Optional[JSONType] = None

    # Unified evaluation result
    evaluation_result: Optional[EvaluateResult] = None

    # Correlation identifiers grouped under execution metadata
    execution_metadata: ExecutionMetadata = Field(default_factory=lambda: ExecutionMetadata(run_id=None))

    # Timestamps and evaluation metadata
    created_at: datetime = Field(default_factory=datetime.now)
    eval_metadata: Optional[EvalMetadata] = None

    # Process info for watchdogs
    pid: Optional[int] = None

from eval_protocol.pytest import evaluation_test, SingleTurnRolloutProcessor

@evaluation_test(
    input_dataset=["tests/pytest/data/markdown_dataset.jsonl"],
    dataset_adapter=markdown_dataset_to_evaluation_row,
    completion_params=[{
        "model": "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct",
        "temperature": 0.0,
        "max_tokens": 4096,
    }],
    passed_threshold={"success": 0.5},
    rollout_processor=SingleTurnRolloutProcessor(),
    num_runs=1,
    mode="pointwise",
)
def test_markdown_highlighting_evaluation(row: EvaluationRow) -> EvaluationRow:
    ...

class McpGym(ABC):
    def __init__(self, server_name: str, adapter: EnvironmentAdapter, seed: Optional[int] = None, max_workers: Optional[int] = None):
        ...

    @abstractmethod
    def _register_tools(self):
        ...

    def format_observation(self, obs: Any, env: Any) -> Dict[str, Any]:
        ...

    def run(self, transport: str = "streamable-http", **kwargs):
        ...

class EnvironmentAdapter:
    """
    Environment adapter with default implementations.
    
    Users can either use this class directly by providing an env_class,
    or inherit from it to customize specific methods for their environment.
    This provides a clean separation between the MCP protocol layer
    and the environment implementation.
    """

class LiteLLMPolicy(LLMBasePolicy):
    """
    Unified LiteLLM policy implementation that works with ANY MCP environment via tool calling.
    
    Supports OpenAI, Anthropic, Fireworks AI
    Includes built-in retry logic and caching.
    NO environment-specific logic - everything comes from MCP tools and dataset prompts.
    """

@dataclass
class MCPSession:
    session_id: str
    base_url: str
    seed: Optional[int]
    model_id: str
    dataset_row: Optional[DatasetRow] = None
    terminated: bool = False
    last_observation: Any = None
    _exit_stack: Optional[AsyncExitStack] = None  # persistent connection resources
    _mcp_session: Optional[ClientSession] = None  # persistent MCP client session

@dataclass
class Trajectory:
    session: MCPSession
    observations: List[Any]
    actions: List[str]
    rewards: List[float]
    terminated: bool
    total_reward: float
    steps: int
    duration: float
    control_plane_steps: List[Dict[str, Any]]
    control_plane_summary: Dict[str, Any]
    termination_reason: str
    conversation_history: List[Dict[str, Any]]
    usage: Dict[str, int] = field(default_factory=dict)