hud-python 0.4.14__py3-none-any.whl → 0.4.16__py3-none-any.whl

hud/datasets/execution/runner.py

@@ -0,0 +1,123 @@
+"""Standard asyncio-based dataset runner."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any, cast
+
+from datasets import Dataset, load_dataset
+
+from hud.agents.misc import ResponseAgent
+from hud.datasets.task import Task
+
+if TYPE_CHECKING:
+    from hud.agents import MCPAgent
+
+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 = 50,
+    metadata: 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.
+
+    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
+
+    Returns:
+        List of results from agent.run() in dataset order
+
+    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
+                with hud.trace(task_name, job_id=job_obj.id, task_id=task_dict.get("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
+        )
+
+    return results

hud/datasets/task.py

@@ -0,0 +1,107 @@
+"""Task model for HUD datasets."""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from string import Template
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator
+
+from hud.settings import settings
+from hud.types import MCPToolCall
+
+
+class Task(BaseModel):
+    """
+    A task configuration that can be used to create a task.
+
+    The mcp_config field supports environment variable substitution using
+    template placeholders in the format ${VAR_NAME} or ${VAR_NAME:default_value}.
+
+    Example:
+        mcp_config: {
+            "hud": {
+                "url": "${HUD_MCP_URL:https://mcp.hud.so/v3/mcp}",
+                "headers": {
+                    "Authorization": "Bearer ${HUD_API_KEY}",
+                    "Mcp-Image": "your-mcp-image"
+                }
+            }
+        }
+    """
+
+    id: str | None = None
+    prompt: str
+    mcp_config: dict[str, Any]
+    setup_tool: MCPToolCall | list[MCPToolCall] | None = None
+    evaluate_tool: MCPToolCall | list[MCPToolCall] | None = None
+    system_prompt: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    @field_validator("mcp_config", "metadata", mode="before")
+    @classmethod
+    def parse_json_strings(cls, v: Any) -> Any:
+        """Parse JSON strings into dictionaries."""
+        if isinstance(v, str):
+            try:
+                return json.loads(v)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON string: {e}") from e
+        return v
+
+    @field_validator("setup_tool", "evaluate_tool", mode="before")
+    @classmethod
+    def convert_dict_to_tool_call(cls, v: Any) -> Any:
+        """Convert dict to MCPToolCall instance, parsing JSON strings first."""
+        if v is None:
+            return None
+
+        # Parse JSON string if needed
+        if isinstance(v, str):
+            try:
+                v = json.loads(v)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON string: {e}") from e
+
+        if isinstance(v, dict):
+            return MCPToolCall(**v)
+        if isinstance(v, list):
+            return [MCPToolCall(**item) if isinstance(item, dict) else item for item in v]
+        return v
+
+    @field_validator("mcp_config", mode="before")
+    @classmethod
+    def resolve_env_vars(cls, v: dict[str, Any]) -> dict[str, Any]:
+        """
+        Automatically resolve environment variables in mcp_config using Template.
+
+        Supports ${VAR_NAME} syntax with variable substitution from
+        System environment variables (including HUD_API_KEY, etc.)
+
+        Missing variables resolve to empty strings.
+        """
+        import os
+
+        # Start with current environment variables
+        mapping = dict(os.environ)
+        mapping.update(settings.model_dump())
+
+        if settings.api_key:
+            mapping["HUD_API_KEY"] = settings.api_key
+
+        def substitute_in_value(obj: Any) -> Any:
+            """Recursively substitute variables in nested structures."""
+            if isinstance(obj, str):
+                # Use Template's substitute with defaultdict - missing vars become empty strings
+                safe_mapping = defaultdict(str, mapping)
+                return Template(obj).substitute(safe_mapping)
+            elif isinstance(obj, dict):
+                return {k: substitute_in_value(v) for k, v in obj.items()}
+            elif isinstance(obj, list):
+                return [substitute_in_value(item) for item in obj]
+            else:
+                return obj
+
+        return substitute_in_value(v)

hud/datasets/utils.py

@@ -0,0 +1,118 @@
+"""Dataset utilities for loading, saving, and fetching datasets."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from datasets import Dataset
+
+from .task import Task
+
+logger = logging.getLogger("hud.datasets")
+
+
+async def fetch_system_prompt_from_dataset(dataset_id: str) -> str | None:
+    """
+    Fetch system_prompt.txt from a HuggingFace dataset repository.
+
+    Args:
+        dataset_id: HuggingFace dataset identifier (e.g., "hud-evals/SheetBench-50")
+
+    Returns:
+        System prompt text if found, None otherwise
+    """
+    try:
+        # Import here to avoid unnecessary dependency
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.errors import EntryNotFoundError
+
+        # Try to download the system_prompt.txt file
+        try:
+            file_path = hf_hub_download(
+                repo_id=dataset_id, filename="system_prompt.txt", repo_type="dataset"
+            )
+
+            # Read and return the content
+            with open(file_path, encoding="utf-8") as f:  # noqa: ASYNC230
+                content = f.read().strip()
+                if content:
+                    logger.info(
+                        "Loaded system prompt from %s (length: %d chars)", dataset_id, len(content)
+                    )
+                    return content
+                else:
+                    logger.warning("System prompt file is empty in %s", dataset_id)
+                    return None
+
+        except EntryNotFoundError:
+            logger.debug("No system_prompt.txt found in dataset %s", dataset_id)
+            return None
+
+    except ImportError:
+        logger.warning(
+            "huggingface_hub not installed. Install it to fetch system prompts from datasets."
+        )
+        return None
+    except Exception as e:
+        logger.error("Error fetching system prompt from %s: %s", dataset_id, e)
+        return None
+
+
+def save_tasks(
+    tasks: list[dict[str, Any]], repo_id: str, fields: list[str] | None = None, **kwargs: Any
+) -> None:
+    """
+    Save data to HuggingFace dataset with JSON string serialization.
+
+    Complex fields (dicts, lists) are serialized as JSON strings to maintain clean schema
+    and avoid null value pollution in HuggingFace datasets.
+
+    Args:
+        tasks: List of dictionaries to save
+        repo_id: HuggingFace repository ID (e.g., "hud-evals/my-tasks")
+        fields: Optional list of fields to save. If None, saves all fields from each dict.
+        **kwargs: Additional arguments passed to dataset.push_to_hub()
+    """
+    # Safety check: Ensure we're not saving Task objects (which have resolved env vars)
+    if tasks and isinstance(tasks[0], Task):
+        raise ValueError(
+            "save_tasks expects dictionaries, not Task objects. "
+            "Task objects have resolved environment variables which would expose secrets. "
+            "Please pass raw dictionaries with template strings like '${HUD_API_KEY}' preserved."
+        )
+
+    # Convert to rows with JSON string fields
+    data = []
+    for i, tc_dict in enumerate(tasks):
+        # Additional safety check for each item
+        if isinstance(tc_dict, Task):
+            raise ValueError(
+                f"Item {i} is a Task object, not a dictionary. "
+                "This would expose resolved environment variables. "
+                "Please convert to dictionary format with template strings preserved."
+            )
+
+        row = {}
+
+        # Determine which fields to process
+        fields_to_process = fields if fields is not None else list(tc_dict.keys())
+
+        for field in fields_to_process:
+            if field in tc_dict:
+                value = tc_dict[field]
+                # Serialize complex types as JSON strings
+                if isinstance(value, (dict | list)):
+                    row[field] = json.dumps(value)
+                elif isinstance(value, (str | int | float | bool | type(None))):
+                    row[field] = value if value is not None else ""
+                else:
+                    # For other types, convert to string
+                    row[field] = str(value)
+
+        data.append(row)
+
+    # Create and push dataset
+    dataset = Dataset.from_list(data)
+    dataset.push_to_hub(repo_id, **kwargs)

hud/otel/instrumentation.py

@@ -32,8 +32,9 @@ def install_mcp_instrumentation(provider: TracerProvider) -> None:
     try:
         # First, patch the _instruments to use our fork
         import opentelemetry.instrumentation.mcp.instrumentation as mcp_inst
+
         mcp_inst._instruments = ("hud-mcp-python-sdk >= 3.13.1",)
-        
+
         from opentelemetry.instrumentation.mcp.instrumentation import (
             McpInstrumentor,
         )

hud/server/server.py

@@ -33,6 +33,8 @@ def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) ->
     """Run *coro_fn* via anyio.run() and cancel on SIGTERM or SIGINT (POSIX)."""
     global _sigterm_received
 
+    sys.stderr.flush()
+
     async def _runner() -> None:
         stop_evt: asyncio.Event | None = None
         if sys.platform != "win32" and os.getenv("FASTMCP_DISABLE_SIGTERM_HANDLER") != "1":
@@ -43,28 +45,46 @@ def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) ->
             def handle_sigterm() -> None:
                 global _sigterm_received
                 _sigterm_received = True
-                logger.info("Received SIGTERM signal")
+                logger.info("Received SIGTERM signal, setting shutdown flag")
+                stop_evt.set()
+
+            # Handle SIGINT for hot-reload
+            def handle_sigint() -> None:
+                logger.info("Received SIGINT signal, triggering hot reload...")
+                # Don't set _sigterm_received for SIGINT
                 stop_evt.set()
 
             # Handle both SIGTERM and SIGINT for graceful shutdown
-            if signal.getsignal(signal.SIGTERM) is signal.SIG_DFL:
+            # In Docker containers, we always want to register our handlers
+            try:
                 loop.add_signal_handler(signal.SIGTERM, handle_sigterm)
-            if signal.getsignal(signal.SIGINT) is signal.SIG_DFL:
-                loop.add_signal_handler(signal.SIGINT, stop_evt.set)
-
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(coro_fn, *args, **kwargs)
-
-            if stop_evt is not None:
-
-                async def _watch() -> None:
-                    logger.info("Waiting for SIGTERM or SIGINT")
-                    if stop_evt is not None:
-                        await stop_evt.wait()
-                    logger.debug("Received shutdown signal, cancelling tasks...")
-                    tg.cancel_scope.cancel()
-
-                tg.start_soon(_watch)
+                logger.info("SIGTERM handler registered")
+            except (ValueError, OSError) as e:
+                logger.warning("Could not register SIGTERM handler: %s", e)
+
+            try:
+                loop.add_signal_handler(signal.SIGINT, handle_sigint)
+                logger.info("SIGINT handler registered")
+            except (ValueError, OSError) as e:
+                logger.warning("Could not register SIGINT handler: %s", e)
+
+        try:
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(coro_fn, *args, **kwargs)
+
+                if stop_evt is not None:
+
+                    async def _watch() -> None:
+                        logger.info("Signal handler ready, waiting for SIGTERM or SIGINT")
+                        if stop_evt is not None:
+                            await stop_evt.wait()
+                        logger.info("Shutdown signal received, initiating graceful shutdown...")
+                        tg.cancel_scope.cancel()
+
+                    tg.start_soon(_watch)
+        except* asyncio.CancelledError:
+            # This ensures the task group cleans up properly
+            logger.info("Task group cancelled, cleaning up...")
 
     anyio.run(_runner)
 
@@ -101,12 +121,29 @@ class MCPServer(FastMCP):
                     yield {}
                 finally:
                     # Only call shutdown handler if SIGTERM was received
+                    logger.info("Lifespan `finally` block reached. Checking for SIGTERM.")
+                    # Force flush logs to ensure they're visible
+                    sys.stderr.flush()
+
                     if self._shutdown_fn is not None and _sigterm_received:
-                        logger.info("SIGTERM received, calling shutdown handler")
-                        await self._shutdown_fn()
+                        logger.info("SIGTERM detected! Calling @mcp.shutdown handler...")
+                        sys.stderr.flush()
+                        try:
+                            await self._shutdown_fn()
+                            logger.info("@mcp.shutdown handler completed successfully.")
+                            sys.stderr.flush()
+                        except Exception as e:
+                            logger.error("Error during @mcp.shutdown: %s", e)
+                            sys.stderr.flush()
                         _sigterm_received = False
                     elif self._shutdown_fn is not None:
-                        logger.debug("Normal shutdown (hot reload), skipping shutdown handler")
+                        logger.info(
+                            "No SIGTERM. This is a hot reload (SIGINT) or normal exit. Skipping @mcp.shutdown handler."  # noqa: E501
+                        )
+                        sys.stderr.flush()
+                    else:
+                        logger.info("No shutdown handler registered.")
+                        sys.stderr.flush()
 
             fastmcp_kwargs["lifespan"] = _lifespan
 

hud/settings.py

@@ -44,6 +44,18 @@ class Settings(BaseSettings):
         validation_alias="OPENAI_API_KEY",
     )
 
+    wandb_api_key: str | None = Field(
+        default=None,
+        description="API key for Weights & Biases",
+        validation_alias="WANDB_API_KEY",
+    )
+
+    prime_api_key: str | None = Field(
+        default=None,
+        description="API key for Prime Intellect",
+        validation_alias="PRIME_API_KEY",
+    )
+
     telemetry_enabled: bool = Field(
         default=True,
         description="Enable telemetry for the HUD SDK",

hud/types.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import json
 import uuid
 from typing import Any, Literal
 
+import mcp.types as types
 from mcp.types import CallToolRequestParams, CallToolResult
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -13,22 +15,41 @@ class MCPToolCall(CallToolRequestParams):
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))  # Unique identifier for reference
 
     def __str__(self) -> str:
-        response = f"Tool: {self.name}"
-        if self.arguments:
-            response += f"\nArguments: {self.arguments}"
-        return response
+        """Format tool call with Rich markup for HUD design."""
+        from hud.utils.design import design
+
+        return design.format_tool_call(self.name, self.arguments)
 
 
 class MCPToolResult(CallToolResult):
     """A tool result."""
 
     def __str__(self) -> str:
-        response = f"Content: {self.content}"
-        if self.structuredContent:
-            response += f"\nStructured Content: {self.structuredContent}"
-        if self.isError:
-            response += f"\nError: {self.isError}"
-        return response
+        """Format tool result with Rich markup for HUD design - compact version."""
+        from hud.utils.design import design
+
+        # Extract content summary
+        content_summary = ""
+        if self.content:
+            for block in self.content:
+                if isinstance(block, types.TextContent):
+                    # Get first line or truncate
+                    text = block.text.strip()
+                    first_line = text.split("\n")[0] if "\n" in text else text
+                    content_summary = first_line
+                    break
+                elif isinstance(block, types.ImageContent):
+                    content_summary = "📷 Image"
+                    break
+
+        # Or use structured content if no text content
+        if not content_summary and self.structuredContent:
+            try:
+                content_summary = json.dumps(self.structuredContent, separators=(",", ":"))
+            except (TypeError, ValueError):
+                content_summary = str(self.structuredContent)
+
+        return design.format_tool_result(content_summary, self.isError)
 
 
 class AgentResponse(BaseModel):