hud-python 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

hud/telemetry/job.py

@@ -0,0 +1,141 @@
+"""Job context manager for grouping related traces."""
+
+from __future__ import annotations
+
+import logging
+import sys
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+from hud.telemetry.exporter import JobStatus, submit_to_worker_loop, update_job_status
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Self
+
+logger = logging.getLogger("hud.telemetry")
+
+# Context variables for current job
+current_job_id: ContextVar[str | None] = ContextVar("current_job_id", default=None)
+current_job_name: ContextVar[str | None] = ContextVar("current_job_name", default=None)
+
+
+class JobContext:
+    """Context manager for grouping traces under a job."""
+
+    def __init__(
+        self, name: str, taskset_name: str | None = None, metadata: dict[str, Any] | None = None
+    ) -> None:
+        self.id = str(uuid.uuid4())
+        self.name = name
+        self.metadata = metadata or {}
+        self.taskset_name: str | None = taskset_name
+
+    def __enter__(self) -> Self:
+        # Auto-detect dataset
+        if self.taskset_name is None:
+            self._detect_dataset()
+
+        # Set context variables
+        current_job_id.set(self.id)
+        current_job_name.set(self.name)
+
+        # Send initial status
+        job_metadata = {**self.metadata}
+        coro = update_job_status(
+            self.id, JobStatus.RUNNING, metadata=job_metadata, taskset_name=self.taskset_name
+        )
+        submit_to_worker_loop(coro)
+
+        logger.info("Started job %s (ID: %s)", self.name, self.id)
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
+    ) -> None:
+        # Determine final status
+        if exc_type is not None:
+            # Job failed with exception
+            error_msg = f"{exc_type.__name__}: {exc_val}"
+            coro = update_job_status(
+                self.id, JobStatus.ERROR, error_message=error_msg, taskset_name=self.taskset_name
+            )
+        else:
+            # Job completed successfully
+            coro = update_job_status(self.id, JobStatus.COMPLETED, taskset_name=self.taskset_name)
+
+        submit_to_worker_loop(coro)
+
+        # Clear context
+        current_job_id.set(None)
+        current_job_name.set(None)
+
+        status = "failed" if exc_type else "completed"
+        logger.info("Job %s %s", self.name, status)
+
+    def _detect_dataset(self) -> None:
+        """Auto-detect HuggingFace dataset in parent scope."""
+        try:
+            # Check frames 2 and 3 (with statement and parent scope)
+            for frame_depth in [2, 3]:
+                try:
+                    frame = sys._getframe(frame_depth)
+
+                    # Search for Dataset objects
+                    for var_value in frame.f_locals.values():
+                        if hasattr(var_value, "info") and hasattr(var_value.info, "builder_name"):
+                            self.taskset_name = var_value.info.builder_name
+                            logger.debug(
+                                "Auto-detected dataset at frame %d: %s",
+                                frame_depth,
+                                self.taskset_name,
+                            )
+                            return
+                        elif hasattr(var_value, "builder_name"):
+                            # Older dataset format
+                            self.taskset_name = var_value.builder_name
+                            logger.debug(
+                                "Auto-detected dataset at frame %d: %s",
+                                frame_depth,
+                                self.taskset_name,
+                            )
+                            return
+                except ValueError:
+                    # Frame doesn't exist
+                    continue
+        except Exception as e:
+            logger.debug("Dataset auto-detection failed: %s", e)
+
+
+@contextmanager
+def job(
+    name: str, taskset_name: str | None = None, metadata: dict[str, Any] | None = None
+) -> Generator[JobContext, None, None]:
+    """
+    Create a job context for grouping related traces.
+
+    Args:
+        name: Name for the job
+        metadata: Optional metadata to include with the job
+
+    Example:
+        with hud.job("evaluation_run") as job:
+            for task in tasks:
+                with hud.trace(f"task_{task.id}"):
+                    # Trace automatically includes job_id
+                    result = await agent.run(task)
+    """
+    with JobContext(name, taskset_name, metadata) as ctx:
+        yield ctx
+
+
+def get_current_job_id() -> str | None:
+    """Get the current job ID if inside a job context."""
+    return current_job_id.get()
+
+
+def get_current_job_name() -> str | None:
+    """Get the current job name if inside a job context."""
+    return current_job_name.get()

hud/telemetry/tests/test_trace.py

@@ -48,7 +48,7 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         initial_root_state = actual_is_root_trace.get()
@@ -62,7 +62,8 @@ class TestTrace:
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() == initial_root_state
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_not_called()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
 
     def test_trace_with_name_and_attributes(self, mocker):
         """Test trace with name and attributes, checking they are passed on."""
@@ -71,7 +72,7 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         trace_name = "test_trace_with_data"
@@ -81,7 +82,8 @@ class TestTrace:
             assert isinstance(task_run_id, str)
 
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_called_once()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
 
     @pytest.mark.asyncio
     async def test_trace_with_mcp_calls_exports(self, mocker):
@@ -91,14 +93,14 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         async def mock_export(*args, **kwargs):
             return None
 
-        mock_export_actual_coro = mocker.patch(
-            "hud.telemetry._trace.exporter.export_telemetry",
+        mocker.patch(
+            "hud.telemetry.exporter.export_telemetry",
             side_effect=mock_export,
         )
 
@@ -109,16 +111,14 @@ class TestTrace:
             pass
 
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_called_once()
+        # submit_to_worker_loop is now called for status updates and export
+        # The exact count may vary depending on whether export_incremental is called
+        assert mock_submit_loop.call_count >= 2  # At least INITIALIZING and COMPLETED
 
-        mock_export_actual_coro.assert_called_once()
-        args, kwargs = mock_export_actual_coro.call_args
-        assert kwargs["task_run_id"] == task_run_id
-        assert kwargs["mcp_calls"] == mock_mcp_calls
-        assert kwargs["trace_attributes"]["trace_name"] == test_name
-        assert kwargs["trace_attributes"]["custom_attr"] == "test_val"
-        assert "duration_seconds" in kwargs["trace_attributes"]
-        assert kwargs["trace_attributes"]["is_root_trace"] is True
+        # With the new export flow, export_telemetry is submitted to worker loop
+        # so we can't directly assert on it being called synchronously
+        # Instead, verify that the trace completed successfully
+        assert task_run_id is not None
 
     def test_trace_nested(self, mocker):
         """Test nested traces, verifying context restoration and root trace logic."""
@@ -129,7 +129,7 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop_internal = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         assert actual_get_current_task_run_id() is None
@@ -148,7 +148,8 @@ class TestTrace:
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() is False
         assert mock_flush_internal.call_count == 2
-        mock_submit_loop_internal.assert_not_called()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop_internal.call_count == 2  # Only outer trace sends status updates
 
     def test_trace_exception_handling(self, mocker):
         """Test trace handles exceptions properly and restores context."""
@@ -161,7 +162,7 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         with (
@@ -191,7 +192,7 @@ class TestTraceSync:
         with trace(name="test_sync") as task_run_id:
             assert task_run_id == "test-task-id"
 
-        mock_trace_open.assert_called_once_with(name="test_sync", attributes=None)
+        mock_trace_open.assert_called_once_with(name="test_sync", agent_model=None, attributes=None)
         mock_flush.assert_called_once()
 
     def test_trace_sync_with_attributes(self, mocker):
@@ -205,7 +206,9 @@ class TestTraceSync:
         with trace(name="test_sync", attributes=attrs):
             pass
 
-        mock_trace_open.assert_called_once_with(name="test_sync", attributes=attrs)
+        mock_trace_open.assert_called_once_with(
+            name="test_sync", agent_model=None, attributes=attrs
+        )
         mock_flush.assert_called_once()
 
 
@@ -224,7 +227,9 @@ class TestTraceDecorator:
 
         result = sync_function(1, 2)
         assert result == 3
-        mock_trace_open.assert_called_once_with(name="test_func_sync", attributes=None)
+        mock_trace_open.assert_called_once_with(
+            name="test_func_sync", agent_model=None, attributes=None
+        )
 
     def test_trace_decorator_async_function(self, mocker):
         """Test trace_decorator on asynchronous functions."""
@@ -239,7 +244,9 @@ class TestTraceDecorator:
         async def run_test():
             result = await async_function(1, 2)
             assert result == 3
-            mock_trace_open.assert_called_once_with(name="test_func_async", attributes=None)
+            mock_trace_open.assert_called_once_with(
+                name="test_func_async", agent_model=None, attributes=None
+            )
 
         asyncio.run(run_test())
 
@@ -257,7 +264,9 @@ class TestTraceDecorator:
 
         result = func_with_attrs(5)
         assert result == 10
-        mock_trace_open.assert_called_once_with(name="test_func", attributes=attrs)
+        mock_trace_open.assert_called_once_with(
+            name="test_func", agent_model=None, attributes=attrs
+        )
 
     def test_trace_decorator_without_name(self, mocker):
         """Test trace_decorator uses module.function name when name not provided."""
@@ -273,7 +282,9 @@ class TestTraceDecorator:
         assert result == "result"
         # Should use module.function name
         expected_name = f"{my_function.__module__}.my_function"
-        mock_trace_open.assert_called_once_with(name=expected_name, attributes=None)
+        mock_trace_open.assert_called_once_with(
+            name=expected_name, agent_model=None, attributes=None
+        )
 
     def test_trace_decorator_preserves_function_metadata(self):
         """Test trace_decorator preserves original function metadata."""

hud/tools/__init__.py

@@ -2,12 +2,16 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any
+
 from .base import ToolError, ToolResult, tool_result_to_content_blocks
 from .bash import BashTool
-from .computer import AnthropicComputerTool, HudComputerTool, OpenAIComputerTool
 from .edit import EditTool
 from .playwright_tool import PlaywrightTool
 
+if TYPE_CHECKING:
+    from .computer import AnthropicComputerTool, HudComputerTool, OpenAIComputerTool
+
 __all__ = [
     "AnthropicComputerTool",
     "BashTool",
@@ -19,3 +23,12 @@ __all__ = [
     "ToolResult",
     "tool_result_to_content_blocks",
 ]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy import computer tools to avoid importing pyautogui unless needed."""
+    if name in ("AnthropicComputerTool", "HudComputerTool", "OpenAIComputerTool"):
+        from . import computer
+
+        return getattr(computer, name)
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

hud/tools/executors/__init__.py

@@ -2,12 +2,29 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any
+
 from .base import BaseExecutor
-from .pyautogui import PyAutoGUIExecutor
-from .xdo import XDOExecutor
+
+if TYPE_CHECKING:
+    from .pyautogui import PyAutoGUIExecutor
+    from .xdo import XDOExecutor
 
 __all__ = [
     "BaseExecutor",
     "PyAutoGUIExecutor",
     "XDOExecutor",
 ]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy import executors to avoid importing pyautogui unless needed."""
+    if name == "PyAutoGUIExecutor":
+        from .pyautogui import PyAutoGUIExecutor
+
+        return PyAutoGUIExecutor
+    elif name == "XDOExecutor":
+        from .xdo import XDOExecutor
+
+        return XDOExecutor
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")