agent-learning-kit 0.1.0__py3-none-any.whl

fi/evals/otel/tracer.py

@@ -0,0 +1,506 @@
+"""
+OpenTelemetry Tracer Setup.
+
+Factory functions and utilities for setting up OpenTelemetry
+tracing with LLM-specific processors and exporters.
+"""
+
+from typing import Any, Dict, Optional
+from contextlib import contextmanager
+import logging
+import os
+
+from .config import (
+    TraceConfig,
+    ExporterConfig,
+    SamplingStrategy,
+)
+from .types import ExporterType, ProcessorType
+from .processors import (
+    OTEL_AVAILABLE,
+    CompositeSpanProcessor,
+    LLMSpanProcessor,
+    EvaluationSpanProcessor,
+    CostSpanProcessor,
+)
+
+logger = logging.getLogger(__name__)
+
+# Global tracer provider instance
+_tracer_provider: Optional[Any] = None
+_is_initialized = False
+
+
+def setup_tracing(
+    config: Optional[TraceConfig] = None,
+    service_name: Optional[str] = None,
+    project_name: Optional[str] = None,
+    otlp_endpoint: Optional[str] = None,
+    **kwargs,
+) -> Any:
+    """
+    Set up OpenTelemetry tracing for LLM applications.
+
+    This is the main entry point for configuring OTEL tracing.
+    It creates exporters, processors, and a tracer provider.
+
+    Args:
+        config: Full TraceConfig, or None to use defaults/env vars
+        service_name: Service name (overrides config)
+        otlp_endpoint: OTLP endpoint (overrides config)
+        **kwargs: Additional arguments passed to TraceConfig
+
+    Returns:
+        The configured TracerProvider
+
+    Example:
+        # Simple setup with environment variables
+        setup_tracing(service_name="my-llm-service")
+
+        # Full configuration
+        from fi.evals.otel import TraceConfig, ExporterConfig, ExporterType
+
+        config = TraceConfig(
+            service_name="my-service",
+            exporters=[
+                ExporterConfig(type=ExporterType.OTLP_GRPC, endpoint="localhost:4317"),
+                ExporterConfig(type=ExporterType.CONSOLE),
+            ],
+        )
+        setup_tracing(config)
+    """
+    global _tracer_provider, _is_initialized
+
+    if not OTEL_AVAILABLE:
+        logger.warning(
+            "OpenTelemetry SDK not available. Install with: "
+            "pip install opentelemetry-sdk opentelemetry-exporter-otlp"
+        )
+        return None
+
+    # Build configuration
+    if config is None:
+        config = _build_config_from_env(service_name, otlp_endpoint, **kwargs)
+    else:
+        # Override with explicit args
+        if service_name:
+            config.service_name = service_name
+        if otlp_endpoint:
+            config.exporters = [ExporterConfig(
+                type=ExporterType.OTLP_GRPC,
+                endpoint=otlp_endpoint,
+            )]
+
+    if not config.enabled:
+        logger.info("Tracing is disabled")
+        return None
+
+    # Import OTEL components
+    from opentelemetry import trace
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.resources import Resource, SERVICE_NAME
+
+    # Create resource
+    resource_attrs = {
+        SERVICE_NAME: config.service_name,
+        "service.version": config.service_version or "unknown",
+        "deployment.environment": config.deployment_environment,
+        **config.resource.attributes,
+    }
+
+    # Always set trace source
+    resource_attrs["fi.trace.source"] = "traceai"
+
+    # Auto-inject FutureAGI project attributes when using FUTUREAGI exporter
+    has_fi_exporter = any(
+        e.type == ExporterType.FUTUREAGI for e in config.exporters
+    )
+    if has_fi_exporter:
+        import json as _json
+        import uuid as _uuid
+        resource_attrs.setdefault(
+            "project_name",
+            project_name or os.environ.get("FI_PROJECT_NAME", config.service_name),
+        )
+        resource_attrs.setdefault("project_type", "observe")
+        resource_attrs.setdefault(
+            "project_version_name",
+            os.environ.get("FI_PROJECT_VERSION_NAME", "v1"),
+        )
+        resource_attrs.setdefault("project_version_id", str(_uuid.uuid4()))
+        resource_attrs.setdefault("eval_tags", _json.dumps([]))
+        resource_attrs.setdefault("metadata", _json.dumps({}))
+
+    resource = Resource.create(resource_attrs)
+
+    # Create tracer provider
+    provider = TracerProvider(
+        resource=resource,
+        sampler=_create_sampler(config),
+    )
+
+    # Add exporters
+    for exporter_config in config.exporters:
+        exporter = _create_exporter(exporter_config)
+        if exporter:
+            from opentelemetry.sdk.trace.export import BatchSpanProcessor
+            provider.add_span_processor(BatchSpanProcessor(exporter))
+
+    # Add custom processors
+    custom_processors = _create_custom_processors(config)
+    if custom_processors:
+        provider.add_span_processor(custom_processors)
+
+    # Set as global provider
+    trace.set_tracer_provider(provider)
+    _tracer_provider = provider
+    _is_initialized = True
+
+    logger.info(
+        f"OpenTelemetry tracing initialized for '{config.service_name}' "
+        f"with {len(config.exporters)} exporter(s)"
+    )
+
+    return provider
+
+
+def _build_config_from_env(
+    service_name: Optional[str] = None,
+    otlp_endpoint: Optional[str] = None,
+    **kwargs,
+) -> TraceConfig:
+    """Build TraceConfig from environment variables."""
+    # Service name
+    service_name = service_name or os.getenv("OTEL_SERVICE_NAME", "llm-service")
+
+    # OTLP endpoint
+    endpoint = otlp_endpoint or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
+
+    # Build exporters
+    exporters = []
+    if endpoint:
+        exporters.append(ExporterConfig(
+            type=ExporterType.OTLP_GRPC,
+            endpoint=endpoint,
+            headers=_parse_headers(os.getenv("OTEL_EXPORTER_OTLP_HEADERS", "")),
+        ))
+    else:
+        # Default to console
+        exporters.append(ExporterConfig(type=ExporterType.CONSOLE))
+
+    return TraceConfig(
+        service_name=service_name,
+        exporters=exporters,
+        deployment_environment=os.getenv("OTEL_DEPLOYMENT_ENVIRONMENT", "development"),
+        **kwargs,
+    )
+
+
+def _parse_headers(headers_str: str) -> Dict[str, str]:
+    """Parse headers from comma-separated key=value string."""
+    if not headers_str:
+        return {}
+
+    headers = {}
+    for pair in headers_str.split(","):
+        if "=" in pair:
+            key, value = pair.split("=", 1)
+            headers[key.strip()] = value.strip()
+    return headers
+
+
+def _create_sampler(config: TraceConfig) -> Any:
+    """Create a sampler based on configuration."""
+    from opentelemetry.sdk.trace.sampling import (
+        ALWAYS_ON,
+        ALWAYS_OFF,
+        TraceIdRatioBased,
+        ParentBased,
+    )
+
+    if config.sampling_strategy == SamplingStrategy.ALWAYS_ON:
+        return ALWAYS_ON
+    elif config.sampling_strategy == SamplingStrategy.ALWAYS_OFF:
+        return ALWAYS_OFF
+    elif config.sampling_strategy == SamplingStrategy.RATIO:
+        return TraceIdRatioBased(config.sampling_ratio)
+    elif config.sampling_strategy == SamplingStrategy.PARENT_BASED:
+        return ParentBased(TraceIdRatioBased(config.sampling_ratio))
+    else:
+        return ALWAYS_ON
+
+
+def _create_exporter(config: ExporterConfig) -> Optional[Any]:
+    """Create an exporter based on configuration."""
+    try:
+        if config.type == ExporterType.CONSOLE:
+            from opentelemetry.sdk.trace.export import ConsoleSpanExporter
+            return ConsoleSpanExporter()
+
+        elif config.type == ExporterType.OTLP_GRPC:
+            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+                OTLPSpanExporter,
+            )
+            return OTLPSpanExporter(
+                endpoint=config.endpoint or "localhost:4317",
+                headers=config.headers or None,
+                insecure=config.insecure,
+                timeout=config.timeout_ms // 1000,
+            )
+
+        elif config.type == ExporterType.OTLP_HTTP:
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+                OTLPSpanExporter,
+            )
+            return OTLPSpanExporter(
+                endpoint=config.endpoint or "http://localhost:4318/v1/traces",
+                headers=config.headers or None,
+                timeout=config.timeout_ms // 1000,
+            )
+
+        elif config.type == ExporterType.JAEGER:
+            try:
+                from opentelemetry.exporter.jaeger.thrift import JaegerExporter
+                return JaegerExporter(
+                    agent_host_name=config.endpoint.split(":")[0] if config.endpoint else "localhost",
+                    agent_port=int(config.endpoint.split(":")[1]) if config.endpoint and ":" in config.endpoint else 6831,
+                )
+            except ImportError:
+                logger.warning("Jaeger exporter not available. Install opentelemetry-exporter-jaeger")
+                return None
+
+        elif config.type == ExporterType.ZIPKIN:
+            try:
+                from opentelemetry.exporter.zipkin.json import ZipkinExporter
+                return ZipkinExporter(
+                    endpoint=config.endpoint or "http://localhost:9411/api/v2/spans",
+                )
+            except ImportError:
+                logger.warning("Zipkin exporter not available. Install opentelemetry-exporter-zipkin")
+                return None
+
+        elif config.type in (ExporterType.ARIZE, ExporterType.PHOENIX, ExporterType.LANGFUSE):
+            # These platforms support OTLP
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+                OTLPSpanExporter,
+            )
+            return OTLPSpanExporter(
+                endpoint=config.endpoint,
+                headers=config.headers or None,
+            )
+
+        elif config.type == ExporterType.FUTUREAGI:
+            # Use OTLP HTTP for FutureAGI backend
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+                OTLPSpanExporter,
+            )
+            if config.endpoint:
+                endpoint = config.endpoint
+            else:
+                base = os.environ.get("FI_BASE_URL", "https://api.futureagi.com")
+                endpoint = f"{base.rstrip('/')}/tracer/v1/traces"
+            # Auto-inject FI auth headers if not provided
+            headers = config.headers
+            if not headers:
+                api_key = os.environ.get("FI_API_KEY", "")
+                secret_key = os.environ.get("FI_SECRET_KEY", "")
+                if api_key and secret_key:
+                    headers = {
+                        "X-Api-Key": api_key,
+                        "X-Secret-Key": secret_key,
+                    }
+            return OTLPSpanExporter(
+                endpoint=endpoint,
+                headers=headers or None,
+            )
+
+        else:
+            logger.warning(f"Unsupported exporter type: {config.type}")
+            return None
+
+    except ImportError as e:
+        logger.warning(f"Failed to create exporter {config.type}: {e}")
+        return None
+    except Exception as e:
+        logger.error(f"Error creating exporter {config.type}: {e}")
+        return None
+
+
+def _create_custom_processors(config: TraceConfig) -> Optional[CompositeSpanProcessor]:
+    """Create custom span processors based on configuration."""
+    processors = []
+
+    # Check for processor configs
+    processor_types = {p.type for p in config.processors if p.enabled}
+
+    # LLM processor
+    if ProcessorType.LLM in processor_types:
+        processors.append(LLMSpanProcessor(
+            capture_prompts=config.content.capture_prompts,
+            capture_completions=config.content.capture_completions,
+            max_content_length=config.content.max_content_length,
+            redact_patterns=config.content.redact_patterns,
+        ))
+
+    # Evaluation processor
+    if config.evaluation.enabled:
+        processors.append(EvaluationSpanProcessor(
+            metrics=config.evaluation.metrics,
+            sample_rate=config.evaluation.sample_rate,
+            async_evaluation=config.evaluation.async_evaluation,
+            timeout_ms=config.evaluation.timeout_ms,
+            cache_enabled=config.evaluation.cache_enabled,
+            cache_ttl_seconds=config.evaluation.cache_ttl_seconds,
+            evaluator_model=config.evaluation.evaluator_model,
+        ))
+
+    # Cost processor
+    if config.cost.enabled:
+        processors.append(CostSpanProcessor(
+            custom_pricing=None,  # Could extend config for this
+            alert_threshold_usd=config.cost.alert_threshold_usd,
+        ))
+
+    if not processors:
+        return None
+
+    return CompositeSpanProcessor(processors)
+
+
+def get_tracer(name: str = "fi.evals") -> Any:
+    """
+    Get a tracer instance.
+
+    Args:
+        name: Tracer name (typically module name)
+
+    Returns:
+        Tracer instance
+    """
+    if not OTEL_AVAILABLE:
+        return _NoOpTracer()
+
+    from opentelemetry import trace
+    return trace.get_tracer(name)
+
+
+def get_current_span() -> Any:
+    """Get the current active span."""
+    if not OTEL_AVAILABLE:
+        return None
+
+    from opentelemetry import trace
+    return trace.get_current_span()
+
+
+def is_tracing_enabled() -> bool:
+    """Check if tracing is enabled and initialized."""
+    return _is_initialized and _tracer_provider is not None
+
+
+@contextmanager
+def trace_llm_call(
+    name: str,
+    model: Optional[str] = None,
+    system: Optional[str] = None,
+    **attributes,
+):
+    """
+    Context manager for tracing an LLM call.
+
+    Creates a span with LLM-specific attributes.
+
+    Args:
+        name: Span name
+        model: Model name
+        system: Provider system (openai, anthropic, etc.)
+        **attributes: Additional span attributes
+
+    Yields:
+        The active span
+
+    Example:
+        with trace_llm_call("chat", model="gpt-4", system="openai") as span:
+            response = client.chat.completions.create(...)
+            span.set_attribute("gen_ai.usage.input_tokens", response.usage.prompt_tokens)
+    """
+    from .conventions import GenAIAttributes, OPERATION_CHAT
+
+    tracer = get_tracer()
+
+    # Build initial attributes
+    span_attrs = {
+        GenAIAttributes.SPAN_KIND: "LLM",
+        GenAIAttributes.OPERATION_NAME: OPERATION_CHAT,
+    }
+    if model:
+        span_attrs[GenAIAttributes.REQUEST_MODEL] = model
+    if system:
+        span_attrs[GenAIAttributes.PROVIDER_NAME] = system
+    span_attrs.update(attributes)
+
+    with tracer.start_as_current_span(name, attributes=span_attrs) as span:
+        yield span
+
+
+def shutdown_tracing() -> None:
+    """Shutdown the tracer provider and flush pending spans."""
+    global _tracer_provider, _is_initialized
+
+    if _tracer_provider:
+        try:
+            _tracer_provider.shutdown()
+        except Exception as e:
+            logger.warning(f"Error shutting down tracer: {e}")
+
+    _tracer_provider = None
+    _is_initialized = False
+
+
+class _NoOpTracer:
+    """No-op tracer for when OTEL is not available."""
+
+    def start_span(self, name: str, **kwargs):
+        return _NoOpSpan()
+
+    def start_as_current_span(self, name: str, **kwargs):
+        return _NoOpContextManager()
+
+
+class _NoOpSpan:
+    """No-op span."""
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        pass
+
+    def add_event(self, name: str, attributes: Optional[Dict] = None) -> None:
+        pass
+
+    def set_status(self, status: Any) -> None:
+        pass
+
+    def end(self, end_time: Optional[int] = None) -> None:
+        pass
+
+    def record_exception(self, exception: Exception) -> None:
+        pass
+
+
+class _NoOpContextManager:
+    """No-op context manager."""
+
+    def __enter__(self):
+        return _NoOpSpan()
+
+    def __exit__(self, *args):
+        pass
+
+
+__all__ = [
+    "setup_tracing",
+    "get_tracer",
+    "get_current_span",
+    "is_tracing_enabled",
+    "trace_llm_call",
+    "shutdown_tracing",
+]