Create environments.py

environments.py

@@ -0,0 +1,121 @@
+import inspect
+import uuid
+from abc import ABC, abstractmethod
+from functools import wraps
+
+import gradio as gr
+from gradio_client import Client
+
+
+__version__ = "0.1.0"
+__all__ = ["Environment", "load", "register_env"]
+
+
+class Environment(ABC):
+    @abstractmethod
+    def reset(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def step(self, *args, **kwargs):
+        pass
+
+
+class _RemoteEnvironment(Environment):
+    def __init__(self, env_id: str):
+        username, repo = env_id.split("/")
+        self.client = Client(f"https://{username}-{repo}.hf.space/")
+        self.session_id = self.client.predict(api_name="/init")
+
+    def reset(self, *args, **kwargs):
+        return self.client.predict(self.session_id, api_name="/reset", *args, **kwargs)
+
+    def step(self, *args, **kwargs):
+        return self.client.predict(self.session_id, api_name="/step", *args, **kwargs)
+
+
+def load(env_id: str) -> _RemoteEnvironment:
+    return _RemoteEnvironment(env_id)
+
+
+def bind_method_to_session(method, registry: dict):
+    sig = inspect.signature(method)
+    params = list(sig.parameters.values())
+
+    @wraps(method)
+    def wrapper(session_id: str, *args, **kwargs):
+        instance = registry.get(session_id)
+        if instance is None:
+            raise ValueError(f"Invalid session_id: {session_id}")
+        m = getattr(instance, method.__func__.__name__)
+        return m(*args, **kwargs)
+
+    # --- update __annotations__ ---
+    wrapper.__annotations__ = method.__annotations__.copy()
+    wrapper.__annotations__["session_id"] = str
+
+    # --- build signature ---
+    new_params = (
+        inspect.Parameter(
+            "session_id",
+            kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=str,
+        ),
+        *params,
+    )
+    wrapper.__signature__ = inspect.Signature(
+        parameters=new_params,
+        return_annotation=sig.return_annotation,
+    )
+
+    return wrapper
+
+
+def register_env(env_cls: type[Environment]) -> gr.Blocks:
+    """
+    Register an environment class with Gradio APIs.
+
+    Example:
+
+    ```python
+    from environments import register_env, Environment
+    import gradio as gr
+
+    class MyEnvironmentClass(Environment):
+        def reset(self) -> str:
+            return "Reset called!"
+
+        def step(self, action: str) -> str:
+            return f"Step called with action: {action}!"
+
+    with gr.Blocks() as demo:
+        register_env(MyEnvironmentClass)
+
+    demo.launch(mcp_server_name=True)
+    ```
+    """
+    sessions = {}
+
+    def init_env() -> str:
+        """
+        Initialize a new environment instance and return a session ID.
+        Returns:
+            A unique session ID for the new environment instance.
+        """
+        session_id = str(uuid.uuid4())
+        env = env_cls()
+        sessions[session_id] = env
+        return session_id
+
+    # Bind methods to session dict
+    reset_api = bind_method_to_session(env_cls().reset, sessions)
+    step_api = bind_method_to_session(env_cls().step, sessions)
+
+    # Create Gradio APIs
+    gr.api(
+        init_env,
+        api_name="init",
+        api_description="Initialize a new environment session",
+    )
+    gr.api(reset_api, api_name="reset")
+    gr.api(step_api, api_name="step")