models.py

"""Data models for the Research -> Interactive Explainer environment."""

from typing import Any, Literal

from openenv.core.env_server.types import Action, Observation
from pydantic import Field

try:
    from .constants import MAX_EXPLORE_STEPS, MAX_REPAIR_STEPS
except ImportError:  # pragma: no cover - supports direct test execution
    from constants import MAX_EXPLORE_STEPS, MAX_REPAIR_STEPS


ResearchTool = Literal[
    "search_wikipedia",
    "search_hf_papers",
    "search_arxiv",
    "search_scholar",
    "fetch_docs",
    "search_hf_hub",
]


class ExplainerAction(Action):
    """Action: agent explores, generates, or repairs an artifact."""

    action_type: Literal["explore", "generate", "repair"] = Field(
        ...,
        description="'explore' to research, 'generate' to produce code, 'repair' to fix code",
    )

    # -- explore fields --
    tool: ResearchTool | None = Field(
        default=None,
        description="Research tool to call when action_type='explore'",
    )
    query: str = Field(
        default="",
        description="Research query used when action_type='explore'",
    )
    intent: str = Field(
        default="",
        description="Brief goal for the research call, e.g. equations or visual intuition",
    )

    # -- generate / repair fields --
    format: Literal["marimo", "manim"] | None = Field(
        default=None,
        description="Output format (required for generate/repair)",
    )
    code: str = Field(
        default="",
        description="Complete Python source code (required for generate/repair)",
    )
    narration: str = Field(
        default="",
        description="Narration script (used when format='manim')",
    )
    repair_notes: str = Field(
        default="",
        description="Short explanation of what changed when action_type='repair'",
    )


class ExplainerObservation(Observation):
    """Observation returned to the agent after each step."""

    # -- task info (set on reset, echoed back each step) --
    topic: str = Field(default="", description="Title of the topic or paper")
    content: str = Field(default="", description="Abstract or concept description")
    tier: Literal["beginner", "intermediate", "advanced"] = Field(
        default="beginner", description="Explanation depth tier"
    )
    keywords: str = Field(default="", description="Comma-separated key terms")
    data_available: bool = Field(
        default=False, description="Whether the topic references datasets"
    )
    difficulty: Literal["easy", "medium", "hard"] = Field(
        default="easy", description="Task difficulty tier"
    )

    # -- per-step feedback --
    phase: Literal["explore", "generate", "repair", "done"] = Field(
        default="explore", description="Current episode phase"
    )
    feedback: str = Field(default="", description="Feedback on the last action")
    search_results: str = Field(
        default="", description="Papers/snippets returned from an explore step"
    )
    top_chunks: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Ranked top chunks returned from the last explore step",
    )
    explored_context: str = Field(
        default="",
        description="Accumulated research context from all explore steps so far",
    )
    explore_steps_left: int = Field(
        default=MAX_EXPLORE_STEPS, description="Remaining explore steps before forced generate"
    )
    repair_attempts_left: int = Field(
        default=MAX_REPAIR_STEPS, description="Remaining repair attempts after failed generation"
    )
    last_errors: str = Field(
        default="", description="Latest lint/build errors available for repair"
    )
    available_tools: list[str] = Field(
        default_factory=list, description="Research tools available during explore"
    )