litellm-rs 0.6.0

A high-performance AI Gateway written in Rust, providing OpenAI-compatible APIs with intelligent routing, load balancing, and enterprise features
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

//! Adapter from the canonical integration lifecycle to Langfuse logging.

use async_trait::async_trait;

use super::{LangfuseConfig, LangfuseError, LangfuseLogger, LlmError, LlmRequest, LlmResponse};
use crate::core::traits::integration::{
    Integration, IntegrationError, IntegrationResult, LlmEndEvent, LlmErrorEvent, LlmStartEvent,
};

/// Canonical callback adapter for the existing Langfuse logger.
pub struct LangfuseIntegration {
    logger: LangfuseLogger,
}

impl LangfuseIntegration {
    /// Build a Langfuse callback integration.
    pub fn new(config: LangfuseConfig) -> Result<Self, LangfuseError> {
        LangfuseLogger::new(config).map(|logger| Self { logger })
    }
}

#[async_trait]
impl Integration for LangfuseIntegration {
    fn name(&self) -> &'static str {
        "langfuse"
    }

    fn is_enabled(&self) -> bool {
        true
    }

    async fn on_llm_start(&self, event: &LlmStartEvent) -> IntegrationResult<()> {
        let mut request = LlmRequest::new(&event.request_id, &event.model);
        request.provider = event.provider.clone();
        request.user_id = event.user_id.clone();
        request.session_id = event.session_id.clone();
        request.metadata = event.metadata.clone();
        request.tags = event.tags.clone();
        request.parameters = event.parameters.clone();
        self.logger.on_llm_start(request);
        Ok(())
    }

    async fn on_llm_end(&self, event: &LlmEndEvent) -> IntegrationResult<()> {
        let mut response = LlmResponse::new(&event.request_id);
        response.input_tokens = event.input_tokens;
        response.output_tokens = event.output_tokens;
        response.cost = event.cost_usd;
        response.metadata = event.metadata.clone();
        if let Some(provider) = &event.provider {
            response
                .metadata
                .insert("provider".to_string(), serde_json::json!(provider));
        }
        response.metadata.insert(
            "latency_ms".to_string(),
            serde_json::json!(event.latency_ms),
        );
        self.logger.on_llm_end(response);
        Ok(())
    }

    async fn on_llm_error(&self, event: &LlmErrorEvent) -> IntegrationResult<()> {
        let mut error = LlmError::new(&event.request_id, &event.error_message);
        error.error_type = event.error_type.clone();
        error.metadata = event.metadata.clone();
        if let Some(provider) = &event.provider {
            error
                .metadata
                .insert("provider".to_string(), serde_json::json!(provider));
        }
        self.logger.on_llm_error(error);
        Ok(())
    }

    async fn flush(&self) -> IntegrationResult<()> {
        self.logger
            .flush()
            .await
            .map_err(|error| IntegrationError::other(format!("Langfuse flush failed: {error}")))
    }

    async fn shutdown(&self) -> IntegrationResult<()> {
        self.logger
            .shutdown_ref()
            .await
            .map_err(IntegrationError::other)
    }
}

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

    fn debug_config() -> LangfuseConfig {
        LangfuseConfig {
            public_key: Some("pk-test".to_string()),
            secret_key: Some("sk-test".to_string()),
            debug: true,
            flush_interval_ms: 60_000,
            ..LangfuseConfig::default()
        }
    }

    #[tokio::test]
    async fn adapter_tracks_canonical_start_and_end_without_content() {
        let integration = match LangfuseIntegration::new(debug_config()) {
            Ok(integration) => integration,
            Err(error) => panic!("Langfuse adapter should initialize: {error}"),
        };
        let start = LlmStartEvent::new("req-1", "model").provider("provider");
        assert!(integration.on_llm_start(&start).await.is_ok());
        assert_eq!(integration.logger.active_count(), 1);

        let end = LlmEndEvent::new("req-1", "model")
            .provider("provider")
            .tokens(10, 5)
            .cost(0.25)
            .latency(20);
        assert!(integration.on_llm_end(&end).await.is_ok());
        assert_eq!(integration.logger.active_count(), 0);
    }

    #[tokio::test]
    async fn adapter_shutdown_stops_and_joins_logger_worker() {
        let integration = LangfuseIntegration::new(debug_config()).unwrap();

        assert!(Integration::shutdown(&integration).await.is_ok());
        assert!(integration.logger.worker_stopped());
    }
}