models.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the Mentalhealthpatientenv Environment.

The mentalHealthPatientenv environment is a simple test environment that echoes back messages.
"""

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


class MentalhealthpatientenvAction(Action):
    """
    Represents an action taken by the agent in the mentalHealthPatientenv environment.
        Attributes:
        action_type (Literal["ask_open", "ask_direct", "ask_risk", "reflect", "diagnose","Initialize"]): The type of action the agent is taking, which can be one of the following:
            - "ask_open": An open-ended question to encourage the patient to share more information.
            - "ask_direct": A direct question to obtain specific information from the patient.
            - "ask_risk": A question aimed at assessing any potential risks related to the patient's mental state.
            - "reflect": A reflective statement to show understanding and encourage further discussion.
            - "diagnose": An action where the agent provides a diagnosis based on the information gathered.
            - "Initialize": An action to initialize the environment.
        task_difficulty (Literal["easy", "medium", "hard"]): The difficulty level of the task the agent is trying to accomplish, which can be "easy", "medium", or "hard". This can influence the complexity of the agent's actions and the expected responses from the patient.
        message (str): The message content that the agent sends to the patient. Defaults to an empty string.
    """

    action_type: Literal[
        "ask_open", "ask_direct", "ask_risk", "reflect", "diagnose","Initialize"
    ] = Field(
        default="ask_open",
        description="Type of action"
    )
    task_difficulty: Literal[
        "easy", "medium", "hard"
    ] = Field(
        default="medium",
        description="Difficulty level of the task"
    )
    message: str = Field(
        default="",
        description="Agent message to patient"
    )



class MentalhealthpatientenvObservation(Observation):
    """
    Represents the patient's response and state after an action is taken in the mentalHealthPatientenv environment.
        Attributes: 
        response (str): The patient's reply to the agent's message. Defaults to an empty string.
        clarity (float): A value between 0.0 and 1.0 indicating how clear and detailed the patient's response is, where 0.0 means vague and 1.0 means very detailed. Defaults to 0.5.
        emotional_state (Literal["sad", "anxious", "neutral", "angry"]): The patient's emotional state based on their response. Defaults to "neutral".
        trust_level (float): A value between 0.0 and 1.0 indicating the patient's trust level towards the agent, where 0.0 means no trust and 1.0 means full trust. Defaults to 0.5.
        risk_flag (bool): A boolean flag indicating whether the patient's response contains any risk-related information. Defaults to False.
        done (bool): A boolean flag indicating whether the interaction session is complete. Defaults to False.
    """
    response: str = Field(
        default="",
        description="Patient reply"
    )

    clarity: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="0 = vague, 1 = detailed"
    )

    emotional_state: Literal[
        "sad", "anxious", "neutral", "angry"
    ] = Field(
        default="neutral"
    )

    trust_level: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Trust towards agent"
    )

    risk_flag: bool = Field(
        default=False
    )

    done: bool = Field(
        default=False
    )