hud-python 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

hud/agent/tests/test_base.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+from typing import Any
+
+import pytest
+
+from hud.agent.base import Agent
+from hud.adapters import Adapter
+from hud.adapters.common.types import ClickAction, Point
+from hud.utils.common import Observation
+
+
+class ConcreteAgent(Agent[Any, dict[str, Any]]):
+    """Concrete implementation of Agent for testing."""
+
+    def __init__(self, client: Any = None, adapter: Adapter | None = None):
+        super().__init__(client, adapter)
+        self.mock_responses = []
+        self.call_count = 0
+
+    async def fetch_response(self, observation: Observation) -> tuple[list[dict[str, Any]], bool]:
+        """Mock implementation that returns predefined responses."""
+        if self.call_count < len(self.mock_responses):
+            response = self.mock_responses[self.call_count]
+            self.call_count += 1
+            return response
+        return [], True
+
+
+class TestAgentBase:
+    """Test the base Agent class."""
+
+    @pytest.fixture
+    def mock_client(self):
+        """Mock client for testing."""
+        return MagicMock()
+
+    @pytest.fixture
+    def mock_adapter(self):
+        """Mock adapter for testing."""
+        adapter = MagicMock(spec=Adapter)
+        adapter.rescale.return_value = "rescaled_screenshot"
+        adapter.adapt_list.return_value = [ClickAction(point=Point(x=100, y=200))]
+        return adapter
+
+    @pytest.fixture
+    def agent_with_adapter(self, mock_client, mock_adapter):
+        """Agent with both client and adapter."""
+        return ConcreteAgent(client=mock_client, adapter=mock_adapter)
+
+    @pytest.fixture
+    def agent_without_adapter(self, mock_client):
+        """Agent with client but no adapter."""
+        return ConcreteAgent(client=mock_client, adapter=None)
+
+    def test_init_with_client_and_adapter(self, mock_client, mock_adapter):
+        """Test agent initialization with client and adapter."""
+        agent = ConcreteAgent(client=mock_client, adapter=mock_adapter)
+        assert agent.client == mock_client
+        assert agent.adapter == mock_adapter
+
+    def test_init_with_none_values(self):
+        """Test agent initialization with None values."""
+        agent = ConcreteAgent(client=None, adapter=None)
+        assert agent.client is None
+        assert agent.adapter is None
+
+    def test_preprocess_without_adapter(self, agent_without_adapter):
+        """Test preprocess when no adapter is available."""
+        observation = Observation(text="test", screenshot="screenshot_data")
+        result = agent_without_adapter.preprocess(observation)
+
+        # Should return original observation unchanged
+        assert result == observation
+        assert result.text == "test"
+        assert result.screenshot == "screenshot_data"
+
+    def test_preprocess_without_screenshot(self, agent_with_adapter):
+        """Test preprocess when no screenshot is available."""
+        observation = Observation(text="test", screenshot=None)
+        result = agent_with_adapter.preprocess(observation)
+
+        # Should return original observation unchanged
+        assert result == observation
+        assert result.text == "test"
+        assert result.screenshot is None
+
+    def test_preprocess_with_adapter_and_screenshot(self, agent_with_adapter, mock_adapter):
+        """Test preprocess with adapter and screenshot (covers missing lines 48-55)."""
+        observation = Observation(text="test", screenshot="original_screenshot")
+        result = agent_with_adapter.preprocess(observation)
+
+        # Should create new observation with rescaled screenshot
+        mock_adapter.rescale.assert_called_once_with("original_screenshot")
+        assert result.text == "test"
+        assert result.screenshot == "rescaled_screenshot"
+        # Should be a new object, not the original
+        assert result is not observation
+
+    def test_postprocess_without_adapter(self, agent_without_adapter):
+        """Test postprocess when no adapter is available (covers missing lines 82-85)."""
+        actions = [{"type": "click", "x": 100, "y": 200}]
+
+        with pytest.raises(ValueError, match="Cannot postprocess actions without an adapter"):
+            agent_without_adapter.postprocess(actions)
+
+    def test_postprocess_with_adapter(self, agent_with_adapter, mock_adapter):
+        """Test postprocess with adapter."""
+        actions = [{"type": "click", "x": 100, "y": 200}]
+        result = agent_with_adapter.postprocess(actions)
+
+        mock_adapter.adapt_list.assert_called_once_with(actions)
+        assert len(result) == 1
+        assert isinstance(result[0], ClickAction)
+
+    @pytest.mark.asyncio
+    async def test_predict_without_verbose(self, agent_with_adapter):
+        """Test predict method without verbose logging."""
+        observation = Observation(text="test", screenshot="screenshot")
+        agent_with_adapter.mock_responses = [([{"type": "click", "x": 100, "y": 200}], False)]
+
+        actions, done = await agent_with_adapter.predict(observation, verbose=False)
+
+        assert len(actions) == 1
+        assert isinstance(actions[0], ClickAction)
+        assert done is False
+
+    @pytest.mark.asyncio
+    @patch("hud.agent.base.logger")
+    async def test_predict_with_verbose_logging(self, mock_logger, agent_with_adapter):
+        """Test predict method with verbose logging (covers missing lines 100-116)."""
+        observation = Observation(text="test", screenshot="screenshot")
+        agent_with_adapter.mock_responses = [([{"type": "click", "x": 100, "y": 200}], True)]
+
+        actions, done = await agent_with_adapter.predict(observation, verbose=True)
+
+        # Verify verbose logging was called
+        mock_logger.info.assert_any_call("Predicting action...")
+        mock_logger.info.assert_any_call("Raw action: %s", [{"type": "click", "x": 100, "y": 200}])
+
+        assert len(actions) == 1
+        assert isinstance(actions[0], ClickAction)
+        assert done is True
+
+    @pytest.mark.asyncio
+    async def test_predict_without_adapter_returns_raw_actions(self, agent_without_adapter):
+        """Test predict without adapter returns raw actions."""
+        observation = Observation(text="test", screenshot=None)
+        raw_actions = [{"type": "click", "x": 100, "y": 200}]
+        agent_without_adapter.mock_responses = [(raw_actions, True)]
+
+        actions, done = await agent_without_adapter.predict(observation, verbose=False)
+
+        # Should return raw actions, not processed ones
+        assert actions == raw_actions
+        assert done is True
+
+    @pytest.mark.asyncio
+    async def test_predict_with_empty_actions(self, agent_with_adapter):
+        """Test predict when fetch_response returns empty actions."""
+        observation = Observation(text="test", screenshot="screenshot")
+        agent_with_adapter.mock_responses = [([], True)]
+
+        actions, done = await agent_with_adapter.predict(observation, verbose=False)
+
+        # Should return empty actions without calling adapter
+        assert actions == []
+        assert done is True
+
+    @pytest.mark.asyncio
+    async def test_predict_full_pipeline(self, agent_with_adapter, mock_adapter):
+        """Test the complete predict pipeline with all stages."""
+        # Set up observation with screenshot that will be rescaled
+        observation = Observation(text="test input", screenshot="original_screenshot")
+        raw_actions = [{"type": "click", "x": 150, "y": 250}]
+        agent_with_adapter.mock_responses = [(raw_actions, False)]
+
+        actions, done = await agent_with_adapter.predict(observation, verbose=True)
+
+        # Verify all stages were called
+        # Stage 1: Preprocessing
+        mock_adapter.rescale.assert_called_once_with("original_screenshot")
+
+        # Stage 3: Postprocessing
+        mock_adapter.adapt_list.assert_called_once_with(raw_actions)
+
+        assert len(actions) == 1
+        assert isinstance(actions[0], ClickAction)
+        assert done is False
+
+    @pytest.mark.asyncio
+    async def test_predict_integration_without_screenshot(self, agent_with_adapter):
+        """Test predict integration when observation has no screenshot."""
+        observation = Observation(text="test input", screenshot=None)
+        raw_actions = [{"type": "response", "text": "Task completed"}]
+        agent_with_adapter.mock_responses = [(raw_actions, True)]
+
+        actions, done = await agent_with_adapter.predict(observation, verbose=False)
+
+        assert len(actions) == 1
+        assert done is True

hud/env/docker_client.py

@@ -26,6 +26,8 @@ STATUS_MESSAGES = {
     EnvironmentStatus.COMPLETED.value: "completed",
 }
 
+PACKAGE_NAME = "hud_controller"
+
 
 class InvokeError(Exception):
     """
@@ -63,20 +65,12 @@ class DockerClient(Client):
     _last_update_time: int = 0
     _last_file_mtimes: dict[str, float] = {}  # noqa: RUF012 - Not recognized as Pydantic model
     _source_path: Path | None = None
-    _package_name: str | None = None
 
     @property
     def source_path(self) -> Path | None:
         """Get the source path."""
         return self._source_path
 
-    @property
-    def package_name(self) -> str:
-        """Get the package name."""
-        if not self._package_name:
-            raise ValueError("Package name not set")
-        return self._package_name
-
     def set_source_path(self, source_path: Path) -> None:
         """
         Set the source path for this environment controller.
@@ -102,21 +96,37 @@ class DockerClient(Client):
         if not pyproject_path.exists():
             raise FileNotFoundError(f"pyproject.toml not found in {source_path}")
 
+        # validate package name
         pyproject_data = toml.load(pyproject_path)
-        self._package_name = pyproject_data.get("project", {}).get("name")
-        if not self._package_name:
+        package_name = pyproject_data.get("project", {}).get("name")
+        if not package_name:
             raise ValueError("Could not find package name in pyproject.toml")
+        if package_name != PACKAGE_NAME:
+            raise ValueError(f"Package name in pyproject.toml must be {PACKAGE_NAME}")
 
         self._source_path = source_path
 
+        # set current mtimes
+        self._last_file_mtimes = self._get_all_file_mtimes()
+
+    @classmethod
+    @abc.abstractmethod
+    async def build_image(cls, build_context: Path) -> tuple[str, dict[str, Any]]:
+        """
+        Build an image from a build context.
+
+        Returns:
+            tuple[str, dict[str, Any]]: The image tag and build output
+        """
+
     @classmethod
     @abc.abstractmethod
-    async def create(cls, dockerfile: str) -> DockerClient:
+    async def create(cls, image: str) -> DockerClient:
         """
-        Creates an environment client from a dockerfile.
+        Creates an environment client from an image.
 
         Args:
-            dockerfile: The dockerfile content to build the environment
+            image: The image to build the environment from
 
         Returns:
             EnvClient: An instance of the environment client
@@ -194,8 +204,8 @@ class DockerClient(Client):
 
         # Create tar archive of the source code and send it to the container
         tar_bytes = directory_to_tar_bytes(self._source_path)
-        await self.execute(["mkdir", "-p", "/root/controller"], timeout=5)
-        await self.put_archive("/root/controller", tar_bytes)
+        await self.execute(["mkdir", "-p", "/controller"], timeout=5)
+        await self.put_archive("/controller", tar_bytes)
 
         # Check if pyproject.toml exists and parse it
         pyproject_path = self._source_path / "pyproject.toml"
@@ -213,9 +223,9 @@ class DockerClient(Client):
             self._package_name = pyproject_data.get("project", {}).get("name")
             if not self._package_name:
                 raise ValueError("Could not find package name in pyproject.toml")
-            logger.info("Installing %s in /root/controller", self._package_name)
+            logger.info("Installing %s in /controller", self._package_name)
             result = await self.execute(
-                ["bash", "-c", "cd /root/controller && pip install -e . --break-system-packages"],
+                ["bash", "-c", "cd /controller && pip install -e . --break-system-packages"],
                 timeout=60,
             )
             if result["stdout"]:
@@ -262,7 +272,7 @@ class DockerClient(Client):
         # generate a random uuid as a divider
         divider = str(uuid.uuid4())
 
-        template = invoke_template(config, self.package_name, divider)
+        template = invoke_template(config, PACKAGE_NAME, divider)
         logger.debug("Invoking template: %s", template)
 
         result = await self.execute(["python3", "-c", template])

hud/env/environment.py

@@ -74,14 +74,17 @@ class Environment(BaseModel):
             config: The configuration to use for the setup
         """
         if isinstance(self.client, RemoteClient):
+            await self.get_urls()
             await self._invoke_all(create_remote_config(self, config, REMOTE_SETUP))
         else:
             if config is not None:
                 await self._invoke_all(config)
-            elif self.task and self.task.config is not None:
-                await self._invoke_all(self.task.config)
+            elif self.task and self.task.setup is not None:
+                await self._invoke_all(self.task.setup)
             else:
-                raise ValueError("No config or task provided for local environment")
+                raise ValueError(
+                    "No config, task or task setup function provided for local environment"
+                )
 
     async def evaluate(self, config: FunctionConfigs | None = None) -> Any:
         """
@@ -98,8 +101,8 @@ class Environment(BaseModel):
         else:
             if config is not None:
                 results = await self._invoke_all(config)
-            elif self.task and self.task.config is not None:
-                results = await self._invoke_all(self.task.config)
+            elif self.task and self.task.evaluate is not None:
+                results = await self._invoke_all(self.task.evaluate)
             else:
                 raise ValueError("No config or task provided for local environment")
         if len(results) == 1:
@@ -144,8 +147,7 @@ class Environment(BaseModel):
         args = [[action.model_dump() for action in actions]]
 
         # TODO: Move this into the server side
-        if self._maybe_store_response(actions):
-            return Observation(text=self.final_response), 0, False, {}
+        self._maybe_store_response(actions)
 
         result, stdout, stderr = await self.client.invoke(
             FunctionConfig(function="step", args=args)
@@ -211,20 +213,20 @@ class Environment(BaseModel):
             agent: The agent to run
         """
         if verbose:
-            logger.info("[HUD] Running agent in environment...")
+            logger.info("Running agent in environment...")
         obs, _ = await self.reset()
         for i in range(max_steps):
-            action, done = await agent.predict(obs)
+            action, done = await agent.predict(obs, verbose=verbose)
             if verbose:
-                logger.info("[HUD] Step %d: Action: %s", i, action)
+                logger.info("Step %d: Action: %s", i, action)
             obs, reward, terminated, info = await self.step(action)
             if verbose:
-                logger.info("[HUD] Step %d: Observation: %s", i, obs)
+                logger.info("Step %d: Observation: %s", i, obs)
             if done or terminated:
                 break
         result = await self.evaluate()
         if verbose:
-            logger.info("[HUD] Evaluation result: %s", result)
+            logger.info("Evaluation result: %s", result)
         return result
 
 
@@ -348,7 +350,11 @@ def create_remote_config(
             if not isinstance(expanded_configs[0].args, list):
                 expanded_configs[0].args = [expanded_configs[0].args]
             expanded_configs[0].args.append(env.final_response)  # for remote responses
-        return [FunctionConfig(function=function, args=expanded_configs)]
+        return [
+            FunctionConfig(
+                function=function, args=expanded_configs, metadata={"task": task.model_dump()}
+            )
+        ]
 
     # Case 3: Check for task.config
     if hasattr(task, "config") and task.config:
@@ -363,17 +369,27 @@ def create_remote_config(
             if not isinstance(final_args["args"], list):
                 final_args["args"] = [final_args["args"]]
             final_args["args"].append(env.final_response)
-        return [FunctionConfig(function=function, args=[final_args])]
+        return [
+            FunctionConfig(
+                function=function, args=[final_args], metadata={"task": task.model_dump()}
+            )
+        ]
 
     # Case 4: Use task.id
     if task.id:
         args_list = [task.id]
         if env and env.final_response:
             args_list.append(env.final_response)  # Append final response
-        return [FunctionConfig(function=f"{REMOTE_FUNCTION_PREFIX}{function}", args=args_list)]
+        return [
+            FunctionConfig(
+                function=f"{REMOTE_FUNCTION_PREFIX}{function}",
+                args=args_list,
+                metadata={"task": task.model_dump()},
+            )
+        ]
 
     # Case 5: No valid configuration found
     args_list = []
     if env and env.final_response:
         args_list.append(env.final_response)
-    return [FunctionConfig(function=function, args=args_list)]
+    return [FunctionConfig(function=function, args=args_list, metadata={"task": task.model_dump()})]

hud/env/local_docker_client.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import asyncio
 import io
 import logging
-import tarfile
-import tempfile
+import textwrap
+import time
 import uuid
 from typing import TYPE_CHECKING, Any
 
@@ -12,12 +13,15 @@ from aiohttp import ClientTimeout
 
 from hud.env.docker_client import DockerClient, EnvironmentStatus
 from hud.utils import ExecuteResult
+from hud.utils.common import directory_to_tar_bytes
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from aiodocker.containers import DockerContainer
     from aiodocker.stream import Stream
 
-logger = logging.getLogger("hud.env.docker_env_client")
+logger = logging.getLogger(__name__)
 
 
 class LocalDockerClient(DockerClient):
@@ -26,17 +30,9 @@ class LocalDockerClient(DockerClient):
     """
 
     @classmethod
-    async def create(
-        cls, dockerfile: str, ports: list[int] | None = None
-    ) -> tuple[LocalDockerClient, dict[str, Any]]:
+    async def build_image(cls, build_context: Path) -> tuple[str, dict[str, Any]]:
         """
-        Creates a Docker environment client from a dockerfile.
-
-        Args:
-            dockerfile: The dockerfile content to build the Docker image
-
-        Returns:
-            DockerClient: An instance of the Docker environment client
+        Build an image from a build context.
         """
         # Create a unique image tag
         image_tag = f"hud-env-{uuid.uuid4().hex[:8]}"
@@ -44,32 +40,19 @@ class LocalDockerClient(DockerClient):
         # Initialize Docker client
         docker_client = aiodocker.Docker()
 
-        # Create fileobj for the Dockerfile
-        dockerfile_fileobj = io.BytesIO(dockerfile.encode("utf-8"))
-
-        if ports is None:
-            ports = []
-
-        # Create a tar file from the dockerfile
-        with tempfile.NamedTemporaryFile() as f:
-            with tarfile.open(mode="w:gz", fileobj=f) as t:
-                dfinfo = tarfile.TarInfo("Dockerfile")
-                dfinfo.size = len(dockerfile_fileobj.getvalue())
-                dockerfile_fileobj.seek(0)
-                t.addfile(dfinfo, dockerfile_fileobj)
-
-            # Reset the file pointer to the beginning of the file
-            f.seek(0)
-
-            # Build the image
-            build_stream = await docker_client.images.build(
-                fileobj=f,
-                encoding="gzip",
-                tag=image_tag,
-                rm=True,
-                pull=True,
-                forcerm=True,
-            )
+        # Create a tar file from the path
+        tar_bytes = directory_to_tar_bytes(build_context)
+        logger.info("generated tar file with size: %d KB", len(tar_bytes) // 1024)
+
+        # Build the image
+        build_stream = await docker_client.images.build(
+            fileobj=io.BytesIO(tar_bytes),
+            encoding="gzip",
+            tag=image_tag,
+            rm=True,
+            pull=True,
+            forcerm=True,
+        )
 
         # Print build output
         output = ""
@@ -78,23 +61,63 @@ class LocalDockerClient(DockerClient):
                 logger.info(chunk["stream"])
                 output += chunk["stream"]
 
+        return image_tag, {"build_output": output}
+
+    @classmethod
+    async def create(
+        cls,
+        image: str,
+    ) -> LocalDockerClient:
+        """
+        Creates a Docker environment client from a image.
+
+        Args:
+            image: The image to build the Docker image
+
+        Returns:
+            DockerClient: An instance of the Docker environment client
+        """
+
+        # Initialize Docker client
+        docker_client = aiodocker.Docker()
+
         # Create and start the container
         container_config = {
-            "Image": image_tag,
+            "Image": image,
             "Tty": True,
             "OpenStdin": True,
             "Cmd": None,
             "HostConfig": {
                 "PublishAllPorts": True,
             },
-            "ExposedPorts": {f"{port}/tcp": {} for port in ports},
         }
 
         container = await docker_client.containers.create(config=container_config)
         await container.start()
 
+        inspection = await container.show()
+        if health_check_config := inspection["Config"].get("Healthcheck"):
+            # Using the interval as spinup deadline is a bit implicit - could
+            # consider adding explicitly to API if there's demand
+            window_usecs = health_check_config.get("Interval", int(30 * 1e9))
+            window_secs = window_usecs // 1_000_000
+
+            deadline = time.monotonic() + window_secs
+            logger.debug("Waiting for container %s to become healthy", container.id)
+            while True:
+                state = (await container.show())["State"]
+                if state.get("Health", {}).get("Status") == "healthy":
+                    break
+                if state.get("Status") in {"exited", "dead"}:
+                    raise RuntimeError("Container crashed before becoming healthy")
+                now = time.monotonic()
+                if now > deadline:
+                    raise TimeoutError(f"{container.id} not healthy after {window_secs}s")
+                await asyncio.sleep(1)
+            logger.debug("Container %s is healthy", container.id)
+
         # Return the controller instance
-        return cls(docker_client, container.id), {"build_output": output}
+        return cls(docker_client, container.id)
 
     def __init__(self, docker_conn: aiodocker.Docker, container_id: str) -> None:
         """
@@ -190,6 +213,24 @@ class LocalDockerClient(DockerClient):
             elif message.stream == 2:  # stderr
                 stderr_data.extend(message.data)
 
+        if "No module named 'hud_controller'" in stderr_data.decode():
+            if self._source_path is None:
+                message = textwrap.dedent("""\
+                Your environment is not set up correctly.
+                You are using a prebuilt image, so please ensure the following:
+                1. Your image cannot be a generic python image, it must contain a python package
+                   called hud_controller.
+                """)
+            else:
+                message = textwrap.dedent("""\
+                Your environment is not set up correctly.
+                You are using a local controller, so please ensure the following:
+                1. Your package name is hud_controller
+                2. You installed the package in the Dockerfile.
+                3. The package is visible from the global python environment (no venv, conda, or uv)
+                """)
+            logger.error(message)
+
         return ExecuteResult(
             stdout=bytes(stdout_data),
             stderr=bytes(stderr_data),

hud/env/remote_client.py

@@ -83,9 +83,7 @@ class RemoteClient(Client):
         build_data = response.get("metadata", {})
 
         if response.get("readme"):
-            logger.info(
-                "[HUD] %s gym created, see how to use it at %s", gym_id, response.get("readme")
-            )
+            logger.info("Gym created, see how to use it at %s", response.get("readme"))
 
         return controller, build_data