models.py

"""
Data models for the Enterprise QA Environment.

Supports multi-domain strategic analysis tasks with rich observation fields.
"""

from typing import Any, Dict, Optional
from dataclasses import dataclass, field
from openenv_core.env_server.types import Action, Observation


@dataclass(kw_only=True)
class EnterpriseQaAction(Action):
    """Agent's response to a task prompt."""
    message: str
    agent_id: Optional[str] = None


@dataclass(kw_only=True)
class EnterpriseQaObservation(Observation):
    """
    Full observation returned after reset() or step().

    Fields:
        task_id:        Unique task identifier from data.csv
        domain:         Business domain (Finance, Healthcare, etc.)
        prompt:         The task description / question
        rubric:         Scoring criteria string
        echoed_message: The agent's last response (for debugging)
        done:           True when episode terminates
        reward:         Score in [0.0, 1.0]
        metadata:       Step count, gold_response, episode_id
    """
    task_id: str = ""
    domain: str = ""
    prompt: str = ""
    rubric: str = ""
    echoed_message: str = ""
    gold_response: str = ""