gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/tracing/attributes.py

@@ -0,0 +1,9 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Semantic attributes for GatewayDK tracing.
+
+DEPRECATED: This module re-exports all constants from semconv.py for backward
+compatibility. New code should import directly from gatewaysdk.tracing.semconv.
+"""
+
+from .semconv import *  # noqa: F401, F403

gatewaysdk/tracing/context.py

@@ -0,0 +1,272 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Unified tracing context for GatewayDK.
+
+This module provides context variables and context managers for ALL tracing
+metadata: call_type, tags, environment. The TracingContextSpanProcessor
+reads from these and injects into spans.
+
+Run/rollout context (run_id, session_id) is read directly from the run module.
+
+Usage::
+
+    import gatewaysdk.tracing as tracing
+
+    # Context manager for ad-hoc tagging
+    with tracing.tracing_context(call_type="agent-call", tags=["gepa", "train"]):
+        response = client.chat.completions.create(...)  # Tagged automatically
+
+    # Decorator for consistent function tagging
+    @tracing.call_type_decorator("judge-call")
+    def my_judge_function():
+        response = client.chat.completions.create(...)  # Tagged as judge-call
+"""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import contextmanager
+from contextvars import ContextVar
+from functools import wraps
+from typing import Iterator, Any, Callable, Generator, List, Optional, TypeVar
+
+# Type variable for decorators
+F = TypeVar("F", bound=Callable[..., Any])
+
+# =============================================================================
+# Context Variables for Tracing Metadata
+# =============================================================================
+
+_call_type_context: ContextVar[Optional[str]] = ContextVar("gateway_call_type", default=None)
+_tags_context: ContextVar[Optional[List[str]]] = ContextVar("gateway_tags", default=None)
+_environment_context: ContextVar[Optional[str]] = ContextVar("gateway_environment", default=None)
+# Proxy-specific session context. Used by RolloutAttemptMiddleware in llm_proxy
+# to set the session_id for proxy spans (where there's no Run/Rollout object).
+_proxy_session_context: ContextVar[Optional[str]] = ContextVar("gateway_proxy_session", default=None)
+# Context Hub items (prompts/agents/skills/memory) used in this execution
+# context. Stored as an immutable tuple of dicts with keys: kind,
+# repo_handle, commit_hash, label. Appended automatically by
+# GatewayContext.pull(); injected into every span by
+# TracingContextSpanProcessor so traces/sessions link to the content
+# they used.
+_context_items_context: ContextVar[Optional[tuple]] = ContextVar(
+    "gateway_context_items", default=None
+)
+
+
+# =============================================================================
+# Getters
+# =============================================================================
+
+
+def get_current_call_type() -> Optional[str]:
+    """Get the current LLM call type (e.g., 'agent-call', 'judge-call').
+
+    Returns:
+        The current call type string, or None if not set.
+    """
+    return _call_type_context.get()
+
+
+def get_current_tags() -> Optional[List[str]]:
+    """Get the current tracing tags.
+
+    Returns:
+        The current list of tags, or None if not set.
+    """
+    return _tags_context.get()
+
+
+def get_current_environment() -> Optional[str]:
+    """Get the current tracing environment.
+
+    Returns:
+        The current environment string, or None if not set.
+    """
+    return _environment_context.get()
+
+
+def get_current_context_items() -> tuple:
+    """Get the Context Hub items recorded in this execution context.
+
+    Returns:
+        Tuple of dicts with keys: kind, repo_handle, commit_hash, label.
+    """
+    return _context_items_context.get() or ()
+
+
+def record_context_item(
+    kind: str,
+    repo_handle: Optional[str],
+    commit_hash: Optional[str] = None,
+    label: Optional[str] = None,
+) -> None:
+    """Record that a Context Hub item is being used.
+
+    Called automatically by ``GatewayContext.pull()`` so that all spans
+    created afterwards carry ``gateway.context.*`` attributes linking the
+    trace (and its session) to the prompt/agent/skill/memory content used.
+    Deduplicates on (kind, repo_handle, commit_hash).
+    """
+    if not kind or not repo_handle:
+        return
+    item = {
+        "kind": str(kind),
+        "repo_handle": str(repo_handle),
+        "commit_hash": str(commit_hash) if commit_hash else None,
+        "label": str(label) if label else None,
+    }
+    existing = _context_items_context.get() or ()
+    for candidate in existing:
+        if (
+            candidate.get("kind") == item["kind"]
+            and candidate.get("repo_handle") == item["repo_handle"]
+            and candidate.get("commit_hash") == item["commit_hash"]
+        ):
+            return
+    _context_items_context.set(existing + (item,))
+
+
+def get_proxy_session_id() -> Optional[str]:
+    """Get the proxy-specific session ID.
+
+    This is set by RolloutAttemptMiddleware in the LLM proxy when there's
+    no Run/Rollout object in scope. Falls back to None, in which case
+    TracingContextSpanProcessor won't set session_id on proxy spans.
+
+    Returns:
+        The proxy session ID string (typically the rollout_id), or None.
+    """
+    return _proxy_session_context.get()
+
+
+# =============================================================================
+# Context Managers
+# =============================================================================
+
+
+@contextmanager
+def tracing_context(
+    *,
+    call_type: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    environment: Optional[str] = None,
+    session_id: Optional[str] = None,
+) -> Generator[None, None, None]:
+    """Set tracing context for the duration of the block.
+
+    All spans created within this block will be tagged with the provided
+    metadata via the TracingContextSpanProcessor.
+
+    Args:
+        call_type: LLM call type (e.g., "agent-call", "judge-call", "reflection-call")
+        tags: List of tags for the traces
+        environment: Environment name (e.g., "gatewaysdk-gepa")
+        session_id: Session ID for proxy spans (set by RolloutAttemptMiddleware)
+
+    Example::
+
+        with tracing_context(call_type="agent-call", tags=["gepa", "train"]):
+            response = client.chat.completions.create(...)  # Tagged automatically
+    """
+    tokens: List[tuple[str, Any]] = []
+
+    if call_type is not None:
+        tokens.append(("call_type", _call_type_context.set(call_type)))
+    if tags is not None:
+        tokens.append(("tags", _tags_context.set(tags)))
+    if environment is not None:
+        tokens.append(("environment", _environment_context.set(environment)))
+    if session_id is not None:
+        tokens.append(("session_id", _proxy_session_context.set(session_id)))
+
+    try:
+        yield
+    finally:
+        for name, token in tokens:
+            if name == "call_type":
+                _call_type_context.reset(token)
+            elif name == "tags":
+                _tags_context.reset(token)
+            elif name == "environment":
+                _environment_context.reset(token)
+            elif name == "session_id":
+                _proxy_session_context.reset(token)
+
+
+@contextmanager
+def session(session_id: str) -> Iterator[None]:
+    """Group every trace recorded inside this block under one session id.
+
+    THE way to record a named session — e.g. a session you intend to
+    replay later with gatewaysdk.replay, or a logical conversation spanning
+    several traces:
+
+        import gatewaysdk.tracing as tracing
+
+        with tracing.session("checkout-regression-001"):
+            with tracing.trace("turn-1", input=prompt):
+                ...   # every span in here carries the session id
+
+    Works process-wide via context propagation, so auto-instrumented LLM
+    calls and tool spans inside the block are grouped too. Nesting: the
+    innermost session wins.
+    """
+    with tracing_context(session_id=session_id):
+        yield
+
+
+# =============================================================================
+# Decorators
+# =============================================================================
+
+
+def call_type_decorator(call_type: str) -> Callable[[F], F]:
+    """Factory to create call-type tagging decorators.
+
+    Creates a decorator that sets the call_type context for the duration
+    of the decorated function. All spans created within the function
+    will have the `gateway.call_type` attribute set.
+
+    Args:
+        call_type: The call type to set (e.g., "agent-call", "judge-call")
+
+    Returns:
+        A decorator that wraps functions with tracing context.
+
+    Example::
+
+        # Create reusable decorators
+        agent = call_type_decorator("agent-call")
+        judge = call_type_decorator("judge-call")
+
+        @agent
+        def my_agent_function():
+            # LLM calls here are tagged as "agent-call"
+            response = client.chat.completions.create(...)
+
+        @judge
+        def my_judge_function():
+            # LLM calls here are tagged as "judge-call"
+            response = client.chat.completions.create(...)
+    """
+
+    def decorator(fn: F) -> F:
+        if asyncio.iscoroutinefunction(fn):
+
+            @wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with tracing_context(call_type=call_type):
+                    return await fn(*args, **kwargs)
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with tracing_context(call_type=call_type):
+                    return fn(*args, **kwargs)
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    return decorator

gatewaysdk/tracing/exporters/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Exporters for GatewayDK tracing."""
+
+from .gateway import GatewayOTLPExporter, gateway
+
+__all__ = [
+    "gateway",
+    "GatewayOTLPExporter",
+]

gatewaysdk/tracing/exporters/gateway.py

@@ -0,0 +1,228 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Gateway OTLP exporter for GatewayDK tracing.
+
+This module provides an OTLP/HTTP exporter configured for Gateway,
+with support for environment variable auto-detection and Basic Auth headers.
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import os
+import time
+from typing import Optional
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import SpanExportResult
+
+logger = logging.getLogger(__name__)
+
+# Environment variable names (GATEWAY_* prefix)
+ENV_GATEWAY_HOST = "GATEWAY_HOST"
+ENV_GATEWAY_PUBLIC_KEY = "GATEWAY_PUBLIC_KEY"
+ENV_GATEWAY_SECRET_KEY = "GATEWAY_SECRET_KEY"
+ENV_GATEWAY_OTLP_ENDPOINT = "GATEWAY_OTLP_ENDPOINT"
+
+# Default OTLP endpoint path
+DEFAULT_OTLP_PATH = "/api/public/otel/v1/traces"
+
+
+class GatewayOTLPExporter(OTLPSpanExporter):
+    """OTLP exporter configured for Gateway.
+
+    This exporter automatically handles Basic Auth header construction
+    from public/secret key pairs.
+
+    Example::
+
+        from gatewaysdk.tracing_claude.exporters import GatewayOTLPExporter
+
+        exporter = GatewayOTLPExporter(
+            host="https://app.gateway.dev",
+            public_key="pk-lf-...",
+            secret_key="sk-lf-...",
+        )
+    """
+
+    def __init__(
+        self,
+        host: str,
+        public_key: str,
+        secret_key: str,
+        *,
+        endpoint: Optional[str] = None,
+        **kwargs,
+    ):
+        """Initialize the Gateway OTLP exporter.
+
+        Args:
+            host: The Gateway server URL (e.g., "https://app.gateway.dev").
+            public_key: The public key for authentication (pk-lf-...).
+            secret_key: The secret key for authentication (sk-lf-...).
+            endpoint: Optional full OTLP endpoint URL. If not provided,
+                derived from host + DEFAULT_OTLP_PATH.
+            **kwargs: Additional arguments passed to OTLPSpanExporter.
+        """
+        # Derive endpoint from host if not explicitly provided
+        if endpoint is None:
+            # Remove trailing slash if present
+            host = host.rstrip("/")
+            endpoint = f"{host}{DEFAULT_OTLP_PATH}"
+
+        # Construct Basic Auth header
+        credentials = f"{public_key}:{secret_key}"
+        encoded = base64.b64encode(credentials.encode("utf-8")).decode("utf-8")
+        auth_header = f"Basic {encoded}"
+
+        # Merge auth header with any user-provided headers
+        headers = kwargs.pop("headers", {}) or {}
+        headers["Authorization"] = auth_header
+
+        super().__init__(endpoint=endpoint, headers=headers, **kwargs)
+
+        self._gateway_host = host
+        self._gateway_public_key = public_key
+        # Don't store secret key in instance for security
+
+        logger.debug(
+            "GatewayOTLPExporter initialized: endpoint=%s, public_key=%s",
+            endpoint,
+            public_key[:20] + "..." if len(public_key) > 20 else public_key,
+        )
+
+    def export(self, spans):
+        count = len(spans) if spans else 0
+        t0 = time.monotonic()
+        try:
+            result = super().export(spans)
+        except Exception as exc:
+            elapsed_ms = (time.monotonic() - t0) * 1000
+            logger.warning("Export: %d span(s) -> FAILED (%s, %.0fms) %s", count, self._gateway_host, elapsed_ms, exc)
+            return SpanExportResult.FAILURE
+        elapsed_ms = (time.monotonic() - t0) * 1000
+        if result == SpanExportResult.SUCCESS:
+            logger.info("Export: %d span(s) -> SUCCESS (%s, %.0fms)", count, self._gateway_host, elapsed_ms)
+            self._record_exported(count)
+        else:
+            logger.warning("Export: %d span(s) -> FAILED (%s, %.0fms)", count, self._gateway_host, elapsed_ms)
+        return result
+
+    def _record_exported(self, count: int) -> None:
+        """Notify the init module about successful exports for the shutdown summary."""
+        try:
+            from gatewaysdk.tracing.init import _record_exported
+
+            _record_exported(count)
+        except ImportError:
+            pass
+
+    def __repr__(self) -> str:
+        return (
+            f"GatewayOTLPExporter(" f"host={self._gateway_host!r}, " f"public_key={self._gateway_public_key[:10]}...)"
+        )
+
+
+def gateway(
+    *,
+    host: Optional[str] = None,
+    public_key: Optional[str] = None,
+    secret_key: Optional[str] = None,
+    endpoint: Optional[str] = None,
+    **kwargs,
+) -> GatewayOTLPExporter:
+    """Factory function to create an Gateway OTLP exporter.
+
+    This function supports both explicit configuration and environment variable
+    auto-detection. If any required parameter is not provided, it will be read
+    from the corresponding environment variable.
+
+    Environment variables:
+        - GATEWAY_HOST: The Gateway server URL
+        - GATEWAY_PUBLIC_KEY: The public key for authentication
+        - GATEWAY_SECRET_KEY: The secret key for authentication
+        - GATEWAY_OTLP_ENDPOINT: Optional override for the OTLP endpoint
+
+    Args:
+        host: The Gateway server URL. Falls back to GATEWAY_HOST env var.
+        public_key: The public key. Falls back to GATEWAY_PUBLIC_KEY env var.
+        secret_key: The secret key. Falls back to GATEWAY_SECRET_KEY env var.
+        endpoint: Optional full OTLP endpoint URL. Falls back to GATEWAY_OTLP_ENDPOINT.
+        **kwargs: Additional arguments passed to GatewayOTLPExporter.
+
+    Returns:
+        Configured GatewayOTLPExporter instance.
+
+    Raises:
+        ValueError: If required credentials are not provided and not found in env vars.
+
+    Example::
+
+        import gatewaysdk.tracing_claude as tracing
+
+        # Using environment variables
+        tracing.init(exporters=[tracing.gateway_exporter()])
+
+        # Explicit configuration
+        tracing.init(exporters=[
+            tracing.gateway_exporter(
+                host="https://app.gateway.dev",
+                public_key="pk-lf-...",
+                secret_key="sk-lf-...",
+            )
+        ])
+    """
+    # Read from env vars if not provided
+    host = host or os.environ.get(ENV_GATEWAY_HOST)
+    public_key = public_key or os.environ.get(ENV_GATEWAY_PUBLIC_KEY)
+    secret_key = secret_key or os.environ.get(ENV_GATEWAY_SECRET_KEY)
+    endpoint = endpoint or os.environ.get(ENV_GATEWAY_OTLP_ENDPOINT)
+
+    # Validate required parameters
+    missing = []
+    if not host:
+        missing.append(f"{ENV_GATEWAY_HOST} (or host parameter)")
+    if not public_key:
+        missing.append(f"{ENV_GATEWAY_PUBLIC_KEY} (or public_key parameter)")
+    if not secret_key:
+        missing.append(f"{ENV_GATEWAY_SECRET_KEY} (or secret_key parameter)")
+
+    if missing:
+        raise ValueError(
+            f"Missing required Gateway configuration: {', '.join(missing)}. "
+            "Provide these values via function arguments or environment variables."
+        )
+
+    return GatewayOTLPExporter(
+        host=host,
+        public_key=public_key,
+        secret_key=secret_key,
+        endpoint=endpoint,
+        **kwargs,
+    )
+
+
+def detect_gateway_config() -> tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+    """Detect Gateway configuration from environment variables.
+
+    Returns:
+        Tuple of (host, public_key, secret_key, endpoint) where any value may be None
+        if the corresponding environment variable is not set.
+    """
+    return (
+        os.environ.get(ENV_GATEWAY_HOST),
+        os.environ.get(ENV_GATEWAY_PUBLIC_KEY),
+        os.environ.get(ENV_GATEWAY_SECRET_KEY),
+        os.environ.get(ENV_GATEWAY_OTLP_ENDPOINT),
+    )
+
+
+def is_gateway_configured() -> bool:
+    """Check if Gateway is configured via environment variables.
+
+    Returns:
+        True if all required environment variables are set.
+    """
+    host, public_key, secret_key, _ = detect_gateway_config()
+    return all([host, public_key, secret_key])