reasonkit-core 0.1.8

//! LLM-Specific Observability with Langfuse
//!
//! This module provides AI-specific insights including trace/span models,
//! cost tracking, latency measurement, and evaluation scoring.
//!
//! # Features
//! - Trace/span model for agent execution
//! - Token usage and cost tracking
//! - Latency measurement per step
//! - Quality scores and evaluations
//!
//! Enable with: `cargo build --features llm-observability`

use serde::{Deserialize, Serialize};
use std::time::Duration;
use uuid::Uuid;

// Re-export langfuse for direct access
pub use langfuse;

/// Configuration for Langfuse integration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LangfuseConfig {
    /// Langfuse public key
    pub public_key: String,
    /// Langfuse secret key
    pub secret_key: String,
    /// Langfuse host URL
    pub host: String,
    /// Enable debug mode
    pub debug: bool,
    /// Batch size for async uploads
    pub batch_size: usize,
    /// Flush interval in seconds
    pub flush_interval_secs: u64,
}

impl Default for LangfuseConfig {
    fn default() -> Self {
        Self {
            public_key: std::env::var("LANGFUSE_PUBLIC_KEY").unwrap_or_default(),
            secret_key: std::env::var("LANGFUSE_SECRET_KEY").unwrap_or_default(),
            host: std::env::var("LANGFUSE_HOST")
                .unwrap_or_else(|_| "https://cloud.langfuse.com".to_string()),
            debug: false,
            batch_size: 100,
            flush_interval_secs: 5,
        }
    }
}

/// A trace representing a complete agent execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Trace {
    pub id: String,
    pub name: String,
    pub user_id: Option<String>,
    pub session_id: Option<String>,
    pub metadata: serde_json::Value,
    pub input: Option<String>,
    pub output: Option<String>,
    pub start_time: chrono::DateTime<chrono::Utc>,
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    pub spans: Vec<Span>,
    pub scores: Vec<Score>,
}

impl Trace {
    /// Create a new trace
    pub fn new(name: &str) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            user_id: None,
            session_id: None,
            metadata: serde_json::json!({}),
            input: None,
            output: None,
            start_time: chrono::Utc::now(),
            end_time: None,
            spans: Vec::new(),
            scores: Vec::new(),
        }
    }

    /// Set the user ID
    pub fn with_user(mut self, user_id: &str) -> Self {
        self.user_id = Some(user_id.to_string());
        self
    }

    /// Set the session ID
    pub fn with_session(mut self, session_id: &str) -> Self {
        self.session_id = Some(session_id.to_string());
        self
    }

    /// Set the input
    pub fn with_input(mut self, input: &str) -> Self {
        self.input = Some(input.to_string());
        self
    }

    /// Add a span to the trace
    pub fn add_span(&mut self, span: Span) {
        self.spans.push(span);
    }

    /// Add a score to the trace
    pub fn add_score(&mut self, score: Score) {
        self.scores.push(score);
    }

    /// End the trace
    pub fn end(&mut self, output: Option<&str>) {
        self.end_time = Some(chrono::Utc::now());
        self.output = output.map(|s| s.to_string());
    }

    /// Get the total duration
    pub fn duration(&self) -> Option<Duration> {
        self.end_time.map(|end| {
            let start = self.start_time;
            (end - start).to_std().unwrap_or_default()
        })
    }
}

/// A span representing a single step in the trace
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Span {
    pub id: String,
    pub name: String,
    pub span_type: SpanType,
    pub input: Option<String>,
    pub output: Option<String>,
    pub model: Option<String>,
    pub model_parameters: Option<serde_json::Value>,
    pub usage: Option<TokenUsage>,
    pub start_time: chrono::DateTime<chrono::Utc>,
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    pub metadata: serde_json::Value,
    pub level: SpanLevel,
    pub status: SpanStatus,
}

/// Type of span
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum SpanType {
    /// LLM generation
    Generation,
    /// Tool/function call
    Tool,
    /// Retrieval operation
    Retrieval,
    /// Embedding generation
    Embedding,
    /// Generic span
    Span,
}

/// Span severity level
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum SpanLevel {
    Debug,
    Default,
    Warning,
    Error,
}

/// Span completion status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum SpanStatus {
    Success,
    Error,
    Pending,
}

impl Span {
    /// Create a new generation span
    pub fn generation(name: &str, model: &str) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            span_type: SpanType::Generation,
            input: None,
            output: None,
            model: Some(model.to_string()),
            model_parameters: None,
            usage: None,
            start_time: chrono::Utc::now(),
            end_time: None,
            metadata: serde_json::json!({}),
            level: SpanLevel::Default,
            status: SpanStatus::Pending,
        }
    }

    /// Create a new tool span
    pub fn tool(name: &str) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            span_type: SpanType::Tool,
            input: None,
            output: None,
            model: None,
            model_parameters: None,
            usage: None,
            start_time: chrono::Utc::now(),
            end_time: None,
            metadata: serde_json::json!({}),
            level: SpanLevel::Default,
            status: SpanStatus::Pending,
        }
    }

    /// Create a new retrieval span
    pub fn retrieval(name: &str) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            span_type: SpanType::Retrieval,
            input: None,
            output: None,
            model: None,
            model_parameters: None,
            usage: None,
            start_time: chrono::Utc::now(),
            end_time: None,
            metadata: serde_json::json!({}),
            level: SpanLevel::Default,
            status: SpanStatus::Pending,
        }
    }

    /// Set the input
    pub fn with_input(mut self, input: &str) -> Self {
        self.input = Some(input.to_string());
        self
    }

    /// End the span successfully
    pub fn success(mut self, output: &str) -> Self {
        self.output = Some(output.to_string());
        self.end_time = Some(chrono::Utc::now());
        self.status = SpanStatus::Success;
        self
    }

    /// End the span with an error
    pub fn error(mut self, error: &str) -> Self {
        self.output = Some(error.to_string());
        self.end_time = Some(chrono::Utc::now());
        self.status = SpanStatus::Error;
        self.level = SpanLevel::Error;
        self
    }

    /// Set token usage
    pub fn with_usage(mut self, usage: TokenUsage) -> Self {
        self.usage = Some(usage);
        self
    }
}

/// Token usage statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
    pub prompt_tokens: usize,
    pub completion_tokens: usize,
    pub total_tokens: usize,
}

impl TokenUsage {
    /// Create new token usage
    pub fn new(prompt_tokens: usize, completion_tokens: usize) -> Self {
        Self {
            prompt_tokens,
            completion_tokens,
            total_tokens: prompt_tokens + completion_tokens,
        }
    }

    /// Estimate cost based on model pricing
    pub fn estimate_cost(&self, model: &str) -> f64 {
        // Prices per 1K tokens (approximate)
        let (prompt_price, completion_price) = match model.to_lowercase().as_str() {
            m if m.contains("gpt-4-turbo") => (0.01, 0.03),
            m if m.contains("gpt-4o") => (0.005, 0.015),
            m if m.contains("gpt-4") => (0.03, 0.06),
            m if m.contains("gpt-3.5") => (0.0005, 0.0015),
            m if m.contains("claude-3-opus") => (0.015, 0.075),
            m if m.contains("claude-3-sonnet") => (0.003, 0.015),
            m if m.contains("claude-3-haiku") => (0.00025, 0.00125),
            _ => (0.001, 0.002), // Conservative estimate
        };

        let prompt_cost = (self.prompt_tokens as f64 / 1000.0) * prompt_price;
        let completion_cost = (self.completion_tokens as f64 / 1000.0) * completion_price;

        prompt_cost + completion_cost
    }
}

/// Quality score for a trace
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Score {
    pub id: String,
    pub name: String,
    pub value: f64,
    pub comment: Option<String>,
    pub source: ScoreSource,
}

/// Source of the score
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum ScoreSource {
    /// Automated evaluation
    Auto,
    /// Human evaluation
    Human,
    /// External system
    External,
}

impl Score {
    /// Create a new automated score
    pub fn auto(name: &str, value: f64) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            value,
            comment: None,
            source: ScoreSource::Auto,
        }
    }

    /// Create a new human score
    pub fn human(name: &str, value: f64, comment: Option<&str>) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name: name.to_string(),
            value,
            comment: comment.map(|s| s.to_string()),
            source: ScoreSource::Human,
        }
    }
}

/// Trace builder for ergonomic trace creation
pub struct TraceBuilder {
    trace: Trace,
    current_span: Option<Span>,
}

impl TraceBuilder {
    /// Start a new trace
    pub fn start(name: &str) -> Self {
        Self {
            trace: Trace::new(name),
            current_span: None,
        }
    }

    /// Set user ID
    pub fn user(mut self, user_id: &str) -> Self {
        self.trace = self.trace.with_user(user_id);
        self
    }

    /// Set session ID
    pub fn session(mut self, session_id: &str) -> Self {
        self.trace = self.trace.with_session(session_id);
        self
    }

    /// Set input
    pub fn input(mut self, input: &str) -> Self {
        self.trace = self.trace.with_input(input);
        self
    }

    /// Start a generation span
    pub fn generation(&mut self, name: &str, model: &str) {
        self.flush_current_span();
        self.current_span = Some(Span::generation(name, model));
    }

    /// Start a tool span
    pub fn tool(&mut self, name: &str) {
        self.flush_current_span();
        self.current_span = Some(Span::tool(name));
    }

    /// Start a retrieval span
    pub fn retrieval(&mut self, name: &str) {
        self.flush_current_span();
        self.current_span = Some(Span::retrieval(name));
    }

    /// End the current span successfully
    pub fn span_success(&mut self, output: &str) {
        if let Some(span) = self.current_span.take() {
            self.trace.add_span(span.success(output));
        }
    }

    /// End the current span with error
    pub fn span_error(&mut self, error: &str) {
        if let Some(span) = self.current_span.take() {
            self.trace.add_span(span.error(error));
        }
    }

    /// Add a score
    pub fn score(&mut self, name: &str, value: f64) {
        self.trace.add_score(Score::auto(name, value));
    }

    /// Finish the trace and return it
    pub fn finish(mut self, output: Option<&str>) -> Trace {
        self.flush_current_span();
        self.trace.end(output);
        self.trace
    }

    fn flush_current_span(&mut self) {
        if let Some(span) = self.current_span.take() {
            // If span wasn't explicitly ended, mark as success with empty output
            self.trace.add_span(span.success(""));
        }
    }
}

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

    #[test]
    fn test_trace_creation() {
        let trace = Trace::new("test-trace")
            .with_user("user-123")
            .with_input("test input");

        assert_eq!(trace.name, "test-trace");
        assert_eq!(trace.user_id, Some("user-123".to_string()));
    }

    #[test]
    fn test_token_usage_cost() {
        let usage = TokenUsage::new(1000, 500);
        let cost = usage.estimate_cost("gpt-4-turbo");
        assert!(cost > 0.0);
    }

    #[test]
    fn test_trace_builder() {
        let mut builder = TraceBuilder::start("test").user("user-1").input("hello");

        builder.generation("gen1", "gpt-4");
        builder.span_success("world");

        let trace = builder.finish(Some("complete"));
        assert_eq!(trace.spans.len(), 1);
        assert!(trace.output.is_some());
    }
}