cua-agent 0.4.14__py3-none-any.whl → 0.7.16__py3-none-any.whl

agent/integrations/hud/__init__.py

@@ -1,77 +1,167 @@
-"""HUD integration for ComputerAgent."""
-
-import logging
-from typing import Any, Optional, Dict
-from hud import run_job as hud_run_job
-
-from .agent import ComputerAgent
-from .adapter import ComputerAgentAdapter
-from .computer_handler import HUDComputerHandler
-
-
-async def run_job(
-    model: str,
-    task_or_taskset: Any,
-    job_name: str,
-    # Job kwargs
-    auto_reply_question: bool = False,
-    adapter_cls: Any = None,
-    adapter_kwargs: Optional[Dict[str, Any]] = None,
-    max_steps_per_task: int = 20,
-    run_parallel: bool = True,
-    job_metadata: Optional[Dict[str, Any]] = None,
-    show_progress: bool = True,
-    max_concurrent_env_creations: Optional[int] = 30,  # Limits gym.make calls
-    max_concurrent_agent_predictions: Optional[int] = None,  # No limit on LLM calls
-    max_concurrent_tasks: Optional[int] = 30,  # Limits overall task concurrency
-    **agent_kwargs: Any
-) -> Any:
-    """
-    Run a job using ComputerAgent with the specified model.
-    
-    Args:
-        model: Model string for ComputerAgent (e.g., "anthropic/claude-3-5-sonnet-20241022")
-        task_or_taskset: Task or TaskSet to run
-        job_name: Name for the job
-        auto_reply_question: Whether to auto-reply to questions
-        adapter_cls: Custom adapter class (defaults to ComputerAgentAdapter)
-        adapter_kwargs: Additional kwargs for the adapter
-        max_steps_per_task: Maximum steps per task
-        run_parallel: Whether to run tasks in parallel
-        job_metadata: Additional metadata for the job
-        show_progress: Whether to show progress
-        max_concurrent_env_creations: Max concurrent environment creations
-        max_concurrent_agent_predictions: Max concurrent agent predictions
-        max_concurrent_tasks: Max concurrent tasks
-        **agent_kwargs: Additional kwargs to pass to ComputerAgent
-    
-    Returns:
-        Job instance from HUD
-    """
-    # combine verbose and verbosity kwargs
-    if "verbose" in agent_kwargs:
-        agent_kwargs["verbosity"] = logging.INFO
-        del agent_kwargs["verbose"]
-    verbose = True if agent_kwargs.get("verbosity", logging.WARNING) > logging.INFO else False
-    
-    # run job
-    return await hud_run_job(
-        agent_cls=ComputerAgent,
-        agent_kwargs={"model": model, **agent_kwargs},
-        task_or_taskset=task_or_taskset,
-        job_name=job_name,
-        auto_reply_question=auto_reply_question,
-        adapter_cls=adapter_cls,
-        adapter_kwargs=adapter_kwargs,
-        max_steps_per_task=max_steps_per_task,
-        run_parallel=run_parallel,
-        job_metadata=job_metadata,
-        show_progress=show_progress,
-        verbose=verbose,
-        max_concurrent_env_creations=max_concurrent_env_creations,
-        max_concurrent_agent_predictions=max_concurrent_agent_predictions,
-        max_concurrent_tasks=max_concurrent_tasks
+"""HUD integration: dataset runners and MCP-based computer agent export.
+
+This module exposes helpers to evaluate HUD-compatible datasets and exports
+the MCP-compatible computer agent implementation.
+
+Exports:
+- run_single_task(dataset, ...)
+- run_full_dataset(dataset, ...)
+- MCPComputerAgent
+"""
+
+import time
+from typing import Any, Optional
+
+from agent.computers import is_agent_computer
+from datasets import Dataset, load_dataset
+from hud import trace
+from hud.datasets import Task, run_dataset
+
+from .agent import MCPComputerAgent
+
+# ---------------------------------------------------------------------------
+# Single-task runner
+# ---------------------------------------------------------------------------
+
+
+async def run_single_task(
+    dataset: str | Dataset | list[dict[str, Any]],
+    *,
+    task_id: int = 0,
+    model: str | None = None,
+    allowed_tools: list[str] | None = None,
+    # === ComputerAgent kwargs ===
+    tools: list[Any] | None = None,
+    custom_loop: Any | None = None,
+    only_n_most_recent_images: int | None = None,
+    callbacks: list[Any] | None = None,
+    instructions: str | None = None,
+    verbosity: int | None = None,
+    trajectory_dir: str | dict | None = None,
+    max_retries: int | None = 3,
+    screenshot_delay: float | int = 0.5,
+    use_prompt_caching: bool | None = False,
+    max_trajectory_budget: float | dict | None = None,
+    telemetry_enabled: bool | None = True,
+) -> None:
+    """Load one task from the dataset and execute it with MCPComputerAgent."""
+
+    # Load dataset and pick a sample
+    if isinstance(dataset, str):
+        dataset = load_dataset(dataset, split="train")  # type: ignore[arg-type]
+    elif isinstance(dataset, list):
+        dataset = dataset
+    else:
+        dataset = dataset["train"]
+
+    sample_task = dataset[task_id]  # type: ignore[index]
+    task_prompt = sample_task.get("prompt", f"Task {sample_task.get('id', 0)}")  # type: ignore[attr-defined]
+
+    # Filter any existing Computer tools
+    # The eval framework will add its own Computer tool per task
+    if tools:
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
+    with trace(name=task_prompt):
+        task = Task(**sample_task)  # type: ignore[arg-type]
+
+        agent = MCPComputerAgent(
+            model=model or "computer-use-preview",
+            allowed_tools=allowed_tools or ["openai_computer"],
+            # === ComputerAgent kwargs passthrough ===
+            tools=tools,
+            custom_loop=custom_loop,
+            only_n_most_recent_images=only_n_most_recent_images,
+            callbacks=callbacks,
+            instructions=instructions,
+            verbosity=verbosity,
+            trajectory_dir=trajectory_dir,
+            max_retries=max_retries,
+            screenshot_delay=screenshot_delay,
+            use_prompt_caching=use_prompt_caching,
+            max_trajectory_budget=max_trajectory_budget,
+            telemetry_enabled=telemetry_enabled,
+        )
+        print(f"Running: {task_prompt}")
+        result = await agent.run(task, max_steps=10)
+        print(f"✅ Reward: {result.reward}")
+
+
+# ---------------------------------------------------------------------------
+# Full-dataset runner
+# ---------------------------------------------------------------------------
+
+
+async def run_full_dataset(
+    dataset: str | Dataset | list[dict[str, Any]],
+    *,
+    job_name: Optional[str] = None,
+    model: str | None = None,
+    allowed_tools: list[str] | None = None,
+    max_concurrent: int = 30,
+    max_steps: int = 50,
+    split: str = "train",
+    trajectory_dir: str | dict | None = None,
+    # === ComputerAgent kwargs ===
+    tools: list[Any] | None = None,
+    custom_loop: Any | None = None,
+    only_n_most_recent_images: int | None = 5,
+    callbacks: list[Any] | None = None,
+    instructions: str | None = None,
+    verbosity: int | None = None,
+    max_retries: int | None = 3,
+    screenshot_delay: float | int = 0.5,
+    use_prompt_caching: bool | None = False,
+    max_trajectory_budget: float | dict | None = None,
+    telemetry_enabled: bool | None = True,
+) -> list[Any]:
+    """Run evaluation across the entire dataset using hud.datasets.run_dataset."""
+
+    # Run with our MCP-based agent class.
+    if isinstance(dataset, str):
+        dataset_name = dataset.split("/")[-1]
+        job_name = job_name or f"Evaluation {dataset_name}"
+        dataset = load_dataset(dataset, split=split)  # type: ignore[arg-type]
+    else:
+        dataset_name = "custom"
+        job_name = job_name or f"Evaluation {time.strftime('%H:%M %Y-%m-%d')}"
+
+    # Filter any existing Computer tools
+    # The eval framework will add its own Computer tool per task
+    if tools:
+        tools = [tool for tool in tools if not is_agent_computer(tool)]
+
+    # Execute evaluation
+    return await run_dataset(
+        name=job_name,
+        dataset=dataset,
+        agent_class=MCPComputerAgent,
+        agent_config={
+            "model": model,
+            "allowed_tools": allowed_tools,
+            "trajectory_dir": trajectory_dir,
+            # === ComputerAgent kwargs passthrough ===
+            "tools": tools,
+            "custom_loop": custom_loop,
+            "only_n_most_recent_images": only_n_most_recent_images,
+            "callbacks": callbacks,
+            "instructions": instructions,
+            "verbosity": verbosity,
+            "max_retries": max_retries,
+            "screenshot_delay": screenshot_delay,
+            "use_prompt_caching": use_prompt_caching,
+            "max_trajectory_budget": max_trajectory_budget,
+            "telemetry_enabled": telemetry_enabled,
+        },
+        max_concurrent=max_concurrent,
+        metadata={"dataset": dataset_name},
+        max_steps=max_steps,
+        auto_respond=True,
     )
 
 
-__all__ = ["ComputerAgent", "ComputerAgentAdapter", "HUDComputerHandler", "run_job"]
+__all__ = [
+    "run_single_task",
+    "run_full_dataset",
+    "MCPComputerAgent",
+]