# 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 Tool Use Env Environment.

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

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


# class ToolUseAction(Action):
#     """Action for the Tool Use Env environment - just a message to echo."""

#     message: str = Field(..., description="Message to echo back")


# class ToolUseObservation(Observation):
#     """Observation from the Tool Use Env environment - the echoed message."""

#     echoed_message: str = Field(default="", description="The echoed message")
#     message_length: int = Field(default=0, description="Length of the echoed message")

from openenv.core.env_server import Action, Observation, State
from typing import Optional


class ToolUseAction(Action):
    action_type: str


class ToolUseObservation(Observation):
    query: str
    tool_output: Optional[str]
    message: str


class ToolUseState(State):
    current_query: str = ""
    correct_action: str = ""
    correct_answer: str = ""
    difficulty: str = ""