File size: 5,729 Bytes
4827e61
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
# 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.

"""
Generic environment client that works with raw dictionaries.

This module provides a GenericEnvClient that doesn't require installing
environment-specific packages. It's useful for connecting to remote servers
without running any untrusted code locally.
"""

from typing import Any, Dict

from .client_types import StepResult
from .env_client import EnvClient


class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]]):
    """
    Environment client that works with raw dictionaries instead of typed classes.

    This client doesn't require installing environment-specific packages, making it
    ideal for:
    - Connecting to remote servers without installing their packages
    - Quick prototyping and testing
    - Environments where type safety isn't needed
    - Security-conscious scenarios where you don't want to run remote code

    The trade-off is that you lose type safety and IDE autocomplete for actions
    and observations. Instead of typed objects, you work with plain dictionaries.

    Example:
        >>> # Direct connection to a running server (no installation needed)
        >>> with GenericEnvClient(base_url="http://localhost:8000") as env:
        ...     result = env.reset()
        ...     result = env.step({"code": "print('hello')"})
        ...     print(result.observation)  # Dict[str, Any]
        ...     print(result.observation.get("output"))

        >>> # From local Docker image
        >>> env = GenericEnvClient.from_docker_image("coding-env:latest")
        >>> result = env.reset()
        >>> result = env.step({"code": "x = 1 + 2"})
        >>> env.close()

        >>> # From HuggingFace Hub (pulls Docker image, no pip install)
        >>> env = GenericEnvClient.from_env("user/my-env", use_docker=True)
        >>> result = env.reset()
        >>> env.close()

    Note:
        GenericEnvClient inherits `from_docker_image()` and `from_env()` from
        EnvClient, so you can use it with Docker containers and HuggingFace
        Spaces without any package installation.
    """

    def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert action to payload for the server.

        For GenericEnvClient, this handles both raw dictionaries and
        typed Action objects (Pydantic models). If a Pydantic model is
        passed, it will be converted to a dictionary using model_dump().

        Args:
            action: Action as a dictionary or Pydantic BaseModel

        Returns:
            The action as a dictionary for the server
        """
        # If it's already a dict, return as-is
        if isinstance(action, dict):
            return action

        # If it's a Pydantic model (Action subclass), convert to dict
        if hasattr(action, "model_dump"):
            return action.model_dump()

        # Fallback for other objects with __dict__
        if hasattr(action, "__dict__"):
            return vars(action)

        # Last resort: try to convert to dict
        return dict(action)

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

        Extracts the observation, reward, and done fields from the
        server response.

        Args:
            payload: Response payload from the server

        Returns:
            StepResult with observation as a dictionary
        """
        return StepResult(
            observation=payload.get("observation", {}),
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Parse state response from the server.

        For GenericEnvClient, this returns the payload as-is since
        we're working with dictionaries.

        Args:
            payload: State payload from the server

        Returns:
            The state as a dictionary
        """
        return payload


class GenericAction(Dict[str, Any]):
    """
    A dictionary subclass for creating actions when using GenericEnvClient.

    This provides a semantic wrapper around dictionaries to make code more
    readable when working with GenericEnvClient. It behaves exactly like a
    dict but signals intent that this is an action for an environment.

    Example:
        >>> # Without GenericAction (works fine)
        >>> env.step({"code": "print('hello')"})

        >>> # With GenericAction (more explicit)
        >>> action = GenericAction(code="print('hello')")
        >>> env.step(action)

        >>> # With multiple fields
        >>> action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"})
        >>> env.step(action)

    Note:
        GenericAction is just a dict with a constructor that accepts keyword
        arguments. It's provided for symmetry with typed Action classes and
        to make code more readable.
    """

    def __init__(self, **kwargs: Any) -> None:
        """
        Create a GenericAction from keyword arguments.

        Args:
            **kwargs: Action fields as keyword arguments

        Example:
            >>> action = GenericAction(code="print(1)", timeout=30)
            >>> action["code"]
            'print(1)'
        """
        super().__init__(kwargs)

    def __repr__(self) -> str:
        """Return a readable representation."""
        items = ", ".join(f"{k}={v!r}" for k, v in self.items())
        return f"GenericAction({items})"