hud-python 0.4.51__py3-none-any.whl → 0.4.53__py3-none-any.whl

hud/__init__.py

@@ -5,10 +5,22 @@ tools for building, evaluating, and training AI agents.
 
 from __future__ import annotations
 
-from .telemetry import Trace, clear_trace, create_job, get_trace, instrument, job, trace
+from .telemetry import (
+    Trace,
+    async_job,
+    async_trace,
+    clear_trace,
+    create_job,
+    get_trace,
+    instrument,
+    job,
+    trace,
+)
 
 __all__ = [
     "Trace",
+    "async_job",
+    "async_trace",
     "clear_trace",
     "create_job",
     "get_trace",

hud/agents/base.py

@@ -55,6 +55,7 @@ class MCPAgent(ABC):
         # Filtering
         allowed_tools: list[str] | None = None,
         disallowed_tools: list[str] | None = None,
+        response_tool_name: str | None = None,
         # Messages
         system_prompt: str = GLOBAL_SYSTEM_PROMPT,
         append_setup_output: bool = True,
@@ -74,6 +75,7 @@ class MCPAgent(ABC):
                 that provides `mcp_config`.
             allowed_tools: Names of tools to allow (None means allow all).
             disallowed_tools: Names of tools to always exclude.
+            response_tool_name: Name of the tool to use for response.
             system_prompt: System prompt to seed the conversation.
             append_setup_output: Whether to append setup tool output to the
                 first turn's messages.
@@ -108,7 +110,7 @@ class MCPAgent(ABC):
 
         # Initialize these here so methods can be called before initialize()
         self._tool_map: dict[str, types.Tool] = {}  # Simplified: just name to tool
-        self.response_tool_name = None
+        self.response_tool_name = response_tool_name
 
         # Trace
         self._auto_trace = auto_trace
@@ -135,7 +137,11 @@ class MCPAgent(ABC):
                 "No MCPClient. Please provide one when initializing the agent or pass a Task with mcp_config."  # noqa: E501
             )
 
-        await self._setup_config(self.mcp_client.mcp_config)
+        try:
+            client_cfg = getattr(self.mcp_client, "mcp_config", None)
+        except Exception:
+            client_cfg = None
+        await self._setup_config(client_cfg)
 
         # Initialize client if needed
         try:
@@ -168,6 +174,8 @@ class MCPAgent(ABC):
                     self.disallowed_tools.extend(task.agent_config["disallowed_tools"])
                 else:  # If disallowed_tools is None, we overwrite it
                     self.disallowed_tools = task.agent_config["disallowed_tools"]
+            if "response_tool_name" in task.agent_config:
+                self.response_tool_name = task.agent_config["response_tool_name"]
 
         all_tools = await self.mcp_client.list_tools()
         self._available_tools = []
@@ -614,8 +622,11 @@ class MCPAgent(ABC):
             except Exception as e:
                 self.console.error_log(f"Response lifecycle tool failed: {e}")
 
-    async def _setup_config(self, mcp_config: dict[str, dict[str, Any]]) -> None:
+    async def _setup_config(self, mcp_config: dict[str, dict[str, Any]] | None) -> None:
         """Inject metadata into the metadata of the initialize request."""
+        if not isinstance(mcp_config, dict):
+            return
+
         if self.metadata:
             patch_mcp_config(
                 mcp_config,

hud/agents/lite_llm.py

@@ -47,7 +47,7 @@ class LiteAgent(GenericOpenAIChatAgent):
             **agent_kwargs,
         )
 
-    def get_tool_schemas(self) -> list[dict]:
+    def get_tool_schemas(self) -> list[Any]:
         # Prefer LiteLLM's stricter transformer (handles Bedrock & friends)
         if transform_mcp_tool_to_openai_tool is not None:
             return [

hud/agents/openai_chat_generic.py

@@ -20,6 +20,7 @@ import logging
 from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 import mcp.types as types
+from openai import AsyncOpenAI
 
 from hud import instrument
 from hud.types import AgentResponse, MCPToolCall, MCPToolResult
@@ -28,7 +29,6 @@ from hud.utils.hud_console import HUDConsole
 from .base import MCPAgent
 
 if TYPE_CHECKING:
-    from openai import AsyncOpenAI
     from openai.types.chat import ChatCompletionToolParam
 
 logger = logging.getLogger(__name__)
@@ -42,14 +42,26 @@ class GenericOpenAIChatAgent(MCPAgent):
     def __init__(
         self,
         *,
-        openai_client: AsyncOpenAI | None,
+        openai_client: AsyncOpenAI | None = None,
+        api_key: str | None = None,
+        base_url: str | None = None,
         model_name: str = "gpt-4o-mini",
         completion_kwargs: dict[str, Any] | None = None,
         **agent_kwargs: Any,
     ) -> None:
         # Accept base-agent settings via **agent_kwargs (e.g., mcp_client, system_prompt, etc.)
         super().__init__(**agent_kwargs)
-        self.oai = openai_client
+
+        # Handle client creation - support both patterns
+        if openai_client is not None:
+            # Use provided client (backward compatibility)
+            self.oai = openai_client
+        elif api_key is not None or base_url is not None:
+            # Create client from config (new pattern, consistent with other agents)
+            self.oai = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        else:
+            raise ValueError("Either openai_client or (api_key and base_url) must be provided")
+
         self.model_name = model_name
         self.completion_kwargs: dict[str, Any] = completion_kwargs or {}
         self.mcp_schemas = []

hud/agents/tests/test_base.py

@@ -94,7 +94,7 @@ class TestBaseMCPAgent:
 
         assert agent.mcp_client is not None
         assert agent.allowed_tools is None
-        assert agent.disallowed_tools == []
+        assert agent.disallowed_tools is None
         assert agent.initial_screenshot is True
         assert agent.system_prompt is not None  # Default system prompt is set
 
@@ -241,6 +241,13 @@ class TestBaseMCPAgent:
         assert "tool2" not in tool_names  # Not in allowed list
         assert "tool3" not in tool_names  # In disallowed list
 
+        # Make sure tool schemas are correct
+        schemas = agent.get_tool_schemas()
+        assert len(schemas) == 1
+        assert schemas[0]["name"] == "tool1"
+        assert schemas[0]["description"] == "Tool 1"
+        assert schemas[0]["parameters"] == {"type": "object"}
+
     @pytest.mark.asyncio
     async def test_call_tool_success(self):
         """Test successful tool call."""
@@ -334,7 +341,7 @@ class TestBaseMCPAgent:
         schemas = agent.get_tool_schemas()
 
         # Should include non-lifecycle tools
-        assert len(schemas) == 1
+        assert len(schemas) == 2
         assert schemas[0]["name"] == "tool1"
 
     def test_get_tools_by_server(self):

hud/agents/tests/test_base_runtime.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+from unittest import mock
+
+import mcp.types as types
+import pytest
+
+from hud.agents.base import MCPAgent, find_content, find_reward, text_to_blocks
+from hud.types import AgentResponse, MCPToolCall, MCPToolResult
+
+
+class DummyAgent(MCPAgent):
+    async def get_system_messages(self):
+        return [types.TextContent(text="sys", type="text")]
+
+    async def get_response(self, messages):
+        # Single step: no tool calls -> done
+        return AgentResponse(content="ok", tool_calls=[], done=True)
+
+    async def format_blocks(self, blocks):
+        # Return as-is
+        return blocks
+
+    async def format_tool_results(self, tool_calls, tool_results):
+        return [types.TextContent(text="tools", type="text")]
+
+
+@pytest.mark.asyncio
+async def test_run_with_string_prompt_auto_client(monkeypatch):
+    # Fake MCPClient with required methods
+    fake_client = mock.AsyncMock()
+    fake_client.initialize.return_value = None
+    fake_client.list_tools.return_value = []
+    fake_client.shutdown.return_value = None
+
+    # Patch MCPClient construction inside initialize()
+    with mock.patch("hud.clients.MCPClient", return_value=fake_client):
+        agent = DummyAgent(mcp_client=fake_client, auto_trace=False)
+        result = await agent.run("hello", max_steps=1)
+    assert result.done is True and result.isError is False
+
+
+def test_find_reward_and_content_extractors():
+    # Structured content
+    r = MCPToolResult(
+        content=text_to_blocks("{}"), isError=False, structuredContent={"reward": 0.7}
+    )
+    assert find_reward(r) == 0.7
+
+    # Text JSON
+    r2 = MCPToolResult(content=text_to_blocks('{"score": 0.5, "content": "hi"}'), isError=False)
+    assert find_reward(r2) == 0.5
+    assert find_content(r2) == "hi"
+
+
+@pytest.mark.asyncio
+async def test_call_tools_error_paths():
+    fake_client = mock.AsyncMock()
+    # First call succeeds
+    ok_result = MCPToolResult(content=text_to_blocks("ok"), isError=False)
+    fake_client.call_tool.side_effect = [ok_result, RuntimeError("boom")]
+    agent = DummyAgent(mcp_client=fake_client, auto_trace=False)
+    results = await agent.call_tools(
+        [MCPToolCall(name="a", arguments={}), MCPToolCall(name="b", arguments={})]
+    )
+    assert results[0].isError is False
+    assert results[1].isError is True
+
+
+@pytest.mark.asyncio
+async def test_initialize_without_client_raises_valueerror():
+    agent = DummyAgent(mcp_client=None, auto_trace=False)
+    with pytest.raises(ValueError):
+        await agent.initialize(None)
+
+
+def test_get_available_tools_before_initialize_raises():
+    agent = DummyAgent(mcp_client=mock.AsyncMock(), auto_trace=False)
+    with pytest.raises(RuntimeError):
+        agent.get_available_tools()
+
+
+@pytest.mark.asyncio
+async def test_format_message_invalid_type_raises():
+    agent = DummyAgent(mcp_client=mock.AsyncMock(), auto_trace=False)
+    with pytest.raises(ValueError):
+        await agent.format_message({"oops": 1})  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_call_tools_timeout_error_shutdown_called():
+    fake_client = mock.AsyncMock()
+    fake_client.call_tool.side_effect = TimeoutError("timeout")
+    fake_client.shutdown.return_value = None
+    agent = DummyAgent(mcp_client=fake_client, auto_trace=False)
+    with pytest.raises(TimeoutError):
+        await agent.call_tools(MCPToolCall(name="x", arguments={}))
+    fake_client.shutdown.assert_awaited_once()
+
+
+def test_text_to_blocks_shapes():
+    blocks = text_to_blocks("x")
+    assert isinstance(blocks, list) and blocks and isinstance(blocks[0], types.TextContent)
+
+
+@pytest.mark.asyncio
+async def test_run_returns_connection_error_trace(monkeypatch):
+    fake_client = mock.AsyncMock()
+    fake_client.mcp_config = {}
+    fake_client.initialize.side_effect = RuntimeError("Connection refused http://localhost:1234")
+    fake_client.list_tools.return_value = []
+    fake_client.shutdown.return_value = None
+
+    class DummyCM:
+        def __exit__(self, *args, **kwargs):
+            return False
+
+    monkeypatch.setattr("hud.utils.mcp.setup_hud_telemetry", lambda *args, **kwargs: DummyCM())
+
+    agent = DummyAgent(mcp_client=fake_client, auto_trace=False)
+    result = await agent.run("p", max_steps=1)
+    assert result.isError is True
+    assert "Could not connect" in (result.content or "")
+
+
+@pytest.mark.asyncio
+async def test_run_calls_response_tool_when_configured(monkeypatch):
+    fake_client = mock.AsyncMock()
+    fake_client.mcp_config = {}
+    fake_client.initialize.return_value = None
+    fake_client.list_tools.return_value = []
+    fake_client.shutdown.return_value = None
+    ok = MCPToolResult(content=text_to_blocks("ok"), isError=False)
+    fake_client.call_tool.return_value = ok
+
+    class DummyCM:
+        def __exit__(self, *args, **kwargs):
+            return False
+
+    monkeypatch.setattr("hud.utils.mcp.setup_hud_telemetry", lambda *args, **kwargs: DummyCM())
+
+    agent = DummyAgent(mcp_client=fake_client, auto_trace=False, response_tool_name="submit")
+    result = await agent.run("hello", max_steps=1)
+    assert result.isError is False
+    fake_client.call_tool.assert_awaited()
+
+
+@pytest.mark.asyncio
+async def test_get_available_tools_after_initialize(monkeypatch):
+    fake_client = mock.AsyncMock()
+    fake_client.mcp_config = {}
+    fake_client.initialize.return_value = None
+    fake_client.list_tools.return_value = []
+    fake_client.shutdown.return_value = None
+
+    class DummyCM:
+        def __exit__(self, *args, **kwargs):
+            return False
+
+    monkeypatch.setattr("hud.utils.mcp.setup_hud_telemetry", lambda *args, **kwargs: DummyCM())
+
+    agent = DummyAgent(mcp_client=fake_client, auto_trace=False)
+    await agent.initialize(None)
+    assert agent.get_available_tools() == []

hud/cli/__init__.py

@@ -242,15 +242,18 @@ def debug(
                 if build and not build_environment(directory, image_name):
                     raise typer.Exit(1)
 
-            # Build Docker command
-            from .utils.docker import build_run_command
+            # Build Docker command with folder-mode envs
+            from .utils.docker import create_docker_run_command
 
-            command = build_run_command(image_name, docker_args)
+            command = create_docker_run_command(
+                image_name, docker_args=docker_args, env_dir=directory
+            )
         else:
             # Assume it's an image name
             image = first_param
             from .utils.docker import build_run_command
 
+            # Image-only mode: do not auto-inject local .env
             command = build_run_command(image, docker_args)
     else:
         console.print(
@@ -796,33 +799,19 @@ def eval(
         help="Comma-separated list of allowed tools",
     ),
     max_concurrent: int = typer.Option(
-        50,
+        30,
         "--max-concurrent",
-        help="Max concurrent tasks (prevents rate limits in both asyncio and parallel modes)",
+        help="Maximum concurrent tasks (1-200 recommended, prevents rate limits)",
     ),
     max_steps: int | None = typer.Option(
         None,
         "--max-steps",
         help="Maximum steps per task (default: 10 for single, 50 for full)",
     ),
-    parallel: bool = typer.Option(
-        False,
-        "--parallel",
-        help="Use process-based parallel execution for large datasets (100+ tasks)",
-    ),
-    max_workers: int | None = typer.Option(
-        None,
-        "--max-workers",
-        help="Number of worker processes for parallel mode (auto-optimized if not set)",
-    ),
-    max_concurrent_per_worker: int = typer.Option(
-        20,
-        "--max-concurrent-per-worker",
-        help="Maximum concurrent tasks per worker in parallel mode",
-    ),
     verbose: bool = typer.Option(
         False,
         "--verbose",
+        "-v",
         help="Enable verbose output from the agent",
     ),
     very_verbose: bool = typer.Option(
@@ -867,14 +856,14 @@ def eval(
 
             source = find_tasks_file(None, msg="Select a tasks file to run")
             hud_console.success(f"Selected: {source}")
-        except Exception as e:
+        except (FileNotFoundError, Exception):
             hud_console.error(
                 "No source provided and no task/eval JSON files found in current directory"
             )
             hud_console.info(
                 "Usage: hud eval <source> or create a task JSON file (e.g., task.json, tasks.jsonl)"
             )
-            raise typer.Exit(1) from e
+            raise typer.Exit(1) from None
 
     # Import eval_command lazily to avoid importing agent dependencies
     try:
@@ -950,9 +939,6 @@ def eval(
         allowed_tools=allowed_tools,
         max_concurrent=max_concurrent,
         max_steps=max_steps,
-        parallel=parallel,
-        max_workers=max_workers,
-        max_concurrent_per_worker=max_concurrent_per_worker,
         verbose=verbose,
         very_verbose=very_verbose,
         vllm_base_url=vllm_base_url,
@@ -1126,6 +1112,13 @@ def set(
 
 def main() -> None:
     """Main entry point for the CLI."""
+    # Check for updates (including on --version command)
+    # Skip only on help-only commands
+    if not (len(sys.argv) == 1 or (len(sys.argv) == 2 and sys.argv[1] in ["--help", "-h"])):
+        from .utils.version_check import display_update_prompt
+
+        display_update_prompt()
+
     # Handle --version flag before Typer parses args
     if "--version" in sys.argv:
         try:

hud/cli/build.py

@@ -161,49 +161,42 @@ async def analyze_mcp_environment(
     hud_console = HUDConsole()
     env_vars = env_vars or {}
 
-    # Build Docker command to run the image
-    docker_cmd = ["docker", "run", "--rm", "-i"]
+    # Build Docker command to run the image, injecting any provided env vars
+    from hud.cli.utils.docker import build_env_flags
 
-    # Add environment variables
-    for key, value in env_vars.items():
-        docker_cmd.extend(["-e", f"{key}={value}"])
+    docker_cmd = ["docker", "run", "--rm", "-i", *build_env_flags(env_vars), image]
 
-    docker_cmd.append(image)
+    # Show full docker command being used for analysis
+    hud_console.dim_info("Command:", " ".join(docker_cmd))
 
-    # Create MCP config
-    config = {
-        "server": {"command": docker_cmd[0], "args": docker_cmd[1:] if len(docker_cmd) > 1 else []}
-    }
+    # Create MCP config consistently with analyze helpers
+    from hud.cli.analyze import parse_docker_command
+
+    mcp_config = parse_docker_command(docker_cmd)
 
     # Initialize client and measure timing
     start_time = time.time()
-    client = MCPClient(mcp_config=config, verbose=verbose, auto_trace=False)
+    client = MCPClient(mcp_config=mcp_config, verbose=verbose, auto_trace=False)
     initialized = False
 
     try:
         if verbose:
-            hud_console.info(f"Initializing MCP client with command: {' '.join(docker_cmd)}")
+            hud_console.info("Initializing MCP client...")
 
-        # Add timeout to fail fast instead of hanging (30 seconds)
+        # Add timeout to fail fast instead of hanging (60 seconds)
         await asyncio.wait_for(client.initialize(), timeout=60.0)
         initialized = True
         initialize_ms = int((time.time() - start_time) * 1000)
 
-        # Get tools
-        tools = await client.list_tools()
-
-        # Extract tool information
-        tool_info = []
-        for tool in tools:
-            tool_dict = {"name": tool.name, "description": tool.description}
-            if hasattr(tool, "inputSchema") and tool.inputSchema:
-                tool_dict["inputSchema"] = tool.inputSchema
-            tool_info.append(tool_dict)
+        # Delegate to standard analysis helper for consistency
+        full_analysis = await client.analyze_environment()
 
+        # Normalize to build's expected fields
+        tools_list = full_analysis.get("tools", [])
         return {
             "initializeMs": initialize_ms,
-            "toolCount": len(tools),
-            "tools": tool_info,
+            "toolCount": len(tools_list),
+            "tools": tools_list,
             "success": True,
         }
     except TimeoutError:
@@ -295,6 +288,10 @@ def build_environment(
         hud_console.error(f"Directory not found: {directory}")
         raise typer.Exit(1)
 
+    from hud.cli.utils.docker import require_docker_running
+
+    require_docker_running()
+
     # Step 1: Check for hud.lock.yaml (previous build)
     lock_path = env_dir / "hud.lock.yaml"
     base_name = None
@@ -355,13 +352,24 @@ def build_environment(
 
     hud_console.success(f"Built temporary image: {temp_tag}")
 
-    # Analyze the environment
+    # Analyze the environment (merge folder .env if present)
     hud_console.progress_message("Analyzing MCP environment...")
 
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     try:
-        analysis = loop.run_until_complete(analyze_mcp_environment(temp_tag, verbose, env_vars))
+        # Merge .env from env_dir for analysis only
+        try:
+            from hud.cli.utils.docker import load_env_vars_for_dir
+
+            env_from_file = load_env_vars_for_dir(env_dir)
+        except Exception:
+            env_from_file = {}
+        merged_env_for_analysis = {**env_from_file, **(env_vars or {})}
+
+        analysis = loop.run_until_complete(
+            analyze_mcp_environment(temp_tag, verbose, merged_env_for_analysis)
+        )
     except Exception as e:
         hud_console.error(f"Failed to analyze MCP environment: {e}")
         hud_console.info("")

hud/cli/dev.py

@@ -504,15 +504,12 @@ def run_docker_dev_server(
     base_name = image_name.replace(":", "-").replace("/", "-")
     container_name = f"{base_name}-dev-{pid}"
 
-    # Build docker run command with volume mounts
-    docker_cmd = [
-        "docker",
-        "run",
-        "--rm",
-        "-i",
+    # Build docker run command with volume mounts and folder-mode envs
+    from .utils.docker import create_docker_run_command
+
+    base_args = [
         "--name",
         container_name,
-        # Mount both server and environment for hot-reload
         "-v",
         f"{env_dir.absolute()}/server:/app/server:rw",
         "-v",
@@ -524,29 +521,14 @@ def run_docker_dev_server(
         "-e",
         "HUD_DEV=1",
     ]
+    combined_args = [*base_args, *docker_args] if docker_args else base_args
+    docker_cmd = create_docker_run_command(
+        image_name,
+        docker_args=combined_args,
+        env_dir=env_dir,
+    )
 
-    # Load .env file if present
-    env_file = env_dir / ".env"
-    loaded_env_vars: dict[str, str] = {}
-    if env_file.exists():
-        try:
-            from hud.cli.utils.config import parse_env_file
-
-            env_contents = env_file.read_text(encoding="utf-8")
-            loaded_env_vars = parse_env_file(env_contents)
-            for key, value in loaded_env_vars.items():
-                docker_cmd.extend(["-e", f"{key}={value}"])
-            if verbose and loaded_env_vars:
-                hud_console.info(f"Loaded {len(loaded_env_vars)} env var(s) from .env")
-        except Exception as e:
-            hud_console.warning(f"Failed to load .env file: {e}")
-
-    # Add user-provided Docker arguments
-    if docker_args:
-        docker_cmd.extend(docker_args)
-
-    # Append the image name
-    docker_cmd.append(image_name)
+    # Env flags already injected by create_docker_run_command
 
     # Print startup info
     hud_console.header("HUD Development Mode (Docker)")