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

hud/datasets/loader.py

@@ -0,0 +1,327 @@
+"""Task loading utilities for HUD.
+
+Unified interface for loading evaluation tasks from:
+- HUD API (v5 format)
+- Local JSON/JSONL files (v4 LegacyTask format, auto-converted)
+- HuggingFace datasets (v4 LegacyTask format, auto-converted)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, overload
+
+if TYPE_CHECKING:
+    from hud.eval.task import Task
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["load_dataset", "load_tasks", "save_tasks"]
+
+
+def _load_raw_from_file(path: Path) -> list[dict[str, Any]]:
+    """Load raw task dicts from a local JSON or JSONL file."""
+    raw_items: list[dict[str, Any]] = []
+
+    if path.suffix == ".jsonl":
+        # JSONL: one task per line
+        with open(path, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                item = json.loads(line)
+                # Handle case where line contains a list
+                if isinstance(item, list):
+                    raw_items.extend(i for i in item if isinstance(i, dict))
+                elif isinstance(item, dict):
+                    raw_items.append(item)
+                else:
+                    raise ValueError(
+                        f"Invalid JSONL format: expected dict or list, got {type(item)}"
+                    )
+    else:
+        # JSON: array of tasks
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+
+        if isinstance(data, list):
+            raw_items = [item for item in data if isinstance(item, dict)]
+        elif isinstance(data, dict):
+            raw_items = [data]
+        else:
+            raise ValueError(f"JSON file must contain an array or object, got {type(data)}")
+
+    return raw_items
+
+
+def _load_from_file(path: Path) -> list[Task]:
+    """Load tasks from a local JSON or JSONL file."""
+    from hud.eval.task import Task
+
+    raw_items = _load_raw_from_file(path)
+    return [Task(**item) for item in raw_items]
+
+
+def _load_raw_from_huggingface(dataset_name: str) -> list[dict[str, Any]]:
+    """Load raw task dicts from HuggingFace dataset."""
+    try:
+        from datasets import load_dataset as hf_load_dataset
+    except ImportError as e:
+        raise ImportError(
+            "Please install 'datasets' to load from HuggingFace: uv pip install datasets"
+        ) from e
+
+    # Parse dataset name and optional split
+    if ":" in dataset_name:
+        name, split = dataset_name.split(":", 1)
+    else:
+        name = dataset_name
+        split = "train"  # Default split
+
+    logger.info("Loading from HuggingFace dataset: %s (split=%s)", name, split)
+    dataset = hf_load_dataset(name, split=split)
+
+    raw_items: list[dict[str, Any]] = []
+    for item in dataset:
+        if not isinstance(item, dict):
+            raise ValueError(f"Invalid HuggingFace dataset: expected dict, got {type(item)}")
+        raw_items.append(dict(item))
+
+    return raw_items
+
+
+def _load_from_huggingface(dataset_name: str) -> list[Task]:
+    """Load tasks from HuggingFace dataset."""
+    raw_items = _load_raw_from_huggingface(dataset_name)
+    from hud.eval.task import Task
+
+    return [Task(**item) for item in raw_items]
+
+
+def _load_raw_from_api(dataset_name: str) -> list[dict[str, Any]]:
+    """Load raw task dicts from HUD API."""
+    import httpx
+
+    from hud.settings import settings
+
+    headers = {}
+    if settings.api_key:
+        headers["Authorization"] = f"Bearer {settings.api_key}"
+
+    with httpx.Client() as client:
+        response = client.get(
+            f"{settings.hud_api_url}/tasks/evalset/{dataset_name}",
+            headers=headers,
+            params={"all": "true"},
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract tasks dict from response
+        tasks_dict = data.get("tasks", {})
+
+        raw_items: list[dict[str, Any]] = []
+        for task_id, task_data in tasks_dict.items():
+            if task_data.get("id") is None:
+                task_data["id"] = task_id
+            raw_items.append(task_data)
+
+        return raw_items
+
+
+def _load_from_api(dataset_name: str) -> list[Task]:
+    """Load tasks from HUD API."""
+    from hud.eval.task import Task
+
+    raw_items = _load_raw_from_api(dataset_name)
+    return [Task(**item) for item in raw_items]
+
+
+@overload
+def load_tasks(source: str, *, raw: bool = False) -> list[Task]: ...
+
+
+@overload
+def load_tasks(source: str, *, raw: bool = True) -> list[dict[str, Any]]: ...
+
+
+def load_tasks(source: str, *, raw: bool = False) -> list[Task] | list[dict[str, Any]]:
+    """Load tasks from a source.
+
+    Supports multiple sources with auto-detection:
+    - Local file path (JSON or JSONL)
+    - HUD API dataset slug (e.g., "hud-evals/SheetBench-50")
+    - HuggingFace dataset (e.g., "username/dataset" or "username/dataset:split")
+
+    Automatically detects and converts v4 LegacyTask format to v5 Task.
+
+    Args:
+        source: Task source. Can be:
+            - Path to a local JSON/JSONL file
+            - HUD API dataset slug (e.g., "hud-evals/SheetBench-50")
+            - HuggingFace dataset name (e.g., "hud-evals/tasks" or "hud-evals/tasks:train")
+        raw: If True, return raw dicts without validation or env var substitution.
+            Useful for preserving template strings like "${HUD_API_KEY}".
+
+    Returns:
+        - If raw=False (default): list[Task] ready to use with hud.eval()
+        - If raw=True: list[dict] with raw task data
+
+    Example:
+        ```python
+        import hud
+        from hud.datasets import load_tasks
+
+        # Load from HUD API
+        tasks = load_tasks("hud-evals/SheetBench-50")
+
+        # Load from local file (v4 format auto-converted)
+        tasks = load_tasks("./my-tasks.json")
+
+        # Load from HuggingFace
+        tasks = load_tasks("hud-evals/benchmark:test")
+
+        # Load raw dicts (preserves env var placeholders)
+        raw_tasks = load_tasks("./tasks.json", raw=True)
+
+        # Run evaluation
+        async with hud.eval(tasks) as ctx:
+            await agent.run(ctx)
+        ```
+
+    Raises:
+        ValueError: If task loading fails
+    """
+    # Check if it's a local file
+    path = Path(source)
+    if path.exists() and path.suffix in {".json", ".jsonl"}:
+        logger.info("Loading tasks from file: %s", source)
+        items = _load_raw_from_file(path) if raw else _load_from_file(path)
+        logger.info("Loaded %d tasks from %s", len(items), source)
+        return items
+
+    # Try HUD API first
+    try:
+        logger.info("Trying HUD API: %s", source)
+        items = _load_raw_from_api(source) if raw else _load_from_api(source)
+        logger.info("Loaded %d tasks from HUD API: %s", len(items), source)
+        return items
+    except Exception as hud_error:
+        logger.debug("HUD API load failed (%s), trying HuggingFace", hud_error)
+
+    # Try HuggingFace as fallback
+    try:
+        logger.info("Trying HuggingFace dataset: %s", source)
+        items = _load_raw_from_huggingface(source) if raw else _load_from_huggingface(source)
+        logger.info("Loaded %d tasks from HuggingFace: %s", len(items), source)
+        return items
+    except ImportError:
+        raise ValueError(
+            f"Failed to load tasks from '{source}'. "
+            "Install 'datasets' package for HuggingFace support."
+        ) from None
+    except Exception as hf_error:
+        raise ValueError(f"Failed to load tasks from '{source}': {hf_error}") from hf_error
+
+
+def save_tasks(
+    name: str,
+    tasks: list[Task],
+) -> str:
+    """Save tasks to the HUD API.
+
+    Creates or updates an evalset with the given tasks.
+
+    Args:
+        name: Evalset name/slug (e.g., "my-evals/benchmark-v1").
+            If no org prefix, uses user's default org.
+        tasks: List of Task objects (v5 format) to save.
+
+    Returns:
+        The evalset ID of the created/updated evalset.
+
+    Example:
+        ```python
+        from hud.datasets import save_tasks, load_tasks
+        from hud.eval.task import Task
+        from hud.environment import Environment
+
+        # Create tasks
+        env = Environment("my-env")
+        tasks = [
+            Task(env=env, scenario="checkout", args={"user": "alice"}),
+            Task(env=env, scenario="checkout", args={"user": "bob"}),
+        ]
+
+        # Save to HUD API
+        evalset_id = save_tasks("my-evals/benchmark-v1", tasks)
+
+        # Later, load them back
+        loaded = load_tasks("my-evals/benchmark-v1")
+        ```
+
+    Raises:
+        TypeError: If any task is not a v5 Task object (must have 'scenario')
+        ValueError: If API key is not set or save fails
+    """
+    import httpx
+
+    from hud.settings import settings
+
+    if not settings.api_key:
+        raise ValueError("HUD_API_KEY is required to save tasks")
+
+    # Validate all tasks are v5 format (must have 'scenario')
+    for i, task in enumerate(tasks):
+        if not hasattr(task, "scenario"):
+            raise TypeError(
+                f"Task at index {i} is missing 'scenario' - only v5 Task objects can be saved. "
+                "Use Task.from_v4(legacy_task) to convert from LegacyTask."
+            )
+
+    # Convert tasks to dicts (Task is a Pydantic model)
+    task_dicts = [task.model_dump(mode="json", exclude_none=True) for task in tasks]
+
+    # Build request payload
+    payload: dict[str, Any] = {
+        "name": name,
+        "tasks": task_dicts,
+    }
+
+    headers = {"Authorization": f"Bearer {settings.api_key}"}
+
+    try:
+        with httpx.Client(timeout=60) as client:
+            response = client.post(
+                f"{settings.hud_api_url}/tasks/evalset",
+                json=payload,
+                headers=headers,
+            )
+            response.raise_for_status()
+            data = response.json()
+            evalset_id = data.get("evalset_id") or data.get("id") or name
+            logger.info("Saved %d tasks to evalset: %s", len(tasks), evalset_id)
+            return evalset_id
+    except httpx.HTTPStatusError as e:
+        raise ValueError(f"Failed to save tasks: {e.response.text}") from e
+    except Exception as e:
+        raise ValueError(f"Failed to save tasks: {e}") from e
+
+
+# Deprecated alias for backwards compatibility
+def load_dataset(source: str, *, raw: bool = False) -> list[Task] | list[dict[str, Any]]:
+    """Deprecated: Use load_tasks() instead.
+
+    .. deprecated:: 0.6.0
+        load_dataset() is deprecated. Use load_tasks() instead.
+    """
+    warnings.warn(
+        "load_dataset() is deprecated. Use load_tasks() instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return load_tasks(source, raw=raw)

hud/datasets/runner.py

@@ -1,126 +1,213 @@
-"""Standard asyncio-based dataset runner."""
+"""Core task runner for evaluating agents on datasets.
+
+Requires the [agents] extra: pip install hud-python[agents]
+"""
 
 from __future__ import annotations
 
-import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, cast
-
-from datasets import Dataset, load_dataset
+from typing import TYPE_CHECKING, Any
 
-from hud.agents.misc import ResponseAgent
-from hud.types import Task
+import hud
+from hud.types import AgentType, LegacyTask, TaskInput, Trace
 
 if TYPE_CHECKING:
-    from hud.agents import MCPAgent
+    from collections.abc import Sequence
+
+    from hud.eval.context import EvalContext
+    from hud.eval.task import Task
 
 logger = logging.getLogger("hud.datasets")
 
 
 async def run_dataset(
-    name: str,
-    dataset: str | Dataset | list[dict[str, Any]],
-    agent_class: type[MCPAgent],
-    agent_config: dict[str, Any] | None = None,
-    max_concurrent: int = 30,
-    metadata: dict[str, Any] | None = None,
+    tasks: str | TaskInput | Sequence[TaskInput],
+    agent_type: str | AgentType,
+    *,
+    agent_params: dict[str, Any] | None = None,
     max_steps: int = 10,
-    split: str = "train",
-    auto_respond: bool = False,
-    custom_system_prompt: str | None = None,
-) -> list[Any]:
-    """
-    Run all tasks in a dataset with automatic job tracking.
+    max_concurrent: int = 30,
+    group_size: int = 1,
+    quiet: bool = True,
+) -> list[EvalContext]:
+    """Run an agent on a dataset of tasks.
+
+    This is the primary entry point for running evaluations programmatically.
+    The agent is created fresh for each task context to ensure correct tool initialization.
 
     Args:
-        name: Name for the job
-        dataset: HuggingFace dataset identifier (e.g. "hud-evals/SheetBench-50"),
-                Dataset object, OR list of Task objects
-        agent_class: Agent class to instantiate (e.g., ClaudeAgent)
-        agent_config: Configuration/kwargs for agent (model, etc.)
-        max_concurrent: Maximum parallel task execution
-        metadata: Optional metadata for the job
-        max_steps: Maximum steps per task
-        split: Dataset split to use when loading from string (default: "train")
-        auto_respond: Whether to use auto-response agent
-        custom_system_prompt: Override system prompt for all tasks
+        tasks: Tasks to run. Can be:
+            - A source string (file path, API slug) - loaded via load_tasks()
+            - A single TaskInput (Task, LegacyTask, or dict)
+            - A list of TaskInput objects
+        agent_type: Type of agent to create (e.g., "claude", "openai", AgentType.CLAUDE).
+        agent_params: Parameters to pass to agent.create().
+        max_steps: Maximum steps per task.
+        max_concurrent: Maximum concurrent tasks (for parallel execution).
+        group_size: Number of times to run each task (for variance estimation).
+        quiet: Whether to suppress printing eval links and opening browser (default True).
 
     Returns:
-        List of results from agent.run() in dataset order
+        List of EvalContext results from each task execution. Access `.reward` on each.
 
     Example:
-        >>> from hud.agents import ClaudeAgent
-        >>> # Option 1: From dataset string identifier
-        >>> results = await run_dataset(
-        ...     "SheetBench Eval",
-        ...     "hud-evals/SheetBench-50",
-        ...     ClaudeAgent,
-        ...     {"model": "claude-3-5-sonnet-20241022"},
-        ... )
-        >>> # Option 2: From HuggingFace dataset object
-        >>> from datasets import load_dataset
-        >>> dataset = load_dataset("hud-evals/SheetBench-50", split="train")
-        >>> results = await run_dataset("my_eval", dataset, ClaudeAgent)
-        >>> # Option 3: From list of dicts
-        >>> tasks = [{"prompt": "...", "mcp_config": {...}, ...}, ...]
-        >>> results = await run_dataset("browser_eval", tasks, ClaudeAgent)
-    """
-    # Import here to avoid circular imports
-    import hud
-
-    dataset_link = None
-
-    # Load dataset from string if needed
-    if isinstance(dataset, str):
-        logger.info("Loading dataset %s from HuggingFace...", dataset)
-        dataset_link = dataset
-
-        # Load dataset from HuggingFace
-        dataset = cast("Dataset", load_dataset(dataset, split=split))
-
-    # Create job context
-    job_metadata = metadata or {}
-    job_metadata["agent_class"] = agent_class.__name__
-    job_metadata["agent_config"] = agent_config
-
-    # Extract dataset verification info if available
-    if isinstance(dataset, Dataset) and not dataset_link:
-        try:
-            general_info = next(iter(dataset.info.__dict__["download_checksums"].keys())).split("/")
-            project = general_info[3]
-            dataset_name = general_info[4].split("@")[0]
-            dataset_link = f"{project}/{dataset_name}"
-        except Exception:
-            logger.warning("Failed to extract dataset verification info")
-
-    with hud.job(name, metadata=job_metadata, dataset_link=dataset_link) as job_obj:
-        # Run tasks with semaphore for concurrency control
-        sem = asyncio.Semaphore(max_concurrent)
-        results: list[Any | None] = [None] * len(dataset)
-
-        async def _worker(index: int, task_dict: Any, max_steps: int = 10) -> None:
-            async with sem:
-                # Create trace for this task
-                task_name = task_dict.get("prompt") or f"Task {index}"
-                if custom_system_prompt and "system_prompt" not in task_dict:
-                    task_dict["system_prompt"] = custom_system_prompt
-                # Ensure task_id is a string for baggage propagation
-                raw_task_id = task_dict.get("id")
-                safe_task_id = str(raw_task_id) if raw_task_id is not None else None
-                with hud.trace(task_name, job_id=job_obj.id, task_id=safe_task_id):
-                    # Convert dict to Task here, at trace level
-                    task = Task(**task_dict)
-
-                    agent = agent_class(**(agent_config or {}))
-
-                    if auto_respond:
-                        agent.response_agent = ResponseAgent()
-                    results[index] = await agent.run(task, max_steps=max_steps)
-
-        # Execute all tasks
-        await asyncio.gather(
-            *[_worker(i, task, max_steps=max_steps) for i, task in enumerate(dataset)],
-            return_exceptions=True,  # Don't fail entire batch on one error
+        ```python
+        from hud.datasets import load_tasks, run_dataset
+
+        # Load tasks and run
+        tasks = load_tasks("my-tasks.json")
+        results = await run_dataset(
+            tasks,
+            agent_type="claude",
+            agent_params={"checkpoint_name": "claude-sonnet-4-20250514"},
+            max_steps=50,
         )
 
-    return results
+        for ctx in results:
+            print(f"Reward: {ctx.reward}")
+        ```
+    """
+    from hud.datasets.loader import load_tasks
+    from hud.eval.task import Task
+
+    # Normalize tasks to list[Task]
+    task_list: list[Task]
+    if isinstance(tasks, str):
+        task_list = load_tasks(tasks)
+    elif isinstance(tasks, Task):
+        task_list = [tasks]
+    elif isinstance(tasks, LegacyTask | dict):
+        # Single LegacyTask or dict - convert to Task
+        task_list = [Task.from_v4(tasks)]
+    else:
+        # Sequence of TaskInput - convert each to Task
+        task_list = [t if isinstance(t, Task) else Task.from_v4(t) for t in tasks]
+
+    if not task_list:
+        raise ValueError("No tasks to run")
+
+    # Resolve agent class
+    agent_type_enum = agent_type if isinstance(agent_type, AgentType) else AgentType(agent_type)
+    agent_cls = agent_type_enum.cls
+
+    # Use hud.eval() for both single and parallel execution
+    async with hud.eval(
+        task_list,
+        group=group_size,
+        max_concurrent=max_concurrent,
+        quiet=quiet,
+    ) as ctx:
+        # Create agent fresh for each context (ensures correct tool initialization)
+        agent = agent_cls.create(**(agent_params or {}))
+        await agent.run(ctx, max_steps=max_steps)
+        # Reward is computed by EvalContext.__aexit__ from evaluate tools
+
+    # For parallel execution, results are collected via ctx.results
+    if hasattr(ctx, "results") and ctx.results:
+        return ctx.results
+
+    return [ctx]
+
+
+async def run_single_task(
+    task: Task,
+    *,
+    agent_type: AgentType,
+    agent_params: dict[str, Any] | None = None,
+    max_steps: int = 10,
+    job_id: str | None = None,
+    task_id: str | None = None,
+    group_id: str | None = None,
+    trace_name: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    trace_id: str | None = None,
+    api_key: str | None = None,
+    trace: bool = True,
+    quiet: bool = False,
+) -> Trace:
+    """Run a single task with full control over eval context parameters.
+
+    This is the low-level entry point for running individual tasks with explicit
+    trace/job/group IDs. Used by remote execution workers.
+
+    Args:
+        task: Task object to run. Use Task.from_v4() or load_tasks() to create.
+        agent_type: AgentType enum specifying the agent to use.
+        agent_params: Parameters passed to agent.create(). Should include
+            pre-configured model_client for inference gateway usage.
+        max_steps: Maximum steps allowed for the agent.
+        job_id: HUD job identifier for telemetry association.
+        task_id: Task identifier (used in trace name if trace_name not provided).
+        group_id: Optional group identifier for parallel runs.
+        trace_name: Name for the trace (defaults to task_id or task.id).
+        metadata: Additional metadata for the trace context.
+        trace_id: Pre-assigned trace ID (if provided by backend).
+        api_key: API key override for telemetry and backend calls.
+        trace: Whether to send trace data to backend (default True).
+        quiet: Whether to suppress printing eval link (default False).
+
+    Returns:
+        Trace result from the agent run.
+
+    Example:
+        ```python
+        from hud.datasets import run_single_task
+        from hud.eval.task import Task
+        from hud.types import AgentType
+        from openai import AsyncOpenAI
+
+        # Create task (from v4 dict or directly)
+        task = Task.from_v4({"prompt": "...", "mcp_config": {...}, "evaluate_tool": {...}})
+
+        # Configure agent with inference gateway
+        agent_params = {
+            "checkpoint_name": "gpt-4o",
+            "validate_api_key": False,
+            "model_client": AsyncOpenAI(
+                api_key=hud_api_key,
+                base_url=settings.hud_gateway_url,
+            ),
+        }
+
+        result = await run_single_task(
+            task=task,
+            agent_type=AgentType.OPENAI,
+            agent_params=agent_params,
+            max_steps=20,
+            job_id="job-123",
+            task_id="task-456",
+        )
+        ```
+    """
+    # Determine trace name
+    effective_trace_name = trace_name or task_id or task.id or "single_task"
+
+    # Run with explicit eval context parameters
+    async with hud.eval(
+        task,
+        name=effective_trace_name,
+        job_id=job_id,
+        group_id=group_id,
+        trace_id=trace_id,
+        api_key=api_key,
+        trace=trace,
+        quiet=quiet,
+    ) as ctx:
+        # Build agent params - use system_prompt from ctx (set from task.agent_config)
+        final_agent_params = dict(agent_params or {})
+        if ctx.system_prompt and "system_prompt" not in final_agent_params:
+            final_agent_params["system_prompt"] = ctx.system_prompt
+
+        # Create agent inside ctx so it has access to context-derived values
+        agent_cls = agent_type.cls
+        agent = agent_cls.create(**final_agent_params)
+
+        # Store metadata if provided
+        if metadata:
+            ctx.metadata.update(metadata)
+
+        result = await agent.run(ctx, max_steps=max_steps)
+        # Reward is computed by EvalContext.__aexit__ from evaluate tools
+
+    # Return the Trace (ctx.reward is set by EvalContext.__aexit__)
+    return result