gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/tracing/init.py

@@ -0,0 +1,620 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Initialization and lifecycle management for GatewayDK tracing.
+
+This module provides the core init(), shutdown(), and flush() functions
+for managing the OpenTelemetry TracerProvider lifecycle.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import threading
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+
+from .attributes import GATEWAY_ENVIRONMENT, GATEWAY_VERSION
+from .exporters.gateway import gateway, is_gateway_configured
+from .processors import LoggingSpanProcessor, TracingContextSpanProcessor
+
+if TYPE_CHECKING:
+    from .identity import AgentIdentity
+    from openinference.instrumentation import TraceConfig
+
+logger = logging.getLogger(__name__)
+
+# Module-level state
+_initialized = False
+_init_lock = threading.Lock()
+_provider: Optional[TracerProvider] = None
+_processors: List[BatchSpanProcessor] = []
+_export_count = 0
+_batch_config: Dict[str, Any] = {}  # Batch processor kwargs, set by init()
+
+
+def init(
+    *,
+    agent: Optional[Union[str, dict, "AgentIdentity"]] = None,
+    trace_name: Optional[str] = None,
+    service_name: Optional[str] = None,
+    environment: Optional[str] = None,
+    version: Optional[str] = None,
+    # Gateway configuration (alternative to env vars)
+    gateway_host: Optional[str] = None,
+    gateway_public_key: Optional[str] = None,
+    gateway_secret_key: Optional[str] = None,
+    gateway_endpoint: Optional[str] = None,
+    instrument_default: bool = True,
+    exporters: Optional[List[SpanExporter]] = None,
+    trace_config: Optional["TraceConfig"] = None,
+    force: bool = False,
+    integrations: Optional[Dict[str, Any]] = None,
+    debug: Union[bool, str] = False,
+    # Batch processor tuning (passed through to BatchSpanProcessor)
+    batch_queue_size: Optional[int] = None,
+    batch_flush_interval: Optional[int] = None,
+    batch_max_export_size: Optional[int] = None,
+) -> None:
+    """Initialize GatewaySDK standalone tracing.
+
+    This function configures OpenTelemetry with a connection to Gateway
+    (if configured via environment variables), and automatically instruments
+    installed AI libraries (OpenAI, Anthropic, LiteLLM, OpenAI Agents SDK, etc.).
+
+    Provider Reuse Behavior:
+        - If no TracerProvider exists, creates one.
+        - If a TracerProvider exists and force=False, reuses it and adds
+          processors/exporters (avoiding duplicates).
+        - If force=True, replaces the existing TracerProvider.
+
+    Args:
+        agent: The agent this process runs as — the one identity you declare.
+            A string id, an AgentIdentity, or a dict with id/name/version.
+            Falls back to GATEWAY_AGENT_ID / GATEWAY_AGENT_NAME /
+            GATEWAY_AGENT_VERSION. Every trace is stamped with the agent
+            automatically, and the agent is registered with the platform
+            control plane (best-effort, non-blocking) so it appears on the
+            Agents page before its first trace.
+        trace_name: Optional name for the trace.
+        service_name: Name of the service (mapped to resource attributes).
+            Defaults to "gatewaysdk-tracing".
+        environment: Environment name (e.g., "production", "staging", "development").
+            Sets `gateway.environment` resource attribute for filtering traces in Gateway.
+        version: Application version (e.g., "1.2.3", git SHA). Sets `gateway.version`
+            resource attribute for tracking regressions across versions.
+        gateway_host: Gateway server URL (e.g., "https://app.gateway.dev").
+            Overrides GATEWAY_HOST environment variable.
+        gateway_public_key: Gateway public key (pk-lf-...).
+            Overrides GATEWAY_PUBLIC_KEY environment variable.
+        gateway_secret_key: Gateway secret key (sk-lf-...).
+            Overrides GATEWAY_SECRET_KEY environment variable.
+        gateway_endpoint: Optional OTLP endpoint URL. Overrides GATEWAY_OTLP_ENDPOINT.
+            If not provided, derived from gateway_host.
+        instrument_default: Whether to automatically instrument the "core set"
+            of libraries (OpenAI Agents SDK, OpenAI, Anthropic, LiteLLM, etc.)
+            if installed. Defaults to True.
+        exporters: List of custom OpenTelemetry SpanExporters. If None and
+            Gateway credentials are configured (via params or env vars), auto-creates
+            an Gateway exporter.
+        trace_config: OpenInference TraceConfig for controlling capture behavior
+            (hiding inputs/outputs, etc.). Default captures all.
+        force: If True, replaces any existing global TracerProvider. If False
+            (default), reuses the existing provider and appends processors.
+        integrations: Dictionary for fine-grained control over integrations.
+        debug: Controls console logging of tracing activity. Accepts:
+            - False (default): Silent operation, no console output.
+            - True: Shortcut for "DEBUG" -- full firehose including every
+              span start/end, export results, and internal plumbing.
+            - "DEBUG": Same as True.
+            - "INFO": Recommended. Shows init summary, export success/failure,
+              instrumentor activation, and shutdown summary.
+            - "WARNING": Minimal. Only shows export failures and missing
+              exporter warnings.
+            - "ERROR": Only export failures and exceptions.
+        batch_queue_size: Max spans buffered in memory by BatchSpanProcessor.
+            Default: 2048. Increase for bulk import workloads to prevent
+            silent span drops. Wraps OTEL_BSP_MAX_QUEUE_SIZE.
+        batch_flush_interval: How often (in ms) the processor flushes spans.
+            Default: 5000. Decrease for lower-latency export.
+            Wraps OTEL_BSP_SCHEDULE_DELAY.
+        batch_max_export_size: Max spans per export call. Default: 512.
+            Wraps OTEL_BSP_MAX_EXPORT_BATCH_SIZE.
+
+    Example::
+
+        import gatewaysdk.tracing as tracing
+
+        # Minimal: auto-detect Gateway from env vars, auto-instrument
+        tracing.init()
+
+        # With environment and version for filtering in Gateway
+        tracing.init(
+            service_name="my-agent",
+            environment="production",
+            version="1.2.3",
+            instrument_default=True,
+        )
+
+        # Explicit Gateway credentials (instead of env vars)
+        tracing.init(
+            service_name="my-agent",
+            environment="production",
+            version="1.2.3",
+            gateway_host="https://app.gateway.dev",
+            gateway_public_key="pk-lf-...",
+            gateway_secret_key="sk-lf-...",
+        )
+
+        # Ensure cleanup on shutdown
+        import atexit
+        atexit.register(tracing.shutdown)
+    """
+    global _initialized, _provider
+
+    resolved_level = _resolve_debug_level(debug)
+    if resolved_level is not None:
+        _configure_debug_logging(resolved_level)
+
+    with _init_lock:
+        # Check existing provider
+        current_provider = trace.get_tracer_provider()
+        has_sdk_provider = isinstance(current_provider, TracerProvider)
+
+        if _initialized and not force:
+            logger.debug("GatewaySDK tracing already initialized. Use force=True to reinitialize.")
+            return
+
+        if has_sdk_provider and not force:
+            # Reuse existing provider
+            logger.info("Reusing existing TracerProvider")
+            _provider = current_provider
+        else:
+            # Create new provider
+            resource_attrs = {
+                SERVICE_NAME: service_name or "gatewaysdk-tracing",
+            }
+            if trace_name:
+                resource_attrs["trace.name"] = trace_name
+            if environment:
+                resource_attrs[GATEWAY_ENVIRONMENT] = environment
+            if version:
+                resource_attrs[GATEWAY_VERSION] = version
+
+            resource = Resource.create(resource_attrs)
+            _provider = TracerProvider(resource=resource)
+
+            # Set as global provider
+            trace.set_tracer_provider(_provider)
+            logger.info("Created new TracerProvider with service_name=%s", service_name or "gatewaysdk-tracing")
+
+        # Add TracingContextSpanProcessor for automatic context injection
+        # This must be added BEFORE exporter processors so context is set first
+        _add_context_processor_idempotent(_provider)
+
+        if resolved_level is not None:
+            _add_logging_processor_idempotent(_provider)
+
+        configured_exporters = exporters or []
+
+        # Auto-configure Gateway exporter if no custom exporters provided
+        if not configured_exporters:
+            # Check if Gateway credentials are provided via parameters
+            if gateway_host and gateway_public_key and gateway_secret_key:
+                try:
+                    gateway_exporter = gateway(
+                        host=gateway_host,
+                        public_key=gateway_public_key,
+                        secret_key=gateway_secret_key,
+                        endpoint=gateway_endpoint,
+                    )
+                    configured_exporters.append(gateway_exporter)
+                    logger.info("Configured Gateway exporter from init() parameters")
+                except ValueError as e:
+                    logger.warning("Could not configure Gateway exporter: %s", e)
+            # Fall back to environment variables
+            elif is_gateway_configured():
+                try:
+                    gateway_exporter = gateway()
+                    configured_exporters.append(gateway_exporter)
+                    logger.info("Auto-configured Gateway exporter from environment variables")
+                except ValueError as e:
+                    logger.warning("Could not auto-configure Gateway exporter: %s", e)
+
+        if not configured_exporters:
+            logger.warning("No exporter configured -- spans will not be persisted")
+
+        # Store batch processor config for _add_processor_idempotent
+        _batch_config.clear()
+        if batch_queue_size is not None:
+            _batch_config["max_queue_size"] = batch_queue_size
+        if batch_flush_interval is not None:
+            _batch_config["schedule_delay_millis"] = batch_flush_interval
+        if batch_max_export_size is not None:
+            _batch_config["max_export_batch_size"] = batch_max_export_size
+
+        # Add processors for each exporter (idempotent)
+        for exporter in configured_exporters:
+            _add_processor_idempotent(_provider, exporter)
+
+        # Run instrumentation
+        if instrument_default:
+            instrument()
+
+        # Agent identity: the one declared identity. Everything else
+        # (user, session, context items) flows inherently.
+        from .identity import (
+            register_agent_with_platform,
+            resolve_agent_identity,
+            set_agent_identity,
+        )
+
+        identity = resolve_agent_identity(agent)
+        set_agent_identity(identity)
+        from .identity import initialize_ambient_session
+
+        initialize_ambient_session()
+        if identity is not None:
+            register_agent_with_platform(
+                identity,
+                host=gateway_host,
+                public_key=gateway_public_key,
+                secret_key=gateway_secret_key,
+            )
+            # Agent build identity — "the agent IS the manifest": everything
+            # registered before init (define_tool / register_prompt) plus the
+            # code ref becomes a content-addressed build, stamped on every
+            # span and auto-registered (idempotent on hash, fail-soft).
+            from gatewaysdk.build import activate_build
+
+            activate_build(
+                identity.id,
+                version=identity.version,
+                host=gateway_host,
+                public_key=gateway_public_key,
+                secret_key=gateway_secret_key,
+            )
+
+        _initialized = True
+        logger.info("GatewaySDK tracing initialized successfully")
+
+
+def instrument(
+    names: Optional[List[str]] = None,
+    skip: Optional[List[str]] = None,
+) -> None:
+    """Manually enable specific instrumentations.
+
+    This function looks up instrumentors by name from the registry and
+    activates them. If no names are provided and init() was called with
+    instrument_default=True, this enables the "core set".
+
+    Core Set (enabled by default):
+        - openai
+        - anthropic
+        - langchain
+        - llama_index
+        - litellm
+
+    Additional Available:
+        - google_adk
+        - mistral
+        - groq
+        - bedrock
+
+    Args:
+        names: List of instrumentor names to enable. If None and called from
+            init() with instrument_default=True, enables the core set.
+        skip: List of instrumentor names to explicitly skip.
+
+    Example::
+
+        import gatewaysdk.tracing_claude as tracing
+
+        # Enable only specific instrumentors
+        tracing.init(instrument_default=False)
+        tracing.instrument(names=["openai", "anthropic"])
+
+        # Enable all except langchain
+        tracing.instrument(skip=["langchain"])
+    """
+    from .instrumentors.registry import get_registry
+
+    registry = get_registry()
+    skip_set = set(skip or [])
+
+    # Determine which instrumentors to enable
+    if names is None:
+        # Default core set -- instruments all installed AI libraries automatically
+        target_names = [
+            "openai_agents",
+            "claude_agent_sdk",
+            "openai",
+            "anthropic",
+            "langchain",
+            "llama_index",
+            "litellm",
+        ]
+    else:
+        target_names = names
+
+    # Filter out skipped instrumentors
+    target_names = [name for name in target_names if name not in skip_set]
+
+    # Activate each instrumentor
+    for name in target_names:
+        instrumentor = registry.get(name)
+        if instrumentor is not None:
+            try:
+                instrumentor.instrument()
+                logger.info("Instrumented: %s", name)
+            except Exception as e:
+                logger.debug("Could not instrument %s: %s", name, e)
+        else:
+            logger.debug("Instrumentor not available (not installed?): %s", name)
+            # Warn loudly when the AI library IS installed but its instrumentor
+            # is missing — otherwise the user gets zero telemetry, silently.
+            registry.warn_if_lib_present_but_uninstrumented(name)
+
+    # The openai-agents instrumentor rides on the Agents SDK's own tracing
+    # pipeline. OPENAI_AGENTS_DISABLE_TRACING=1 (commonly set to silence the
+    # OpenAI trace-upload 401s when pointing at a non-OpenAI base URL) turns
+    # that pipeline off entirely — and with it, ALL Gateway telemetry for the
+    # agent. Warn loudly; the 401 noise alone is harmless.
+    if "openai_agents" in target_names and os.environ.get(
+        "OPENAI_AGENTS_DISABLE_TRACING", ""
+    ).lower() in ("1", "true", "yes"):
+        print(
+            "gatewaysdk.tracing: OPENAI_AGENTS_DISABLE_TRACING is set, which "
+            "disables the OpenAI Agents SDK tracing pipeline the Gateway "
+            "instrumentor rides on — NOTHING from openai-agents will reach "
+            "Gateway. Unset it (the OpenAI trace-upload 401s it silences are "
+            "harmless) or set OPENAI_AGENTS_TRACING_EXPORT_API_KEY to a real "
+            "OpenAI key.",
+            file=sys.stderr,
+        )
+
+
+def shutdown(timeout_millis: int = 30000) -> None:
+    """Force flush all pending spans and shutdown the tracer provider.
+
+    This function should be called before process exit to ensure all spans
+    are exported. It is safe to call multiple times.
+
+    Args:
+        timeout_millis: Maximum time to wait for flush completion in milliseconds.
+
+    Example::
+
+        import gatewaysdk.tracing_claude as tracing
+        import atexit
+
+        tracing.init()
+        atexit.register(tracing.shutdown)
+
+        # ... run application ...
+    """
+    global _initialized, _provider, _context_processor_added, _logging_processor_added, _export_count
+
+    with _init_lock:
+        if not _initialized or _provider is None:
+            logger.debug("Tracing not initialized, nothing to shutdown")
+            return
+
+        try:
+            # Shutdown the provider (which flushes and shuts down processors)
+            _provider.shutdown()
+            logger.info("Shutdown complete: exported %d span(s) total", _export_count)
+        except Exception as e:
+            logger.warning("Error during TracerProvider shutdown: %s", e)
+
+        _processors.clear()
+        _batch_config.clear()
+        _initialized = False
+        _provider = None
+        _context_processor_added = False
+        _logging_processor_added = False
+        _export_count = 0
+
+
+def flush(timeout_millis: int = 30000) -> bool:
+    """Force flush all pending spans without shutting down.
+
+    This is useful for ensuring spans are exported at specific points
+    (e.g., after completing a batch of work) without ending the tracing session.
+
+    Args:
+        timeout_millis: Maximum time to wait for flush completion in milliseconds.
+
+    Returns:
+        True if flush completed successfully, False otherwise.
+
+    Example::
+
+        import gatewaysdk.tracing_claude as tracing
+
+        tracing.init()
+
+        # Process batch
+        for item in batch:
+            process(item)
+
+        # Ensure all spans from this batch are exported
+        tracing.flush()
+    """
+    with _init_lock:
+        if _provider is None:
+            logger.debug("No provider to flush")
+            return True
+
+        try:
+            success = _provider.force_flush(timeout_millis=timeout_millis)
+            if success:
+                logger.debug("Flush completed successfully")
+            else:
+                logger.warning("Flush timed out after %d ms", timeout_millis)
+            return success
+        except Exception as e:
+            logger.warning("Error during flush: %s", e)
+            return False
+
+
+# Track if context processor has been added
+_context_processor_added = False
+
+
+def _add_context_processor_idempotent(provider: TracerProvider) -> None:
+    """Add the TracingContextSpanProcessor if not already added.
+
+    This processor injects run context (run_id, session_id) and tracing
+    context (call_type, tags, environment) into all spans automatically.
+    """
+    global _context_processor_added
+
+    if _context_processor_added:
+        logger.debug("TracingContextSpanProcessor already added")
+        return
+
+    processor = TracingContextSpanProcessor()
+    provider.add_span_processor(processor)
+    _context_processor_added = True
+    logger.debug("Added TracingContextSpanProcessor for automatic context injection")
+
+
+def _add_processor_idempotent(provider: TracerProvider, exporter: SpanExporter) -> None:
+    """Add a BatchSpanProcessor for the exporter, avoiding duplicates.
+
+    Checks if a processor for the same exporter type/endpoint already exists
+    to prevent duplicate exports.
+    """
+    # Check for existing processor with same exporter type
+    exporter_id = _get_exporter_id(exporter)
+    for existing in _processors:
+        if _get_exporter_id(existing.span_exporter) == exporter_id:
+            logger.debug("Processor already exists for exporter: %s", exporter_id)
+            return
+
+    # Create and add new processor.
+    # Wrap Gateway exporters with AnnotationFilterSpanProcessor to suppress
+    # gatewaysdk.annotation reward spans from OTLP (rewards are sent as
+    # Langfuse Scores via the API instead).
+    from .exporters.gateway import GatewayOTLPExporter
+    from .processors import AnnotationFilterSpanProcessor
+
+    batch_processor = BatchSpanProcessor(exporter, **_batch_config)
+    if isinstance(exporter, GatewayOTLPExporter):
+        filtered = AnnotationFilterSpanProcessor(batch_processor)
+        provider.add_span_processor(filtered)
+        logger.debug("Added filtered BatchSpanProcessor for: %s", exporter_id)
+    else:
+        provider.add_span_processor(batch_processor)
+        logger.debug("Added BatchSpanProcessor for: %s", exporter_id)
+    _processors.append(batch_processor)
+
+
+def _get_exporter_id(exporter: SpanExporter) -> str:
+    """Generate a unique identifier for an exporter instance.
+
+    Uses class name + endpoint (if available) to identify exporters.
+    """
+    class_name = exporter.__class__.__name__
+
+    # Try to get endpoint for OTLP exporters
+    endpoint = getattr(exporter, "_endpoint", None)
+    if endpoint:
+        return f"{class_name}:{endpoint}"
+
+    return class_name
+
+
+def _resolve_debug_level(debug: Union[bool, str]) -> Optional[str]:
+    """Resolve the debug parameter to a Python logging level name, or None if disabled."""
+    if debug is False or debug is None:
+        return None
+    if debug is True:
+        return "DEBUG"
+    level_name = str(debug).upper()
+    if level_name not in ("DEBUG", "INFO", "WARNING", "ERROR"):
+        raise ValueError(f"Invalid debug level: {debug!r}. " "Use True, False, 'DEBUG', 'INFO', 'WARNING', or 'ERROR'.")
+    return level_name
+
+
+def _record_exported(count: int) -> None:
+    """Increment the global export counter (called by GatewayOTLPExporter)."""
+    global _export_count
+    _export_count += count
+
+
+_logging_processor_added = False
+
+
+def _add_logging_processor_idempotent(provider: TracerProvider) -> None:
+    """Add the LoggingSpanProcessor if not already added."""
+    global _logging_processor_added
+
+    if _logging_processor_added:
+        return
+
+    provider.add_span_processor(LoggingSpanProcessor())
+    _logging_processor_added = True
+    logger.debug("Added LoggingSpanProcessor for debug visibility")
+
+
+def _configure_debug_logging(level: str = "DEBUG") -> None:
+    """Configure Python logging so gatewaysdk.tracing messages reach the console.
+
+    Sets up a StreamHandler on the ``gatewaysdk.tracing`` logger hierarchy
+    (and the OTLP exporter logger) at the requested level.
+    Writes to stdout (not stderr) so output interleaves naturally with
+    application print() statements.
+
+    Args:
+        level: Python logging level name (DEBUG, INFO, WARNING, ERROR).
+    """
+    import sys
+
+    numeric_level = getattr(logging, level, logging.DEBUG)
+
+    tracing_logger = logging.getLogger("gatewaysdk.tracing")
+    if tracing_logger.handlers:
+        return
+
+    handler = logging.StreamHandler(stream=sys.stdout)
+    handler.setLevel(numeric_level)
+    fmt = logging.Formatter("[%(name)s] %(message)s")
+    handler.setFormatter(fmt)
+
+    tracing_logger.addHandler(handler)
+    tracing_logger.setLevel(numeric_level)
+    tracing_logger.propagate = False
+
+    # Also capture OTLP exporter warnings (connection failures, retries)
+    otel_logger = logging.getLogger("opentelemetry.exporter.otlp")
+    otel_logger.addHandler(handler)
+    otel_logger.setLevel(logging.WARNING)
+    otel_logger.propagate = False
+
+    # Suppress noisy OTel warnings about provider overriding (force=True triggers these)
+    for name in ("opentelemetry.trace", "opentelemetry.sdk.trace"):
+        lg = logging.getLogger(name)
+        lg.setLevel(logging.ERROR)
+
+
+def get_tracer(name: str = "gatewaysdk.tracing") -> trace.Tracer:
+    """Get a tracer instance for creating spans.
+
+    This is primarily for internal use. Users should prefer the
+    trace() and span() context managers.
+
+    Args:
+        name: The tracer name (instrumentation scope).
+
+    Returns:
+        OpenTelemetry Tracer instance.
+    """
+    return trace.get_tracer(name)

gatewaysdk/tracing/instrumentors/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Instrumentors for GatewayDK tracing.
+
+This module provides a registry-based system for managing OpenTelemetry
+instrumentors from various sources (OpenInference, AgentOps, etc.).
+"""
+
+from .registry import BaseInstrumentor, InstrumentorRegistry, get_registry
+
+__all__ = [
+    "get_registry",
+    "InstrumentorRegistry",
+    "BaseInstrumentor",
+]