client.py

"""
client.py -- Typed Python client for HypothesisLab.

Built on openenv.core.env_client.EnvClient (WebSocket-based, persistent).
This is what the RL trainer and baseline inference script import.
"""

from __future__ import annotations

from typing import Any, Dict, Optional

try:
    from openenv.core.env_client import EnvClient
    from openenv.core.client_types import StepResult
except ImportError:
    raise ImportError(
        "openenv-core is required. Install with: pip install openenv-core"
    )

try:
    from .models import (
        ActionType,
        ExperimentType,
        HypLabAction,
        HypLabObservation,
        HypLabState,
        NoiseLevelTag,
    )
except ImportError:
    from models import (
        ActionType,
        ExperimentType,
        HypLabAction,
        HypLabObservation,
        HypLabState,
        NoiseLevelTag,
    )


class HypothesisLabEnv(EnvClient[HypLabAction, HypLabObservation, HypLabState]):
    """
    Typed async client for the Scientific Hypothesis Lab environment.

    Usage (async):
        async with HypothesisLabEnv(base_url="http://localhost:8000") as env:
            result = await env.reset(noise_level="low", domain="physics")
            obs = result.observation
            ...

    Usage (sync):
        env = HypothesisLabEnv(base_url="http://localhost:8000").sync()
        with env:
            result = env.reset(noise_level="low")
            ...
    """

    def _step_payload(self, action: HypLabAction) -> Dict[str, Any]:
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[HypLabObservation]:
        obs_data = payload.get("observation", payload)
        obs = HypLabObservation(**obs_data)
        return StepResult(
            observation=obs,
            reward=payload.get("reward", obs.reward or 0.0),
            done=payload.get("done", obs.done),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> HypLabState:
        return HypLabState(**payload)

    async def run_intervention(
        self,
        control_variable: str,
        control_value: float,
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.INTERVENTION,
            control_variable=control_variable,
            target_variable=target_variable,
            control_value=control_value,
        )
        return await self.step(action)

    async def run_correlation(
        self,
        control_variable: str,
        control_range: list[float],
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.CORRELATION,
            control_variable=control_variable,
            control_range=control_range,
            target_variable=target_variable,
        )
        return await self.step(action)

    async def run_counterfactual(
        self,
        control_variable: str,
        delta: float,
        target_variable: str,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.COUNTERFACTUAL,
            control_variable=control_variable,
            control_value=delta,
            target_variable=target_variable,
        )
        return await self.step(action)

    async def run_passive(
        self, target_variable: str
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.EXPERIMENT,
            experiment_type=ExperimentType.PASSIVE,
            target_variable=target_variable,
            control_variable=target_variable,
        )
        return await self.step(action)

    async def submit_hypothesis(
        self,
        hypothesis_text: str,
        hypothesis_equations: Optional[list[str]] = None,
        confidence: float = 0.75,
    ) -> StepResult[HypLabObservation]:
        action = HypLabAction(
            action_type=ActionType.SUBMIT,
            hypothesis_text=hypothesis_text,
            hypothesis_equations=hypothesis_equations,
            confidence=max(0.0, min(1.0, confidence)),
        )
        return await self.step(action)