models.py

"""
Models for the Data Extraction Environment.
Self-contained — no openenv dependency needed at runtime.
"""

from pydantic import BaseModel, Field


class ExtractionAction(BaseModel):
    """What the LLM sends back — a JSON dict of extracted fields."""
    extracted_data: dict = Field(
        ...,
        description="A dictionary mapping field names to extracted values"
    )


class ExtractionObservation(BaseModel):
    """What the environment shows the agent."""
    task_id: str = Field(default="")
    difficulty: str = Field(default="easy")
    raw_text: str = Field(default="")
    fields_to_extract: list = Field(default_factory=list)
    extraction_hints: str = Field(default="")
    feedback: str = Field(default="")
    fields_correct: int = Field(default=0)
    fields_total: int = Field(default=0)
    field_scores: dict = Field(default_factory=dict)
    reward: float = Field(default=0.0)
    done: bool = Field(default=False)