manzz05/kitchenflow-v2

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 the
5# LICENSE file in the root directory of this source tree.
6
7"""Kitchenflow Env Environment Client."""
8
9from typing import Dict
10
11from openenv.core import EnvClient
12from openenv.core.client_types import StepResult
13from openenv.core.env_server.types import State
14
15from .models import KitchenflowAction, KitchenflowObservation
16
17
18class KitchenflowEnv(
19    EnvClient[KitchenflowAction, KitchenflowObservation, State]
20):
21    """
22    Client for the Kitchenflow Env Environment.
23
24    This client maintains a persistent WebSocket connection to the environment server,
25    enabling efficient multi-step interactions with lower latency.
26    Each client instance has its own dedicated environment session on the server.
27
28    Example:
29        >>> # Connect to a running server
30        >>> with KitchenflowEnv(base_url="http://localhost:8000") as client:
31        ...     result = client.reset()
32        ...     print(result.observation.echoed_message)
33        ...
34        ...     result = client.step(KitchenflowAction(message="Hello!"))
35        ...     print(result.observation.echoed_message)
36
37    Example with Docker:
38        >>> # Automatically start container and connect
39        >>> client = KitchenflowEnv.from_docker_image("kitchenflow_env-env:latest")
40        >>> try:
41        ...     result = client.reset()
42        ...     result = client.step(KitchenflowAction(message="Test"))
43        ... finally:
44        ...     client.close()
45    """
46
47    def _step_payload(self, action: KitchenflowAction) -> Dict:
48        """
49        Convert KitchenflowAction to JSON payload for step message.
50
51        Args:
52            action: KitchenflowAction instance
53
54        Returns:
55            Dictionary representation suitable for JSON encoding
56        """
57        return {
58            "message": action.message,
59        }
60
61    def _parse_result(self, payload: Dict) -> StepResult[KitchenflowObservation]:
62        """
63        Parse server response into StepResult[KitchenflowObservation].
64
65        Args:
66            payload: JSON response data from server
67
68        Returns:
69            StepResult with KitchenflowObservation
70        """
71        obs_data = payload.get("observation", {})
72        observation = KitchenflowObservation(
73            echoed_message=obs_data.get("echoed_message", ""),
74            message_length=obs_data.get("message_length", 0),
75            done=payload.get("done", False),
76            reward=payload.get("reward"),
77            metadata=obs_data.get("metadata", {}),
78        )
79
80        return StepResult(
81            observation=observation,
82            reward=payload.get("reward"),
83            done=payload.get("done", False),
84        )
85
86    def _parse_state(self, payload: Dict) -> State:
87        """
88        Parse server response into State object.
89
90        Args:
91            payload: JSON response from state request
92
93        Returns:
94            State object with episode_id and step_count
95        """
96        return State(
97            episode_id=payload.get("episode_id"),
98            step_count=payload.get("step_count", 0),
99        )
100