client.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.

"""Episteme Environment Client.

Provides a typed client for the Episteme environment server.  Extends
:class:`EnvClient` parameterised with :class:`EpistemeAction` and
:class:`EpistemeObservation`.

Example (async)::

    async with EpistemeEnv(base_url="http://localhost:8000") as client:
        result = await client.reset(task_type="fact_check", difficulty=2)
        print(result.observation.result)

        result = await client.step(
            EpistemeAction(tool_name="search_corpus",
                           arguments={"query": "eiffel tower"})
        )

Example (sync wrapper)::

    env = EpistemeEnv(base_url="http://localhost:8000").sync()
    with env:
        result = env.reset(task_type="research_brief", difficulty=1)
        result = env.step(
            EpistemeAction(tool_name="read_notes", arguments={})
        )
"""

from __future__ import annotations

from typing import Any, Dict, Optional

from openenv.core.env_client import EnvClient
from openenv.core.client_types import StepResult
from openenv.core.env_server.types import State

from .models import EpistemeAction, EpistemeObservation


class EpistemeEnv(
    EnvClient[EpistemeAction, EpistemeObservation, State]
):
    """Typed WebSocket client for the Episteme environment.

    Inherits connection management, ``step``, and ``state`` from
    :class:`EnvClient`.  Adds an explicit ``reset`` method that
    exposes ``task_type`` and ``difficulty`` as keyword arguments.
    """

    # ------------------------------------------------------------------
    # reset — surfaces Episteme-specific kwargs
    # ------------------------------------------------------------------

    async def reset(
        self,
        *,
        task_type: str = "fact_check",
        difficulty: int = 1,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> StepResult[EpistemeObservation]:
        """Reset the environment and start a new episode.

        All keyword arguments are forwarded to the server's ``reset``
        endpoint via the WebSocket ``reset`` message.

        Args:
            task_type: One of ``"fact_check"``, ``"research_brief"``,
                ``"market_analysis"``, or ``"notes"``.
            difficulty: Corpus difficulty level (1-5).
            seed: Optional random seed for reproducibility.
            episode_id: Optional custom episode identifier.
            **kwargs: Additional parameters forwarded to the server.

        Returns:
            :class:`StepResult` containing the initial
            :class:`EpistemeObservation`.
        """
        reset_kwargs: Dict[str, Any] = {
            "task_type": task_type,
            "difficulty": difficulty,
            **kwargs,
        }
        if seed is not None:
            reset_kwargs["seed"] = seed
        if episode_id is not None:
            reset_kwargs["episode_id"] = episode_id

        return await super().reset(**reset_kwargs)

    # ------------------------------------------------------------------
    # Serialisation / deserialisation hooks
    # ------------------------------------------------------------------

    def _step_payload(self, action: EpistemeAction) -> Dict[str, Any]:
        """Convert an :class:`EpistemeAction` to the JSON payload
        expected by the environment server's ``step`` endpoint.
        """
        return action.model_dump()

    def _parse_result(
        self, payload: Dict[str, Any]
    ) -> StepResult[EpistemeObservation]:
        """Parse a server response into a typed :class:`StepResult`.

        The framework's ``serialize_observation`` places ``reward``,
        ``done``, and ``metadata`` at the top level of ``payload``
        (outside the nested ``observation`` dict), so we source them
        from there.
        """
        obs_data = payload.get("observation", {})

        # reward / done / metadata live at the payload root,
        # NOT inside obs_data (the framework strips them).
        reward = payload.get("reward", 0.0)
        done = payload.get("done", False)

        observation = EpistemeObservation(
            result=obs_data.get("result", ""),
            reward=reward if reward is not None else 0.0,
            done=done,
            step_count=obs_data.get("step_count", 0),
            steps_remaining=obs_data.get("steps_remaining", 0),
            success=obs_data.get("success", True),
        )

        return StepResult(
            observation=observation,
            reward=reward,
            done=done,
        )

    def _parse_state(self, payload: Dict[str, Any]) -> State:
        """Parse a server state response into a :class:`State` object."""
        return State(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
        )