hud-python 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

hud/telemetry/job.py

@@ -1,141 +1,309 @@
-"""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()
+"""Job management for HUD SDK.
+
+This module provides APIs for managing jobs - logical groupings of related tasks.
+Jobs can be used to track experiments, batch processing, training runs, etc.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from contextlib import contextmanager
+from datetime import UTC, datetime
+from functools import wraps
+from typing import TYPE_CHECKING, Any
+
+from hud.settings import settings
+from hud.shared import make_request, make_request_sync
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator
+
+logger = logging.getLogger(__name__)
+
+
+class Job:
+    """A job represents a collection of related tasks."""
+
+    def __init__(
+        self,
+        job_id: str,
+        name: str,
+        metadata: dict[str, Any] | None = None,
+        dataset_link: str | None = None,
+    ) -> None:
+        self.id = job_id
+        self.name = name
+        self.metadata = metadata or {}
+        self.dataset_link = dataset_link
+        self.status = "created"
+        self.created_at = datetime.now(UTC)
+        self.tasks: list[str] = []
+
+    def add_task(self, task_id: str) -> None:
+        """Associate a task with this job."""
+        self.tasks.append(task_id)
+
+    async def update_status(self, status: str) -> None:
+        """Update job status on the server."""
+        self.status = status
+        if settings.telemetry_enabled:
+            try:
+                payload = {
+                    "name": self.name,
+                    "status": status,
+                    "metadata": self.metadata,
+                }
+                if self.dataset_link:
+                    payload["dataset_link"] = self.dataset_link
+
+                await make_request(
+                    method="POST",
+                    url=f"{settings.hud_telemetry_url}/jobs/{self.id}/status",
+                    json=payload,
+                    api_key=settings.api_key,
+                )
+            except Exception as e:
+                logger.warning("Failed to update job status: %s", e)
+
+    def update_status_sync(self, status: str) -> None:
+        """Synchronously update job status on the server."""
+        self.status = status
+        if settings.telemetry_enabled:
+            try:
+                payload = {
+                    "name": self.name,
+                    "status": status,
+                    "metadata": self.metadata,
+                }
+                if self.dataset_link:
+                    payload["dataset_link"] = self.dataset_link
+
+                make_request_sync(
+                    method="POST",
+                    url=f"{settings.hud_telemetry_url}/jobs/{self.id}/status",
+                    json=payload,
+                    api_key=settings.api_key,
+                )
+            except Exception as e:
+                logger.warning("Failed to update job status: %s", e)
+
+    def __repr__(self) -> str:
+        return f"Job(id={self.id!r}, name={self.name!r}, status={self.status!r})"
+
+
+# Global job registry for the decorator pattern
+_current_job: Job | None = None
+
+
+def _print_job_url(job_id: str, job_name: str) -> None:
+    """Print the job URL in a colorful box."""
+    # Only print HUD URL if HUD telemetry is enabled and has API key
+    if not (settings.telemetry_enabled and settings.api_key):
+        return
+
+    url = f"https://app.hud.so/jobs/{job_id}"
+    header = f"🚀 Job '{job_name}' started:"
+
+    # ANSI color codes
+    DIM = "\033[90m"  # Dim/Gray for border
+    GOLD = "\033[33m"  # Gold/Yellow for URL
+    RESET = "\033[0m"
+    BOLD = "\033[1m"
+
+    # Calculate box width based on the longest line
+    box_width = max(len(url), len(header)) + 6
+
+    # Box drawing characters
+    top_border = "╔" + "═" * (box_width - 2) + "╗"
+    bottom_border = "╚" + "═" * (box_width - 2) + "╝"
+    divider = "╟" + "─" * (box_width - 2) + "╢"
+
+    # Center the content
+    header_padding = (box_width - len(header) - 2) // 2
+    url_padding = (box_width - len(url) - 2) // 2
+
+    # Print the box
+    print(f"\n{DIM}{top_border}{RESET}")  # noqa: T201
+    print(  # noqa: T201
+        f"{DIM}║{RESET}{' ' * header_padding}{header}{' ' * (box_width - len(header) - header_padding - 3)}{DIM}║{RESET}"  # noqa: E501
+    )
+    print(f"{DIM}{divider}{RESET}")  # noqa: T201
+    print(  # noqa: T201
+        f"{DIM}║{RESET}{' ' * url_padding}{BOLD}{GOLD}{url}{RESET}{' ' * (box_width - len(url) - url_padding - 2)}{DIM}║{RESET}"  # noqa: E501
+    )
+    print(f"{DIM}{bottom_border}{RESET}\n")  # noqa: T201
+
+
+def _print_job_complete_url(job_id: str, job_name: str, error_occurred: bool = False) -> None:
+    """Print the job completion URL with appropriate messaging."""
+    # Only print HUD URL if HUD telemetry is enabled and has API key
+    if not (settings.telemetry_enabled and settings.api_key):
+        return
+
+    url = f"https://app.hud.so/jobs/{job_id}"
+
+    # ANSI color codes
+    GREEN = "\033[92m"
+    RED = "\033[91m"
+    GOLD = "\033[33m"
+    RESET = "\033[0m"
+    DIM = "\033[2m"
+    BOLD = "\033[1m"
+
+    if error_occurred:
+        print(  # noqa: T201
+            f"\n{RED}✗ Job '{job_name}' failed!{RESET} {DIM}View details at:{RESET} {BOLD}{GOLD}{url}{RESET}\n"  # noqa: E501
+        )
+    else:
+        print(  # noqa: T201
+            f"\n{GREEN}✓ Job '{job_name}' complete!{RESET} {DIM}View all results at:{RESET} {BOLD}{GOLD}{url}{RESET}\n"  # noqa: E501
+        )
+
+
+def get_current_job() -> Job | None:
+    """Get the currently active job, if any."""
+    return _current_job
+
+
+@contextmanager
+def job(
+    name: str,
+    metadata: dict[str, Any] | None = None,
+    job_id: str | None = None,
+    dataset_link: str | None = None,
+) -> Generator[Job, None, None]:
+    """Context manager for job tracking.
+
+    Groups related tasks together under a single job for tracking and organization.
+
+    Args:
+        name: Human-readable job name
+        metadata: Optional metadata dictionary
+        job_id: Optional job ID (auto-generated if not provided)
+        dataset_link: Optional HuggingFace dataset identifier (e.g. "hud-evals/SheetBench-50")
+
+    Yields:
+        Job: The job object
+
+    Example:
+        with hud.job("training_run", {"model": "gpt-4"}) as job:
+            for epoch in range(10):
+                with hud.trace(f"epoch_{epoch}", job_id=job.id):
+                    train_epoch()
+    """
+    global _current_job
+
+    if not job_id:
+        job_id = str(uuid.uuid4())
+
+    job_obj = Job(job_id, name, metadata, dataset_link)
+
+    # Set as current job
+    old_job = _current_job
+    _current_job = job_obj
+
+    try:
+        # Update status to running synchronously to ensure job is registered before tasks start
+        job_obj.update_status_sync("running")
+        # Print the nice job URL box
+        _print_job_url(job_obj.id, job_obj.name)
+        yield job_obj
+        # Update status to completed synchronously to ensure it completes before process exit
+        job_obj.update_status_sync("completed")
+        # Print job completion message
+        _print_job_complete_url(job_obj.id, job_obj.name, error_occurred=False)
+    except Exception:
+        # Update status to failed synchronously to ensure it completes before process exit
+        job_obj.update_status_sync("failed")
+        # Print job failure message
+        _print_job_complete_url(job_obj.id, job_obj.name, error_occurred=True)
+        raise
+    finally:
+        _current_job = old_job
+
+
+def create_job(
+    name: str, metadata: dict[str, Any] | None = None, dataset_link: str | None = None
+) -> Job:
+    """Create a job without using context manager.
+
+    Useful when you need explicit control over job lifecycle.
+
+    Args:
+        name: Human-readable job name
+        metadata: Optional metadata dictionary
+        dataset_link: Optional HuggingFace dataset identifier (e.g. "hud-evals/SheetBench-50")
+
+    Returns:
+        Job: The created job object
+
+    Example:
+        job = hud.create_job("data_processing")
+        try:
+            for item in items:
+                with hud.trace(f"process_{item.id}", job_id=job.id):
+                    process(item)
+        finally:
+            await job.update_status("completed")
+    """
+    job_id = str(uuid.uuid4())
+    return Job(job_id, name, metadata, dataset_link)
+
+
+def job_decorator(name: str | None = None, **metadata: Any) -> Callable:
+    """Decorator for functions that should be tracked as jobs.
+
+    Args:
+        name: Job name (defaults to function name)
+        **metadata: Additional metadata for the job
+
+    Example:
+        @hud.job_decorator("model_training", model="gpt-4", dataset="v2")
+        async def train_model(config):
+            # This entire function execution is tracked as a job
+            await model.train(config)
+            return model.evaluate()
+    """
+
+    def decorator(func: Callable) -> Callable:
+        job_name = name or func.__name__
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with job(job_name, metadata) as job_obj:
+                # Store job ID in function for access
+                func._current_job_id = job_obj.id
+                try:
+                    return await func(*args, **kwargs)
+                finally:
+                    delattr(func, "_current_job_id")
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with job(job_name, metadata) as job_obj:
+                # Store job ID in function for access
+                func._current_job_id = job_obj.id
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    delattr(func, "_current_job_id")
+
+        # Return appropriate wrapper based on function type
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return decorator
+
+
+# Convenience exports
+__all__ = [
+    "Job",
+    "create_job",
+    "get_current_job",
+    "job",
+    "job_decorator",
+]

hud/telemetry/replay.py

@@ -0,0 +1,74 @@
+"""Trace retrieval and replay functionality.
+
+This module provides APIs to retrieve collected traces for analysis,
+debugging, and replay purposes.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from hud.otel.collector import clear_trace as _clear_trace
+from hud.otel.collector import get_trace as _get_trace
+
+if TYPE_CHECKING:
+    from hud.types import Trace
+
+__all__ = ["clear_trace", "get_trace"]
+
+
+def get_trace(task_run_id: str) -> Trace | None:
+    """Retrieve the collected trace for a task run.
+
+    Returns None if trace collection was disabled or the trace doesn't exist.
+
+    Args:
+        task_run_id: The task run ID to retrieve the trace for
+
+    Returns:
+        Trace object containing all collected steps, or None if not found
+
+    Usage:
+        import hud
+
+        # Run agent with tracing
+        with hud.trace() as task_run_id:
+            agent = MyAgent()
+            result = await agent.run("solve task")
+
+        # Get the trace for analysis
+        trace = hud.get_trace(task_run_id)
+        if trace:
+            print(f"Collected {len(trace.trace)} steps")
+
+            # Analyze agent vs MCP steps
+            agent_steps = [s for s in trace.trace if s.category == "agent"]
+            mcp_steps = [s for s in trace.trace if s.category == "mcp"]
+
+            print(f"Agent steps: {len(agent_steps)}")
+            print(f"MCP steps: {len(mcp_steps)}")
+
+            # Replay or analyze individual steps
+            for step in trace.trace:
+                if step.category == "agent" and step.result:
+                    print(f"Agent: {step.result.get('content') if isinstance(step.result, dict) else step.result}")
+                if step.category == "mcp" and step.request:
+                    print(f"MCP: {step.request.method if hasattr(step.request, 'method') else step.request}")
+    """  # noqa: E501
+    return _get_trace(task_run_id)
+
+
+def clear_trace(task_run_id: str) -> None:
+    """Clear the collected trace for a task run ID.
+
+    Useful for cleaning up memory after processing large traces.
+
+    Args:
+        task_run_id: The task run ID to clear the trace for
+
+    Usage:
+        trace = hud.get_trace(task_run_id)
+        # Process trace...
+        hud.clear_trace(task_run_id)  # Free memory
+    """
+    _clear_trace(task_run_id)

hud/telemetry/trace.py

@@ -0,0 +1,83 @@
+"""User-facing trace context manager for HUD telemetry.
+
+This module provides the simple trace() API that users interact with.
+The actual OpenTelemetry implementation is in hud.otel.
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+from hud.otel import configure_telemetry
+from hud.otel import trace as OtelTrace
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+__all__ = ["trace"]
+
+
+@contextmanager
+def trace(
+    name: str = "Test task from hud",
+    *,
+    root: bool = True,
+    attrs: dict[str, Any] | None = None,
+    job_id: str | None = None,
+    task_id: str | None = None,
+) -> Generator[str, None, None]:
+    """Start a HUD trace context.
+
+    A unique task_run_id is automatically generated for each trace.
+
+    Args:
+        name: Descriptive name for this trace/task
+        root: Whether this is a root trace (updates task status)
+        attrs: Additional attributes to attach to the trace
+        job_id: Optional job ID to associate with this trace
+
+    Yields:
+        str: The auto-generated task run ID
+
+    Usage:
+        import hud
+
+        with hud.trace("My Task") as task_run_id:
+            # Your code here
+            print(f"Running task: {task_run_id}")
+
+        # Or with default name:
+        with hud.trace() as task_run_id:
+            pass
+
+        # Or with job_id:
+        with hud.trace("My Task", job_id="550e8400-e29b-41d4-a716-446655440000") as task_run_id:
+            pass
+    """
+    # Ensure telemetry is configured
+    configure_telemetry()
+
+    # Only generate task_run_id if using HUD backend
+    # For custom OTLP backends, we don't need it
+    from hud.settings import get_settings
+
+    settings = get_settings()
+
+    if settings.telemetry_enabled and settings.api_key:
+        task_run_id = str(uuid.uuid4())
+    else:
+        # Use a placeholder for custom backends
+        task_run_id = "custom-otlp-trace"
+
+    # Delegate to OpenTelemetry implementation
+    with OtelTrace(
+        task_run_id,
+        is_root=root,
+        span_name=name,
+        attributes=attrs or {},
+        job_id=job_id,
+        task_id=task_id,
+    ) as run_id:
+        yield run_id