hud-python 0.4.45__py3-none-any.whl → 0.5.1__py3-none-any.whl

hud/environment/tests/test_scenarios.py

@@ -0,0 +1,280 @@
+"""Tests for Environment scenario decorator."""
+
+from __future__ import annotations
+
+import pytest
+
+from hud.environment import Environment
+
+
+class TestScenarioDecorator:
+    """Tests for @env.scenario decorator."""
+
+    def test_scenario_registers_function(self) -> None:
+        """@env.scenario registers the function."""
+        env = Environment("test-env")
+
+        @env.scenario("greet")
+        async def greet_scenario(name: str):
+            yield f"Hello, {name}!"
+            yield 1.0
+
+        assert "greet" in env._scenarios
+
+    def test_scenario_creates_mcp_prompt(self) -> None:
+        """@env.scenario creates an MCP prompt."""
+        env = Environment("test-env")
+
+        @env.scenario("greet", description="Greeting scenario")
+        async def greet_scenario(name: str):
+            yield f"Hello, {name}!"
+            yield 1.0
+
+        # Check that prompt was registered via prompt manager
+        prompt_names = list(env._prompt_manager._prompts.keys())
+        assert "test-env:greet" in prompt_names
+
+    def test_scenario_creates_mcp_resource(self) -> None:
+        """@env.scenario creates an MCP resource."""
+        env = Environment("test-env")
+
+        @env.scenario("greet")
+        async def greet_scenario(name: str):
+            yield f"Hello, {name}!"
+            yield 1.0
+
+        # Check that resource was registered via resource manager
+        resource_uris = list(env._resource_manager._resources.keys())
+        assert "test-env:greet" in resource_uris
+
+    def test_scenario_extracts_arguments(self) -> None:
+        """@env.scenario extracts function arguments for prompt."""
+        env = Environment("test-env")
+
+        @env.scenario("checkout")
+        async def checkout_scenario(user_id: str, amount: int = 100):
+            yield f"Checkout for {user_id}: ${amount}"
+            yield 1.0
+
+        # Find the prompt
+        prompt = env._prompt_manager._prompts.get("test-env:checkout")
+        assert prompt is not None
+        assert prompt.arguments is not None
+
+        # Check arguments
+        arg_names = [arg.name for arg in prompt.arguments]
+        assert "user_id" in arg_names
+        assert "amount" in arg_names
+
+
+class TestScenarioExecution:
+    """Tests for scenario execution flow."""
+
+    @pytest.mark.asyncio
+    async def test_scenario_setup_phase(self) -> None:
+        """Scenario setup phase yields prompt."""
+        env = Environment("test-env")
+        setup_ran = False
+
+        @env.scenario("test")
+        async def test_scenario():
+            nonlocal setup_ran
+            setup_ran = True
+            yield "Test prompt"
+            yield 1.0
+
+        # Get the prompt handler
+        prompt = env._prompt_manager._prompts.get("test-env:test")
+        assert prompt is not None
+
+        # Run setup via prompt render (which calls fn) - no need for context
+        result = await prompt.render({})
+
+        assert setup_ran
+        # Result is list of PromptMessage
+        assert len(result) > 0
+        assert "Test prompt" in str(result[0].content)
+
+    @pytest.mark.asyncio
+    async def test_scenario_stores_session(self) -> None:
+        """Scenario stores generator in session for evaluate phase."""
+        env = Environment("test-env")
+
+        @env.scenario("test")
+        async def test_scenario():
+            yield "Test prompt"
+            yield 1.0
+
+        # Run setup via prompt - no need for context
+        prompt = env._prompt_manager._prompts.get("test-env:test")
+        assert prompt is not None
+        await prompt.render({})
+
+        # Check session was stored
+        assert "test" in env._scenario_latest
+
+    @pytest.mark.asyncio
+    async def test_scenario_full_flow(self) -> None:
+        """Scenario runs setup and evaluate phases correctly."""
+        env = Environment("test-env")
+        phases = []
+
+        @env.scenario("test")
+        async def test_scenario():
+            phases.append("setup")
+            yield "Test prompt"
+            phases.append("evaluate")
+            yield 0.95
+
+        # Setup phase - no context needed for prompt/resource
+        prompt = env._prompt_manager._prompts.get("test-env:test")
+        assert prompt is not None
+        await prompt.render({})
+        assert "setup" in phases
+        assert "evaluate" not in phases
+
+        # Evaluate phase
+        resource = env._resource_manager._resources.get("test-env:test")
+        assert resource is not None
+        await resource.read()
+        assert "evaluate" in phases
+
+
+class TestScenarioWithArgs:
+    """Tests for scenarios with arguments."""
+
+    @pytest.mark.asyncio
+    async def test_scenario_receives_args(self) -> None:
+        """Scenario receives arguments from prompt call."""
+        env = Environment("test-env")
+        received_args = {}
+
+        @env.scenario("checkout")
+        async def checkout_scenario(user_id: str, amount: int = 100):
+            received_args["user_id"] = user_id
+            received_args["amount"] = amount
+            yield f"Checkout {user_id}: ${amount}"
+            yield 1.0
+
+        prompt = env._prompt_manager._prompts.get("test-env:checkout")
+        assert prompt is not None
+        # No context needed for prompt render
+        await prompt.render({"user_id": "alice", "amount": 50})
+
+        assert received_args["user_id"] == "alice"
+        assert received_args["amount"] == 50
+
+
+class TestScenarioSubmit:
+    """Tests for scenario submit and answer flow."""
+
+    @pytest.mark.asyncio
+    async def test_submit_stores_answer(self) -> None:
+        """submit() stores answer for scenario."""
+        env = Environment("test-env")
+
+        @env.scenario("test")
+        async def test_scenario():
+            yield "What is 2+2?"
+            yield 1.0
+
+        # Run setup
+        prompt = env._prompt_manager._prompts.get("test-env:test")
+        assert prompt is not None
+        await prompt.render({})
+
+        # Submit answer
+        await env.submit("test", "4")
+
+        assert env._scenario_answers.get("test") == "4"
+
+    @pytest.mark.asyncio
+    async def test_scenario_receives_answer(self) -> None:
+        """Scenario receives submitted answer via yield."""
+        env = Environment("test-env")
+        received_answer = None
+
+        @env.scenario("qa")
+        async def qa_scenario():
+            nonlocal received_answer
+            answer = yield "What is 2+2?"
+            received_answer = answer
+            yield 1.0 if answer == "4" else 0.0
+
+        # Run setup
+        prompt = env._prompt_manager._prompts.get("test-env:qa")
+        assert prompt is not None
+        await prompt.render({})
+
+        # Submit answer
+        env._scenario_answers["qa"] = "4"
+
+        # Run evaluate
+        resource = env._resource_manager._resources.get("test-env:qa")
+        assert resource is not None
+        await resource.read()
+
+        assert received_answer == "4"
+
+    @pytest.mark.asyncio
+    async def test_scenario_evaluates_answer(self) -> None:
+        """Scenario evaluates answer and returns reward."""
+        env = Environment("test-env")
+
+        @env.scenario("grading")
+        async def grading_scenario():
+            answer = yield "What is the capital of France?"
+            yield 1.0 if "paris" in answer.lower() else 0.0
+
+        # Run setup
+        prompt = env._prompt_manager._prompts.get("test-env:grading")
+        assert prompt is not None
+        await prompt.render({})
+
+        # Submit correct answer
+        env._scenario_answers["grading"] = "Paris"
+
+        # Run evaluate
+        resource = env._resource_manager._resources.get("test-env:grading")
+        assert resource is not None
+        result = await resource.read()
+
+        import json
+
+        data = json.loads(result)
+        assert data["reward"] == 1.0
+
+
+class TestScenarioMeta:
+    """Tests for scenario _meta containing code."""
+
+    def test_scenario_captures_source_code(self) -> None:
+        """@env.scenario captures function source in meta."""
+        env = Environment("test-env")
+
+        @env.scenario("example")
+        async def example_scenario(x: int):
+            yield f"Process {x}"
+            yield 1.0
+
+        prompt = env._prompt_manager._prompts.get("test-env:example")
+        assert prompt is not None
+        assert prompt.meta is not None
+        assert "code" in prompt.meta
+        assert "async def example_scenario" in prompt.meta["code"]
+        assert "yield" in prompt.meta["code"]
+
+    def test_scenario_meta_on_resource(self) -> None:
+        """Resource also has source code in meta."""
+        env = Environment("test-env")
+
+        @env.scenario("example")
+        async def example_scenario():
+            yield "Test"
+            yield 1.0
+
+        resource = env._resource_manager._resources.get("test-env:example")
+        assert resource is not None
+        assert resource.meta is not None
+        assert "code" in resource.meta
+        assert "async def example_scenario" in resource.meta["code"]

hud/environment/tests/test_tools.py

@@ -0,0 +1,208 @@
+"""Tests for @env.tool() decorator and tool operations."""
+
+from __future__ import annotations
+
+import pytest
+
+from hud.environment import Environment
+
+
+class TestToolDecorator:
+    """Tests for @env.tool() decorator."""
+
+    def test_tool_registers_function(self) -> None:
+        """@env.tool registers the function in tool manager."""
+        env = Environment("test-env")
+
+        @env.tool()
+        def add(a: int, b: int) -> int:
+            """Add two numbers."""
+            return a + b
+
+        # Check tool was registered
+        tool_names = list(env._tool_manager._tools.keys())
+        assert "add" in tool_names
+
+    def test_tool_with_custom_name(self) -> None:
+        """@env.tool(name=...) uses custom name."""
+        env = Environment("test-env")
+
+        @env.tool(name="custom_add")
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        tool_names = list(env._tool_manager._tools.keys())
+        assert "custom_add" in tool_names
+        assert "add" not in tool_names
+
+    def test_tool_preserves_docstring(self) -> None:
+        """@env.tool preserves function docstring as description."""
+        env = Environment("test-env")
+
+        @env.tool()
+        def greet(name: str) -> str:
+            """Greet someone by name."""
+            return f"Hello, {name}!"
+
+        tool = env._tool_manager._tools.get("greet")
+        assert tool is not None
+        assert "Greet someone by name" in (tool.description or "")
+
+    def test_tool_async_function(self) -> None:
+        """@env.tool works with async functions."""
+        env = Environment("test-env")
+
+        @env.tool()
+        async def fetch_data(url: str) -> str:
+            """Fetch data from URL."""
+            return f"Data from {url}"
+
+        tool_names = list(env._tool_manager._tools.keys())
+        assert "fetch_data" in tool_names
+
+    def test_tool_returns_function(self) -> None:
+        """@env.tool returns the original function."""
+        env = Environment("test-env")
+
+        @env.tool()
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        # Should be able to call it directly
+        assert add(2, 3) == 5
+
+
+class TestListTools:
+    """Tests for list_tools and as_tools."""
+
+    @pytest.mark.asyncio
+    async def test_as_tools_returns_registered_tools(self) -> None:
+        """as_tools returns list of registered MCP tools."""
+        env = Environment("test-env")
+
+        @env.tool()
+        def tool1() -> str:
+            return "1"
+
+        @env.tool()
+        def tool2() -> str:
+            return "2"
+
+        async with env:
+            tools = env.as_tools()
+            tool_names = [t.name for t in tools]
+            assert "tool1" in tool_names
+            assert "tool2" in tool_names
+
+    @pytest.mark.asyncio
+    async def test_as_tools_empty_when_no_tools(self) -> None:
+        """as_tools returns empty list when no tools registered."""
+        env = Environment("test-env")
+        async with env:
+            tools = env.as_tools()
+            # May have built-in _hud_submit tool
+            user_tools = [t for t in tools if not t.name.startswith("_")]
+            assert len(user_tools) == 0
+
+
+class TestCallTool:
+    """Tests for call_tool method."""
+
+    @pytest.mark.asyncio
+    async def test_call_tool_executes_function(self) -> None:
+        """call_tool executes registered tool function."""
+        env = Environment("test-env")
+        executed = []
+
+        @env.tool()
+        def greet(name: str) -> str:
+            executed.append(name)
+            return f"Hello, {name}!"
+
+        async with env:
+            result = await env.call_tool("greet", name="Alice")
+
+        assert executed == ["Alice"]
+        assert result is not None
+
+    @pytest.mark.asyncio
+    async def test_call_tool_async_function(self) -> None:
+        """call_tool works with async tool functions."""
+        env = Environment("test-env")
+
+        @env.tool()
+        async def async_greet(name: str) -> str:
+            return f"Hello, {name}!"
+
+        async with env:
+            result = await env.call_tool("async_greet", name="Bob")
+
+        assert result is not None
+
+    @pytest.mark.asyncio
+    async def test_call_tool_not_found(self) -> None:
+        """call_tool raises for unknown tool."""
+        env = Environment("test-env")
+
+        async with env:
+            with pytest.raises(ValueError, match="Tool not found"):
+                await env.call_tool("nonexistent")
+
+
+class TestMockMode:
+    """Tests for mock mode."""
+
+    def test_mock_mode_default_false(self) -> None:
+        """Mock mode is False by default."""
+        env = Environment("test-env")
+        assert env._mock_mode is False
+        assert env.is_mock is False
+
+    def test_mock_enables_mock_mode(self) -> None:
+        """mock() enables mock mode."""
+        env = Environment("test-env")
+        env.mock()
+        assert env._mock_mode is True
+        assert env.is_mock is True
+
+    def test_unmock_disables_mock_mode(self) -> None:
+        """unmock() disables mock mode."""
+        env = Environment("test-env")
+        env.mock()
+        env.unmock()
+        assert env._mock_mode is False
+
+    def test_mock_returns_self_for_chaining(self) -> None:
+        """mock() returns self for chaining."""
+        env = Environment("test-env")
+        result = env.mock()
+        assert result is env
+
+    def test_mock_tool_sets_custom_output(self) -> None:
+        """mock_tool() sets custom output for a tool."""
+        env = Environment("test-env")
+        env.mock_tool("navigate", "Custom result")
+        assert env._mock_outputs["navigate"] == "Custom result"
+
+    @pytest.mark.asyncio
+    async def test_mock_mode_returns_mock_response(self) -> None:
+        """Mock mode returns mock response instead of executing tool."""
+        env = Environment("test-env")
+        call_count = 0
+
+        @env.tool()
+        def real_tool() -> str:
+            nonlocal call_count
+            call_count += 1
+            return "real result"
+
+        env.mock()
+        env.mock_tool("real_tool", "mocked result")
+
+        async with env:
+            result = await env.call_tool("real_tool")
+
+        # Tool should not be called in mock mode
+        assert call_count == 0
+        # Should get the mock result
+        assert result is not None

hud/environment/types.py

@@ -0,0 +1,23 @@
+"""Environment types for configuration and tracing."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+__all__ = ["EnvConfig"]
+
+
+class EnvConfig(BaseModel):
+    """Environment configuration for Tasks.
+
+    Specifies which hub to connect to and optional tool filtering.
+
+    Attributes:
+        name: Hub name to connect via connect_hub() (e.g., "browser", "sheets")
+        include: Optional whitelist of tool names to include
+        exclude: Optional blacklist of tool names to exclude
+    """
+
+    name: str = Field(description="Hub name to connect to")
+    include: list[str] | None = Field(default=None, description="Whitelist of tool names")
+    exclude: list[str] | None = Field(default=None, description="Blacklist of tool names")

hud/environment/utils/__init__.py

@@ -0,0 +1,35 @@
+"""Environment utilities."""
+
+from hud.environment.utils.formats import (
+    ToolFormat,
+    format_result,
+    parse_tool_call,
+    parse_tool_calls,
+    result_to_string,
+)
+from hud.environment.utils.schema import (
+    ensure_strict_schema,
+    json_type_to_python,
+    schema_to_pydantic,
+)
+from hud.environment.utils.tool_wrappers import (
+    create_async_tool_fn,
+    create_sync_tool_fn,
+    create_tool_fns,
+    stringify_result,
+)
+
+__all__ = [
+    "ToolFormat",
+    "create_async_tool_fn",
+    "create_sync_tool_fn",
+    "create_tool_fns",
+    "ensure_strict_schema",
+    "format_result",
+    "json_type_to_python",
+    "parse_tool_call",
+    "parse_tool_calls",
+    "result_to_string",
+    "schema_to_pydantic",
+    "stringify_result",
+]