openenv/chat_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Data models for the Chat Environment.9 10The Chat environment provides a chat-based interface for LLMs with support11for tokenization and message history management.12"""13 14from openenv.core.env_server.types import Action, Observation, State15from pydantic import Field, field_validator16 17 18def _flatten_tokens(value) -> list[int]:19    """Coerce nested tensor-like or sequence inputs into a flat token list."""20    if hasattr(value, "tolist") and callable(value.tolist):21        value = value.tolist()22 23    if isinstance(value, tuple):24        value = list(value)25 26    if isinstance(value, list):27        flattened: list[int] = []28        for item in value:29            flattened.extend(_flatten_tokens(item))30        return flattened31 32    return [int(value)]33 34 35class ChatAction(Action):36    """Action for chat environments.37 38    Contains tokens that represent the action to be taken.39    This interfaces directly with models.40    """41 42    tokens: list[int] = Field(..., min_length=1)43 44    @field_validator("tokens", mode="before")45    @classmethod46    def _coerce_tokens(cls, value):47        """Accept either tensors or JSON arrays on the public HTTP surface."""48        if isinstance(value, (list, tuple)) or hasattr(value, "tolist"):49            return _flatten_tokens(value)50        raise TypeError("tokens must be provided as a sequence of token ids")51 52 53class ChatState(State):54    """State of the ChatEnvironment containing message history."""55 56    # TODO: revert to list[Message] once openenv-core ships typing_extensions.TypedDict57    # in interfaces.py and chat_env/pyproject.toml pins to that release.58    history_messages: list[dict[str, str]] = Field(default_factory=list)59    history_tokens: list[list[int]] = Field(default_factory=list)  # Same len as messages60 61 62class ChatObservation(Observation):63    """Observation returned by ChatEnvironment.64 65    Contains the message history in Huggingface format (list of dicts with role/content)66    and the tokenized representation of the entire conversation.67 68    The environment owns the tokenizer and generates the tokens from the messages.69 70    Example:71    messages = [72     {"role": "system", "content": "You are a helpful assistant"},73     {"role": "user", "content": "How tall is the Eiffel Tower?"},74    ]75    tokens = tensor([1, 2, 3, 4, 5, ...])  # tokenized entire conversation76    """77 78    # TODO: revert to list[Message] (same as above)79    messages: list[dict[str, str]] = Field(default_factory=list)80    tokens: list[int] = Field(default_factory=list)81    # Inherited Fields from Observation ABC: reward, done, metadata82