reasonkit-core 0.1.8

The Reasoning Engine — Auditable Reasoning for Production AI | Rust-Native | Turn Prompts into Protocols
//! OpenTelemetry Observability Integration
//!
//! This module provides distributed tracing, metrics, and logging integration
//! using OpenTelemetry for ReasonKit components.
//!
//! # Features
//! - Distributed tracing with context propagation
//! - Custom metrics for ThinkTools and reasoning operations
//! - OTLP export to collectors (Jaeger, Prometheus, etc.)
//! - Integration with existing tracing subscriber
//!
//! Enable with: `cargo build --features otel`

use anyhow::Result;
use serde::{Deserialize, Serialize};

// Re-exports for convenience
#[cfg(feature = "otel")]
pub use opentelemetry;
#[cfg(feature = "otel")]
pub use tracing_opentelemetry;

/// Configuration for OpenTelemetry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OtelConfig {
    /// Service name for tracing
    pub service_name: String,
    /// Service version
    pub service_version: String,
    /// OTLP endpoint for traces
    pub otlp_endpoint: Option<String>,
    /// Enable tracing
    pub tracing_enabled: bool,
    /// Enable metrics
    pub metrics_enabled: bool,
    /// Sampling ratio (0.0 - 1.0)
    pub sampling_ratio: f64,
    /// Batch export timeout
    pub export_timeout_secs: u64,
}

impl Default for OtelConfig {
    fn default() -> Self {
        Self {
            service_name: "reasonkit".to_string(),
            service_version: env!("CARGO_PKG_VERSION").to_string(),
            otlp_endpoint: None,
            tracing_enabled: true,
            metrics_enabled: true,
            sampling_ratio: 1.0,
            export_timeout_secs: 30,
        }
    }
}

/// Initialize OpenTelemetry with the given configuration.
///
/// This sets up the global tracer provider with optional OTLP export.
/// For most use cases, prefer using `init_subscriber()` which sets up
/// the complete tracing stack including log formatting.
pub fn init_otel(config: &OtelConfig) -> Result<OtelGuard> {
    use opentelemetry::global;
    use opentelemetry::KeyValue;
    use opentelemetry_sdk::trace::TracerProvider;
    use opentelemetry_sdk::Resource;

    // Create resource with service information
    let resource = Resource::new(vec![
        KeyValue::new("service.name", config.service_name.clone()),
        KeyValue::new("service.version", config.service_version.clone()),
    ]);

    // Build tracer provider
    let provider = TracerProvider::builder().with_resource(resource).build();

    global::set_tracer_provider(provider);

    tracing::info!(
        service = %config.service_name,
        endpoint = ?config.otlp_endpoint,
        "OpenTelemetry initialized (basic mode)"
    );

    // Note: OTLP export requires additional runtime configuration.
    // For production use, configure the OTEL_EXPORTER_OTLP_ENDPOINT env var
    // and use the opentelemetry-otlp crate's pipeline builder.
    if config.otlp_endpoint.is_some() {
        tracing::warn!(
            "OTLP endpoint configured but export requires additional setup. \
             Set OTEL_EXPORTER_OTLP_ENDPOINT env var for automatic export."
        );
    }

    Ok(OtelGuard { _private: () })
}

/// Guard that shuts down OpenTelemetry on drop
pub struct OtelGuard {
    _private: (),
}

impl Drop for OtelGuard {
    fn drop(&mut self) {
        opentelemetry::global::shutdown_tracer_provider();
        tracing::info!("OpenTelemetry shut down");
    }
}

/// Custom span attributes for ReasonKit operations
pub mod attributes {
    use opentelemetry::KeyValue;

    /// ThinkTool name attribute
    pub fn thinktool_name(name: &str) -> KeyValue {
        KeyValue::new("reasonkit.thinktool.name", name.to_string())
    }

    /// ThinkTool profile attribute
    pub fn thinktool_profile(profile: &str) -> KeyValue {
        KeyValue::new("reasonkit.thinktool.profile", profile.to_string())
    }

    /// Reasoning confidence attribute
    pub fn reasoning_confidence(confidence: f32) -> KeyValue {
        KeyValue::new("reasonkit.reasoning.confidence", confidence as f64)
    }

    /// Token count attribute
    pub fn token_count(tokens: usize) -> KeyValue {
        KeyValue::new("reasonkit.tokens.total", tokens as i64)
    }

    /// LLM provider attribute
    pub fn llm_provider(provider: &str) -> KeyValue {
        KeyValue::new("reasonkit.llm.provider", provider.to_string())
    }

    /// LLM model attribute
    pub fn llm_model(model: &str) -> KeyValue {
        KeyValue::new("reasonkit.llm.model", model.to_string())
    }

    /// Cache hit attribute
    pub fn cache_hit(hit: bool) -> KeyValue {
        KeyValue::new("reasonkit.cache.hit", hit)
    }

    /// Document count for RAG
    pub fn document_count(count: usize) -> KeyValue {
        KeyValue::new("reasonkit.rag.document_count", count as i64)
    }
}

/// Metrics for ReasonKit operations
pub mod metrics {
    use std::sync::OnceLock;

    /// Counter for ThinkTool executions
    static THINKTOOL_EXECUTIONS: OnceLock<()> = OnceLock::new();

    /// Histogram for reasoning latency
    static REASONING_LATENCY: OnceLock<()> = OnceLock::new();

    /// Counter for cache hits/misses
    static CACHE_OPERATIONS: OnceLock<()> = OnceLock::new();

    /// Record a ThinkTool execution
    pub fn record_thinktool_execution(tool: &str, success: bool) {
        let _ = THINKTOOL_EXECUTIONS.get_or_init(|| ());
        tracing::debug!(
            tool = tool,
            success = success,
            "ThinkTool execution recorded"
        );
        // Full implementation would use opentelemetry::metrics
    }

    /// Record reasoning latency
    pub fn record_reasoning_latency(tool: &str, latency_ms: u64) {
        let _ = REASONING_LATENCY.get_or_init(|| ());
        tracing::debug!(
            tool = tool,
            latency_ms = latency_ms,
            "Reasoning latency recorded"
        );
    }

    /// Record cache operation
    pub fn record_cache_operation(cache: &str, hit: bool) {
        let _ = CACHE_OPERATIONS.get_or_init(|| ());
        tracing::debug!(cache = cache, hit = hit, "Cache operation recorded");
    }
}

/// Create a tracing span for a ThinkTool operation
#[macro_export]
macro_rules! thinktool_span {
    ($name:expr) => {
        tracing::info_span!("thinktool", otel.name = $name, thinktool.name = $name,)
    };
    ($name:expr, $profile:expr) => {
        tracing::info_span!(
            "thinktool",
            otel.name = $name,
            thinktool.name = $name,
            thinktool.profile = $profile,
        )
    };
}

/// Create a tracing span for a reasoning operation
#[macro_export]
macro_rules! reasoning_span {
    ($operation:expr) => {
        tracing::info_span!(
            "reasoning",
            otel.name = $operation,
            reasoning.operation = $operation,
        )
    };
}

/// Initialize a tracing subscriber with OpenTelemetry support.
///
/// This function sets up a global subscriber with:
/// - Environment-based log filtering (RUST_LOG)
/// - Formatted console output
/// - OpenTelemetry tracing layer (when enabled in config)
///
/// # Example
/// ```ignore
/// use reasonkit_core::integrations::observability::{OtelConfig, init_subscriber};
///
/// let config = OtelConfig::default();
/// let _guard = init_subscriber(&config)?;
/// // Subscriber is now active globally
/// ```
pub fn init_subscriber(config: &OtelConfig) -> Result<SubscriberGuard> {
    use tracing_subscriber::prelude::*;
    use tracing_subscriber::EnvFilter;

    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));

    let fmt_layer = tracing_subscriber::fmt::layer()
        .with_target(true)
        .with_thread_ids(true);

    // Build subscriber with optional OpenTelemetry layer
    if config.tracing_enabled {
        // Get tracer from the SDK provider (not global) for proper trait bounds
        use opentelemetry::trace::TracerProvider as _;
        use opentelemetry_sdk::trace::TracerProvider;

        let provider = TracerProvider::builder().build();
        let tracer = provider.tracer("reasonkit");

        let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);

        let subscriber = tracing_subscriber::registry()
            .with(env_filter)
            .with(fmt_layer)
            .with(otel_layer);

        tracing::subscriber::set_global_default(subscriber)
            .map_err(|e| anyhow::anyhow!("Failed to set subscriber: {}", e))?;
    } else {
        let subscriber = tracing_subscriber::registry()
            .with(env_filter)
            .with(fmt_layer);

        tracing::subscriber::set_global_default(subscriber)
            .map_err(|e| anyhow::anyhow!("Failed to set subscriber: {}", e))?;
    }

    Ok(SubscriberGuard { _private: () })
}

/// Guard that indicates the subscriber is active.
/// The subscriber remains active for the lifetime of the program.
pub struct SubscriberGuard {
    _private: (),
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_default() {
        let config = OtelConfig::default();
        assert_eq!(config.service_name, "reasonkit");
        assert!(config.tracing_enabled);
        assert!(config.metrics_enabled);
    }

    #[test]
    fn test_attributes() {
        let attr = attributes::thinktool_name("GigaThink");
        assert_eq!(attr.key.as_str(), "reasonkit.thinktool.name");
    }
}