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

hud/telemetry/_trace.py

@@ -1,347 +0,0 @@
-from __future__ import annotations
-
-# ruff: noqa: T201
-import asyncio
-import logging
-import time
-import uuid
-from contextlib import contextmanager
-from functools import wraps
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    ParamSpec,
-    TypeVar,
-)
-
-from hud.telemetry.context import (
-    flush_buffer,
-    get_current_task_run_id,
-    is_root_trace,
-    set_current_task_run_id,
-)
-from hud.telemetry.instrumentation.registry import registry
-
-if TYPE_CHECKING:
-    from collections.abc import Generator
-
-
-logger = logging.getLogger("hud.telemetry")
-T = TypeVar("T")
-P = ParamSpec("P")
-
-# Track whether telemetry has been initialized
-_telemetry_initialized = False
-
-
-def init_telemetry() -> None:
-    """Initialize telemetry instrumentors and ensure worker is started if telemetry is active."""
-    global _telemetry_initialized
-    if _telemetry_initialized:
-        return
-
-    registry.install_all()
-    logger.info("Telemetry initialized.")
-    _telemetry_initialized = True
-
-
-def _ensure_telemetry_initialized() -> None:
-    """Ensure telemetry is initialized - called lazily by trace functions."""
-    from hud.settings import settings
-
-    if settings.telemetry_enabled and not _telemetry_initialized:
-        init_telemetry()
-
-
-def _detect_agent_model() -> str | None:
-    """
-    Try to auto-detect agent model from parent frames.
-    This is a best-effort approach and may not work in all cases.
-    """
-    import sys
-
-    try:
-        # Try different frame depths (2-3 typically covers most cases)
-        for depth in range(2, 3):
-            try:
-                frame = sys._getframe(depth)
-                # Check local variables for agent objects
-                for var_value in frame.f_locals.values():
-                    # Look for objects with model_name attribute
-                    if hasattr(var_value, "model_name") and hasattr(var_value, "run"):
-                        # Likely an agent object
-                        model_name = getattr(var_value, "model_name", None)
-                        if model_name:
-                            logger.debug(
-                                "Found agent with model_name in frame %d: %s", depth, model_name
-                            )
-                            return str(model_name)
-
-                # Also check self in case we're in a method
-                if "self" in frame.f_locals:
-                    self_obj = frame.f_locals["self"]
-                    if hasattr(self_obj, "model_name"):
-                        model_name = getattr(self_obj, "model_name", None)
-                        if model_name:
-                            logger.debug(
-                                "Found agent model_name in self at frame %d: %s", depth, model_name
-                            )
-                            return str(model_name)
-
-            except (ValueError, AttributeError):
-                # Frame doesn't exist at this depth or other issues
-                continue
-
-    except Exception as e:
-        logger.debug("Agent model detection failed: %s", e)
-
-    return None
-
-
-def _print_trace_url(task_run_id: str) -> None:
-    """Print the trace URL in a colorful box."""
-    url = f"https://app.hud.so/trace/{task_run_id}"
-    header = "🚀 See your agent live at:"
-
-    # ANSI color codes
-    DIM = "\033[90m"  # Dim/Gray for border (visible on both light and dark terminals)
-    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}")
-    print(
-        f"{DIM}║{RESET}{' ' * header_padding}{header}{' ' * (box_width - len(header) - header_padding - 3)}{DIM}║{RESET}"  # noqa: E501
-    )
-    print(f"{DIM}{divider}{RESET}")
-    print(
-        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")
-
-
-def _print_trace_complete_url(task_run_id: str) -> None:
-    """Print the trace completion URL in a simple colorful format."""
-    url = f"https://app.hud.so/trace/{task_run_id}"
-
-    # ANSI color codes
-    GREEN = "\033[92m"
-    GOLD = "\033[33m"
-    RESET = "\033[0m"
-    DIM = "\033[2m"
-    BOLD = "\033[1m"
-
-    print(f"\n{GREEN}✓ Trace complete!{RESET} {DIM}View at:{RESET} {BOLD}{GOLD}{url}{RESET}\n")
-
-
-@contextmanager
-def trace_open(
-    name: str | None = None,
-    agent_model: str | None = None,
-    run_id: str | None = None,
-    attributes: dict[str, Any] | None = None,
-) -> Generator[str, None, None]:
-    """
-    Context manager for tracing a block of code.
-
-    Args:
-        name: Optional name for this trace, will be added to attributes.
-        attributes: Optional dictionary of attributes to associate with this trace
-
-    Returns:
-        The generated task run ID (UUID string) used for this trace
-    """
-    # Lazy initialization - only initialize telemetry when trace() is actually called
-    _ensure_telemetry_initialized()
-
-    task_run_id = run_id or str(uuid.uuid4())
-
-    _print_trace_url(task_run_id)
-
-    local_attributes = attributes.copy() if attributes is not None else {}
-    if name is not None:
-        local_attributes["trace_name"] = name
-
-    # Auto-detect agent if not explicitly provided
-    if agent_model is None:
-        agent_model = _detect_agent_model()
-
-    start_time = time.time()
-    logger.debug("Starting trace %s (Name: %s)", task_run_id, name if name else "Unnamed")
-
-    previous_task_id = get_current_task_run_id()
-    was_root = is_root_trace.get()
-
-    set_current_task_run_id(task_run_id)
-    is_root = previous_task_id is None
-    is_root_trace.set(is_root)
-
-    # Update status to initializing for root traces
-    if is_root:
-        from hud.telemetry.exporter import (
-            TaskRunStatus,
-            submit_to_worker_loop,
-            update_task_run_status,
-        )
-        from hud.telemetry.job import get_current_job_id
-
-        # Include metadata in the initial status update
-        initial_metadata = local_attributes.copy()
-        initial_metadata["is_root_trace"] = is_root
-        if agent_model:
-            initial_metadata["agent_model"] = agent_model
-
-        # Get job_id if we're in a job context
-        job_id = get_current_job_id()
-
-        coro = update_task_run_status(
-            task_run_id, TaskRunStatus.INITIALIZING, metadata=initial_metadata, job_id=job_id
-        )
-        submit_to_worker_loop(coro)
-        logger.debug("Updated task run %s status to INITIALIZING with metadata", task_run_id)
-
-    error_occurred = False
-    error_message = None
-
-    try:
-        yield task_run_id
-    except Exception as e:
-        error_occurred = True
-        error_message = str(e)
-        raise
-    finally:
-        end_time = time.time()
-        duration = end_time - start_time
-        local_attributes["duration_seconds"] = duration
-        local_attributes["is_root_trace"] = is_root
-
-        logger.debug("Finishing trace %s after %.2f seconds", task_run_id, duration)
-
-        # Update status for root traces
-        if is_root:
-            from hud.telemetry.exporter import (
-                TaskRunStatus,
-                submit_to_worker_loop,
-                update_task_run_status,
-            )
-
-            # Include final metadata with duration
-            final_metadata = local_attributes.copy()
-
-            if error_occurred:
-                coro = update_task_run_status(
-                    task_run_id, TaskRunStatus.ERROR, error_message, metadata=final_metadata
-                )
-                logger.debug("Updated task run %s status to ERROR: %s", task_run_id, error_message)
-            else:
-                coro = update_task_run_status(
-                    task_run_id, TaskRunStatus.COMPLETED, metadata=final_metadata
-                )
-                logger.debug("Updated task run %s status to COMPLETED with metadata", task_run_id)
-
-            # Wait for the status update to complete
-            future = submit_to_worker_loop(coro)
-            if future:
-                try:
-                    # Wait up to 5 seconds for the status update
-                    import concurrent.futures
-
-                    future.result(timeout=5.0)
-                    logger.debug("Status update completed successfully")
-                except concurrent.futures.TimeoutError:
-                    logger.warning("Timeout waiting for status update to complete")
-                except Exception as e:
-                    logger.error("Error waiting for status update: %s", e)
-
-        # Export any remaining records before flushing
-        if is_root:
-            from hud.telemetry.context import export_incremental
-
-            export_incremental()
-
-        # Always flush the buffer for the current task
-        mcp_calls = flush_buffer(export=True)
-        logger.debug("Flushed %d MCP calls for trace %s", len(mcp_calls), task_run_id)
-
-        # Restore previous context
-        set_current_task_run_id(previous_task_id)
-        is_root_trace.set(was_root)
-
-        # Log at the end
-        if is_root:
-            _print_trace_complete_url(task_run_id)
-
-
-@contextmanager
-def trace(
-    name: str | None = None,
-    agent_model: str | None = None,
-    attributes: dict[str, Any] | None = None,
-) -> Generator[str, None, None]:
-    """
-    Synchronous context manager that traces and blocks until telemetry is sent.
-
-    This is the "worry-free" option when you want to ensure telemetry is
-    sent immediately before continuing, rather than relying on background workers.
-
-    Args:
-        name: Optional name for this trace
-        attributes: Optional attributes for the trace
-
-    Returns:
-        The generated task run ID (UUID string) used for this trace
-    """
-    with trace_open(name=name, agent_model=agent_model, attributes=attributes) as task_run_id:
-        yield task_run_id
-
-    # Ensure telemetry is flushed synchronously
-    from hud import flush
-
-    flush()
-
-
-def trace_decorator(
-    name: str | None = None,
-    agent_model: str | None = None,
-    attributes: dict[str, Any] | None = None,
-) -> Any:
-    """
-    Decorator for tracing functions.
-
-    Can be used on both sync and async functions.
-    """
-
-    def decorator(func: Any) -> Any:
-        if asyncio.iscoroutinefunction(func):
-
-            @wraps(func)
-            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-                func_name = name or f"{func.__module__}.{func.__name__}"
-                with trace_open(name=func_name, agent_model=agent_model, attributes=attributes):
-                    return await func(*args, **kwargs)
-
-            return async_wrapper
-        else:
-
-            @wraps(func)
-            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
-                func_name = name or f"{func.__module__}.{func.__name__}"
-                with trace_open(name=func_name, agent_model=agent_model, attributes=attributes):
-                    return func(*args, **kwargs)
-
-            return sync_wrapper
-
-    return decorator

hud/telemetry/context.py

@@ -1,230 +0,0 @@
-from __future__ import annotations
-
-import contextvars
-import logging
-from collections import defaultdict
-from datetime import datetime
-from typing import Any, TypeVar
-
-from hud.telemetry.mcp_models import (
-    BaseMCPCall,
-    MCPNotificationCall,
-    MCPRequestCall,
-    MCPResponseCall,
-    StatusType,
-)
-
-logger = logging.getLogger("hud.telemetry")
-
-# Context variables for tracing
-current_task_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
-    "current_task_run_id", default=None
-)
-# Global dictionary for buffering, keyed by task_run_id
-_GLOBAL_MCP_CALL_BUFFERS: defaultdict[str, list[BaseMCPCall]] = defaultdict(list)
-# Track the last exported index for each task_run_id
-_GLOBAL_EXPORT_INDICES: defaultdict[str, int] = defaultdict(int)
-# Track whether we've seen a non-init request for each task_run_id
-_GLOBAL_HAS_NON_INIT_REQUEST: defaultdict[str, bool] = defaultdict(bool)
-is_root_trace: contextvars.ContextVar[bool] = contextvars.ContextVar("is_root_trace", default=False)
-
-# Maximum buffer size before automatic flush
-MAX_BUFFER_SIZE = 100
-
-# Type variable for record factories
-T = TypeVar("T", bound=BaseMCPCall)
-
-
-def get_current_task_run_id() -> str | None:
-    """Get the task_run_id for the current trace context."""
-    return current_task_run_id.get()
-
-
-def set_current_task_run_id(task_run_id: str | None) -> None:
-    """Set the task_run_id for the current trace context."""
-    current_task_run_id.set(task_run_id)
-
-
-def buffer_mcp_call(record: BaseMCPCall | dict[str, Any]) -> None:
-    """Buffer an MCP call record for the current trace."""
-    task_run_id = get_current_task_run_id()
-
-    if not task_run_id:
-        logger.warning(
-            "BUFFER_MCP_CALL: No task_run_id. Skipping buffer for %s", type(record).__name__
-        )
-        return
-
-    # Ensure 'record' is a Pydantic model instance
-    if isinstance(record, dict):
-        try:
-            record_model = BaseMCPCall.from_dict(record)
-            record = record_model
-        except Exception as e_conv:
-            logger.exception("BUFFER_MCP_CALL: Failed to convert dict to BaseMCPCall: %s", e_conv)
-            return
-
-    _GLOBAL_MCP_CALL_BUFFERS[task_run_id].append(record)
-    buffer_len = len(_GLOBAL_MCP_CALL_BUFFERS[task_run_id])
-
-    if buffer_len >= MAX_BUFFER_SIZE:
-        flush_buffer(export=True)
-
-
-def export_incremental() -> list[BaseMCPCall]:
-    """
-    Export only new MCP calls since last export without clearing the buffer.
-
-    Returns:
-        The list of newly exported MCP calls
-    """
-    task_run_id = get_current_task_run_id()
-    if not task_run_id or not is_root_trace.get():
-        return []
-
-    buffer = _GLOBAL_MCP_CALL_BUFFERS.get(task_run_id, [])
-    last_exported_idx = _GLOBAL_EXPORT_INDICES.get(task_run_id, 0)
-
-    # Get only the new records since last export
-    new_records = buffer[last_exported_idx:]
-
-    if new_records:
-        # Update the export index
-        _GLOBAL_EXPORT_INDICES[task_run_id] = len(buffer)
-
-        # Trigger export
-        from hud.telemetry import exporter
-        from hud.telemetry.exporter import submit_to_worker_loop
-
-        # Get current trace attributes if available
-        attributes = {"incremental": True}
-
-        coro = exporter.export_telemetry(
-            task_run_id=task_run_id,
-            trace_attributes=attributes,
-            mcp_calls=new_records.copy(),  # Copy to avoid modification during export
-        )
-        submit_to_worker_loop(coro)
-
-        logger.debug(
-            "Incremental export: %d new MCP calls for trace %s", len(new_records), task_run_id
-        )
-
-    return new_records
-
-
-def flush_buffer(export: bool = False) -> list[BaseMCPCall]:
-    """
-    Clear the MCP calls buffer and return its contents.
-
-    Args:
-        export: Whether to trigger export of this buffer
-
-    Returns:
-        The list of buffered MCP calls
-    """
-    task_run_id = get_current_task_run_id()
-    if not task_run_id:
-        logger.warning("FLUSH_BUFFER: No current task_run_id. Cannot flush.")
-        return []
-
-    buffer_for_task = _GLOBAL_MCP_CALL_BUFFERS.pop(task_run_id, [])
-    # Clean up export index when buffer is flushed
-    _GLOBAL_EXPORT_INDICES.pop(task_run_id, None)
-    # Clean up non-init request tracking
-    _GLOBAL_HAS_NON_INIT_REQUEST.pop(task_run_id, None)
-    return buffer_for_task
-
-
-def create_request_record(
-    method: str, status: StatusType = StatusType.STARTED, **kwargs: Any
-) -> MCPRequestCall:
-    """Create and buffer a request record"""
-    task_run_id = get_current_task_run_id()
-    if not task_run_id:
-        logger.warning("No active task_run_id, request record will not be created")
-        raise ValueError("No active task_run_id")
-
-    # Check if this is the first non-init request and update status
-    if is_root_trace.get() and not _GLOBAL_HAS_NON_INIT_REQUEST[task_run_id]:
-        # Common initialization method patterns
-        init_methods = {"initialize", "session/new", "init", "setup", "connect"}
-        method_lower = method.lower()
-
-        # Check if this is NOT an initialization method
-        if not any(init_pattern in method_lower for init_pattern in init_methods):
-            _GLOBAL_HAS_NON_INIT_REQUEST[task_run_id] = True
-
-            # Update status to running
-            from hud.telemetry.exporter import (
-                TaskRunStatus,
-                submit_to_worker_loop,
-                update_task_run_status,
-            )
-
-            coro = update_task_run_status(task_run_id, TaskRunStatus.RUNNING)
-            submit_to_worker_loop(coro)
-            logger.debug(
-                "Updated task run %s status to RUNNING on first non-init request: %s",
-                task_run_id,
-                method,
-            )
-
-    record = MCPRequestCall(
-        task_run_id=task_run_id,
-        method=method,
-        status=status,
-        start_time=kwargs.pop("start_time", None) or datetime.now().timestamp(),
-        **kwargs,
-    )
-    buffer_mcp_call(record)
-    return record
-
-
-def create_response_record(
-    method: str, related_request_id: str | int | None = None, is_error: bool = False, **kwargs: Any
-) -> MCPResponseCall:
-    """Create and buffer a response record"""
-    task_run_id = get_current_task_run_id()
-    if not task_run_id:
-        logger.warning("No active task_run_id, response record will not be created")
-        raise ValueError("No active task_run_id")
-
-    # Default to COMPLETED status if not provided
-    if "status" not in kwargs:
-        kwargs["status"] = StatusType.COMPLETED
-
-    record = MCPResponseCall(
-        task_run_id=task_run_id,
-        method=method,
-        related_request_id=related_request_id,
-        is_error=is_error,
-        **kwargs,
-    )
-
-    buffer_mcp_call(record)
-
-    # Trigger incremental export when we receive a response
-    export_incremental()
-
-    return record
-
-
-def create_notification_record(
-    method: str, status: StatusType = StatusType.STARTED, **kwargs: Any
-) -> MCPNotificationCall:
-    """Create and buffer a notification record"""
-    task_run_id = get_current_task_run_id()
-    if not task_run_id:
-        logger.warning("No active task_run_id, notification record will not be created")
-        raise ValueError("No active task_run_id")
-
-    record = MCPNotificationCall(
-        task_run_id=task_run_id,
-        method=method,
-        status=status,
-        start_time=kwargs.pop("start_time", None) or datetime.now().timestamp(),
-        **kwargs,
-    )
-    buffer_mcp_call(record)
-    return record