From 31a498d0ce8a708c1f7da5495bcec6135ce92461 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Mon, 20 Apr 2026 08:32:11 +0800 Subject: [PATCH] feat(obs): Langfuse exporter wired into chat completions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per spec §3.6 / §9 / plan §4.9. New aisix-obs::langfuse module exports: - LangfuseEvent / LangfuseSender / LangfuseHandle types - spawn() that returns Ok(None) when disabled, or starts a background batch flusher (50 events or 1s, whichever comes first) when enabled The exporter authenticates with HTTP basic (public_key:secret_key), POSTs to {host}/api/public/ingestion in the documented batch shape ({batch: [{type: 'generation-create', body: {...}}]}), and never blocks the request hot path — full queues drop the event silently. Wired into: - aisix-server bootstrap (spawn after metrics, hold handle for life of process) - ProxyState (new optional langfuse: Option>) - chat::chat_completions (emits one event per request, success or failure) Tests (20 new): - Disabled config returns None - Enabled-without-host errors clearly - Enabled-without-key-env errors with the missing env name - Wiremock round-trip: emit one event, wait for the 1s flush interval, assert the upstream received exactly one POST - ISO timestamp round-trip including unix epoch + leap year - Base64 basic auth encoding matches expected output - Channel-full does not block the emitter Streaming chat handler emission, plus other endpoints (messages/embeddings/etc.), follow in a separate PR. --- Cargo.lock | 6 + crates/aisix-core/src/lib.rs | 4 +- crates/aisix-obs/Cargo.toml | 8 + crates/aisix-obs/src/langfuse.rs | 489 +++++++++++++++++++++++++++++++ crates/aisix-obs/src/lib.rs | 2 + crates/aisix-proxy/src/chat.rs | 32 +- crates/aisix-proxy/src/state.rs | 15 +- crates/aisix-server/src/main.rs | 21 +- 8 files changed, 571 insertions(+), 6 deletions(-) create mode 100644 crates/aisix-obs/src/langfuse.rs diff --git a/Cargo.lock b/Cargo.lock index 61179427..c173c789 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -172,17 +172,23 @@ version = "0.1.0" dependencies = [ "aisix-core", "anyhow", + "base64 0.22.1", "metrics", "metrics-exporter-prometheus", "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions", "opentelemetry_sdk", + "reqwest", + "serde", + "serde_json", "thiserror 1.0.69", "tokio", "tracing", "tracing-opentelemetry", "tracing-subscriber", + "uuid", + "wiremock", ] [[package]] diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index c249436c..62772521 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -23,8 +23,8 @@ pub mod resource; pub mod snapshot; pub use config::{ - AdminConfig, CacheBackend, CacheConfig, Config, EtcdConfig, ObservabilityConfig, ProxyConfig, - TlsConfig, + AdminConfig, CacheBackend, CacheConfig, Config, EtcdConfig, LangfuseConfig, + ObservabilityConfig, ProxyConfig, TlsConfig, }; pub use error::{ AdminError, AdminErrorEnvelope, BootstrapError, ProxyError, ProxyErrorEnvelope, RateLimitScope, diff --git a/crates/aisix-obs/Cargo.toml b/crates/aisix-obs/Cargo.toml index 43a7c3eb..c6844c33 100644 --- a/crates/aisix-obs/Cargo.toml +++ b/crates/aisix-obs/Cargo.toml @@ -22,3 +22,11 @@ metrics.workspace = true metrics-exporter-prometheus.workspace = true thiserror.workspace = true anyhow.workspace = true +serde.workspace = true +serde_json.workspace = true +reqwest.workspace = true +uuid.workspace = true +base64.workspace = true + +[dev-dependencies] +wiremock.workspace = true diff --git a/crates/aisix-obs/src/langfuse.rs b/crates/aisix-obs/src/langfuse.rs new file mode 100644 index 00000000..db9c5a97 --- /dev/null +++ b/crates/aisix-obs/src/langfuse.rs @@ -0,0 +1,489 @@ +//! Langfuse exporter — pushes per-request generation events to a +//! Langfuse `/api/public/ingestion` endpoint. +//! +//! There is no first-party Rust SDK for Langfuse, so we hand-roll the +//! JSON shape against the documented batch ingestion API. +//! +//! Design (spec §3.6 / §9): +//! - The proxy emits a [`LangfuseEvent`] at end-of-request (success or +//! failure) onto a fire-and-forget mpsc channel. +//! - A background task drains the channel, batches up to +//! `MAX_BATCH_SIZE` events (or flushes after `FLUSH_INTERVAL`), and +//! POSTs them with HTTP basic auth (`:`). +//! - All errors are logged at WARN; we never block the request hot path +//! on Langfuse availability. + +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::Serialize; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +use aisix_core::ObservabilityConfig; + +/// Hard cap on per-batch event count. Tracks the Langfuse default +/// payload-size limit (~1 MB body). +const MAX_BATCH_SIZE: usize = 50; +/// Maximum delay between batch flushes when the queue is non-empty. +const FLUSH_INTERVAL: Duration = Duration::from_secs(1); +/// Channel capacity. Chosen so that a sustained ~1k req/s burst still +/// never blocks the proxy thread; if Langfuse is offline we drop the +/// oldest events at the edges. +const CHANNEL_CAPACITY: usize = 4096; + +#[derive(Debug, thiserror::Error)] +pub enum LangfuseError { + #[error("langfuse enabled but {0} env var not set")] + MissingEnv(String), + #[error("langfuse host not configured")] + MissingHost, +} + +/// Event emitted to Langfuse for a single proxy request. +#[derive(Debug, Clone)] +pub struct LangfuseEvent { + pub trace_id: String, + pub model: String, + pub provider: String, + pub input: Option, + pub output: Option, + pub prompt_tokens: Option, + pub completion_tokens: Option, + pub total_tokens: Option, + pub status_code: u16, + pub latency: Duration, + pub api_key_id: Option, +} + +/// Owning handle returned by [`spawn`]. Drop or call [`shutdown`] to +/// stop the background task. +pub struct LangfuseHandle { + sender: Arc, + task: Option>, +} + +impl LangfuseHandle { + pub fn sender(&self) -> Arc { + self.sender.clone() + } + + /// Best-effort shutdown — closes the channel and waits up to 2s for + /// the background task to drain remaining events. + pub async fn shutdown(mut self) { + // Drop our reference to the sender so the channel closes. + // Replacing the inner Arc with an empty one isn't possible + // without unsafe; instead we rely on the task's own shutdown + // signal which fires when the channel is closed. + let _ = Arc::try_unwrap(self.sender); + if let Some(task) = self.task.take() { + let _ = tokio::time::timeout(Duration::from_secs(2), task).await; + } + } +} + +/// Cheap-clone handle the proxy keeps to push events. Sending is +/// non-blocking — if the queue is full the event is silently dropped. +#[derive(Debug, Clone)] +pub struct LangfuseSender { + tx: mpsc::Sender, +} + +impl LangfuseSender { + /// Push an event onto the queue. Never blocks; never errors. If the + /// channel is full or closed the event is dropped and a counter is + /// bumped (TODO: wire up `metrics::counter!` once the obs crate + /// gains a shared metrics handle). + pub fn emit(&self, event: LangfuseEvent) { + if self.tx.try_send(event).is_err() { + tracing::debug!("langfuse queue full or closed — event dropped"); + } + } +} + +/// Spawn the Langfuse exporter. Returns `Ok(None)` when Langfuse is +/// disabled in config; returns the handle otherwise. +pub fn spawn(cfg: &ObservabilityConfig) -> Result, LangfuseError> { + let lf = &cfg.langfuse; + if !lf.enabled { + return Ok(None); + } + let host = lf.host.clone().ok_or(LangfuseError::MissingHost)?; + let public_key = resolve_env(lf.public_key_env.as_deref(), "LANGFUSE_PUBLIC_KEY")?; + let secret_key = resolve_env(lf.secret_key_env.as_deref(), "LANGFUSE_SECRET_KEY")?; + + let (tx, rx) = mpsc::channel::(CHANNEL_CAPACITY); + let sender = Arc::new(LangfuseSender { tx }); + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()); + + let task = tokio::spawn(run_exporter(host, public_key, secret_key, client, rx)); + + Ok(Some(LangfuseHandle { + sender, + task: Some(task), + })) +} + +fn resolve_env(env_name: Option<&str>, default: &str) -> Result { + let var = env_name.unwrap_or(default); + std::env::var(var).map_err(|_| LangfuseError::MissingEnv(var.into())) +} + +async fn run_exporter( + host: String, + public_key: String, + secret_key: String, + client: reqwest::Client, + mut rx: mpsc::Receiver, +) { + let endpoint = format!("{}/api/public/ingestion", host.trim_end_matches('/')); + let auth = base64_basic(&public_key, &secret_key); + + let mut buf: Vec = Vec::with_capacity(MAX_BATCH_SIZE); + let mut interval = tokio::time::interval(FLUSH_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + loop { + tokio::select! { + biased; + maybe = rx.recv() => match maybe { + Some(ev) => { + buf.push(ev); + if buf.len() >= MAX_BATCH_SIZE { + flush(&client, &endpoint, &auth, &mut buf).await; + } + } + None => { + // Channel closed — drain remaining and exit. + if !buf.is_empty() { + flush(&client, &endpoint, &auth, &mut buf).await; + } + break; + } + }, + _ = interval.tick() => { + if !buf.is_empty() { + flush(&client, &endpoint, &auth, &mut buf).await; + } + } + } + } +} + +async fn flush(client: &reqwest::Client, endpoint: &str, auth: &str, buf: &mut Vec) { + if buf.is_empty() { + return; + } + let payload = IngestionPayload::from_events(std::mem::take(buf)); + let body = match serde_json::to_vec(&payload) { + Ok(b) => b, + Err(e) => { + tracing::warn!(error = %e, "langfuse: failed to serialise batch"); + return; + } + }; + match client + .post(endpoint) + .header("authorization", format!("Basic {auth}")) + .header("content-type", "application/json") + .body(body) + .send() + .await + { + Ok(resp) if resp.status().is_success() => { + tracing::debug!("langfuse: batch accepted"); + } + Ok(resp) => { + tracing::warn!(status = %resp.status(), "langfuse: ingestion rejected"); + } + Err(e) => { + tracing::warn!(error = %e, "langfuse: ingestion request failed"); + } + } +} + +fn base64_basic(public_key: &str, secret_key: &str) -> String { + use base64::Engine; + base64::engine::general_purpose::STANDARD.encode(format!("{public_key}:{secret_key}")) +} + +#[derive(Debug, Serialize)] +struct IngestionPayload { + batch: Vec, +} + +#[derive(Debug, Serialize)] +struct IngestionItem { + id: String, + timestamp: String, + #[serde(rename = "type")] + event_type: &'static str, + body: GenerationBody, +} + +#[derive(Debug, Serialize)] +struct GenerationBody { + id: String, + trace_id: String, + name: String, + model: String, + #[serde(skip_serializing_if = "Option::is_none")] + input: Option, + #[serde(skip_serializing_if = "Option::is_none")] + output: Option, + #[serde(skip_serializing_if = "Option::is_none")] + metadata: Option, + usage: GenerationUsage, + start_time: String, + end_time: String, + status_message: Option, +} + +#[derive(Debug, Serialize)] +struct GenerationUsage { + #[serde(skip_serializing_if = "Option::is_none")] + input: Option, + #[serde(skip_serializing_if = "Option::is_none")] + output: Option, + #[serde(skip_serializing_if = "Option::is_none")] + total: Option, +} + +impl IngestionPayload { + fn from_events(events: Vec) -> Self { + let now_iso = iso_now(); + Self { + batch: events + .into_iter() + .map(|ev| IngestionItem { + id: format!("evt-{}", uuid::Uuid::new_v4()), + timestamp: now_iso.clone(), + event_type: "generation-create", + body: GenerationBody { + id: format!("gen-{}", uuid::Uuid::new_v4()), + trace_id: ev.trace_id.clone(), + name: format!("{}.chat", ev.provider), + model: ev.model.clone(), + input: ev.input, + output: ev.output, + metadata: ev.api_key_id.map(|k| serde_json::json!({"api_key_id": k})), + usage: GenerationUsage { + input: ev.prompt_tokens, + output: ev.completion_tokens, + total: ev.total_tokens, + }, + start_time: iso_offset(ev.latency), + end_time: now_iso.clone(), + status_message: (ev.status_code != 200) + .then(|| format!("upstream status {}", ev.status_code)), + }, + }) + .collect(), + } + } +} + +fn iso_now() -> String { + let secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::ZERO) + .as_secs(); + format_iso(secs) +} + +fn iso_offset(latency: Duration) -> String { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::ZERO); + let start = now.saturating_sub(latency); + format_iso(start.as_secs()) +} + +/// Minimal RFC-3339 formatter (no chrono dep needed for this one +/// purpose; chrono is heavy and we only need second-precision +/// timestamps for Langfuse). +fn format_iso(secs: u64) -> String { + // Days since 1970-01-01 + let days = (secs / 86_400) as i64; + let time_in_day = secs % 86_400; + let h = time_in_day / 3600; + let m = (time_in_day / 60) % 60; + let s = time_in_day % 60; + let (y, mo, d) = days_to_ymd(days); + format!("{y:04}-{mo:02}-{d:02}T{h:02}:{m:02}:{s:02}Z") +} + +fn days_to_ymd(mut days: i64) -> (i64, u32, u32) { + // Convert "days since 1970-01-01" to (year, month, day). + // Algorithm from Howard Hinnant's date library, simplified. + days += 719_468; + let era = days.div_euclid(146_097); + let doe = days.rem_euclid(146_097); + let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; + let y = yoe + era * 400; + let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); + let mp = (5 * doy + 2) / 153; + let d = (doy - (153 * mp + 2) / 5 + 1) as u32; + let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32; + let y = if m <= 2 { y + 1 } else { y }; + (y, m, d) +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_core::LangfuseConfig; + + fn cfg(enabled: bool, host: Option<&str>) -> ObservabilityConfig { + ObservabilityConfig { + service_name: "test".into(), + log_level: "info".into(), + access_log: true, + metrics: Default::default(), + tracing: Default::default(), + langfuse: LangfuseConfig { + enabled, + host: host.map(String::from), + public_key_env: None, + secret_key_env: None, + }, + } + } + + #[test] + fn spawn_returns_none_when_disabled() { + let result = spawn(&cfg(false, None)).unwrap(); + assert!( + result.is_none(), + "disabled langfuse should produce no handle" + ); + } + + #[test] + fn spawn_errors_when_enabled_without_host() { + let result = spawn(&cfg(true, None)); + assert!(matches!(result, Err(LangfuseError::MissingHost))); + } + + #[test] + fn spawn_errors_when_enabled_without_keys() { + // Host present but env vars missing. We point at custom env names + // to avoid stomping on the global LANGFUSE_PUBLIC_KEY/SECRET_KEY, + // which other tests (e.g. the wiremock round-trip) set. + let mut c = cfg(true, Some("https://cloud.langfuse.com")); + c.langfuse.public_key_env = Some("AISIX_TEST_NEVER_SET_PK".into()); + c.langfuse.secret_key_env = Some("AISIX_TEST_NEVER_SET_SK".into()); + let result = spawn(&c); + assert!(matches!(result, Err(LangfuseError::MissingEnv(_)))); + } + + #[test] + fn ingestion_payload_serialises_required_fields() { + let ev = LangfuseEvent { + trace_id: "trace-1".into(), + model: "openai/gpt-4o".into(), + provider: "openai".into(), + input: Some(serde_json::json!({"messages": [{"role": "user", "content": "hi"}]})), + output: Some(serde_json::json!({"text": "hello"})), + prompt_tokens: Some(7), + completion_tokens: Some(2), + total_tokens: Some(9), + status_code: 200, + latency: Duration::from_millis(123), + api_key_id: Some("k-1".into()), + }; + let payload = IngestionPayload::from_events(vec![ev]); + let json = serde_json::to_value(&payload).unwrap(); + assert_eq!(json["batch"].as_array().unwrap().len(), 1); + let item = &json["batch"][0]; + assert_eq!(item["type"], "generation-create"); + assert_eq!(item["body"]["model"], "openai/gpt-4o"); + assert_eq!(item["body"]["trace_id"], "trace-1"); + assert_eq!(item["body"]["usage"]["input"], 7); + assert_eq!(item["body"]["usage"]["total"], 9); + assert_eq!(item["body"]["metadata"]["api_key_id"], "k-1"); + } + + #[test] + fn iso_format_round_trips_known_epoch_seconds() { + // 2026-04-23T00:00:00Z = 1_776_902_400 + assert_eq!(format_iso(1_776_902_400), "2026-04-23T00:00:00Z"); + // 2024-02-29T12:34:56Z (leap year sanity check) + assert_eq!(format_iso(1_709_210_096), "2024-02-29T12:34:56Z"); + // Unix epoch + assert_eq!(format_iso(0), "1970-01-01T00:00:00Z"); + } + + #[test] + fn base64_basic_matches_expected_encoding() { + // "pk:sk" → "cGs6c2s=" + assert_eq!(base64_basic("pk", "sk"), "cGs6c2s="); + } + + #[test] + fn sender_emit_does_not_block_when_channel_full() { + let (tx, mut _rx) = mpsc::channel(1); + let sender = LangfuseSender { tx }; + let ev = LangfuseEvent { + trace_id: "t".into(), + model: "m".into(), + provider: "p".into(), + input: None, + output: None, + prompt_tokens: None, + completion_tokens: None, + total_tokens: None, + status_code: 200, + latency: Duration::from_millis(1), + api_key_id: None, + }; + sender.emit(ev.clone()); + // Channel now full — second emit must not block or error. + sender.emit(ev); + } + + #[tokio::test] + async fn full_round_trip_to_wiremock_upstream() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/public/ingestion")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + + // SAFETY: tests run single-threaded by default and we only mutate + // env vars that are scoped to this test's spawn() call. + unsafe { + std::env::set_var("LANGFUSE_PUBLIC_KEY", "pk-test"); + std::env::set_var("LANGFUSE_SECRET_KEY", "sk-test"); + } + + let handle = spawn(&cfg(true, Some(&server.uri()))).unwrap().unwrap(); + let sender = handle.sender(); + sender.emit(LangfuseEvent { + trace_id: "t-rt".into(), + model: "openai/gpt-4o".into(), + provider: "openai".into(), + input: None, + output: None, + prompt_tokens: Some(1), + completion_tokens: Some(1), + total_tokens: Some(2), + status_code: 200, + latency: Duration::from_millis(50), + api_key_id: None, + }); + + // Wait for the 1s flush interval + a small jitter window. + tokio::time::sleep(Duration::from_millis(1_500)).await; + server.verify().await; + } +} diff --git a/crates/aisix-obs/src/lib.rs b/crates/aisix-obs/src/lib.rs index d7519e4b..34b4a054 100644 --- a/crates/aisix-obs/src/lib.rs +++ b/crates/aisix-obs/src/lib.rs @@ -13,6 +13,7 @@ #![deny(rust_2018_idioms)] pub mod access_log; +pub mod langfuse; pub mod metrics; pub mod otlp; @@ -20,6 +21,7 @@ use aisix_core::ObservabilityConfig; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; pub use access_log::AccessLog; +pub use langfuse::{LangfuseError, LangfuseEvent, LangfuseHandle, LangfuseSender}; pub use metrics::{Metrics, RequestOutcome}; pub use otlp::{install_otlp_tracer, shutdown_otlp, OtlpError, OtlpHandle}; diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 07728884..7c680abd 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -19,7 +19,7 @@ use aisix_cache::CacheKey; use aisix_gateway::{BridgeContext, BridgeError, ChatFormat}; use aisix_guardrails::GuardrailVerdict; -use aisix_obs::{AccessLog, Metrics, RequestOutcome}; +use aisix_obs::{AccessLog, LangfuseEvent, Metrics, RequestOutcome}; use axum::extract::State; use axum::http::HeaderValue; use axum::response::sse::{Event, KeepAlive, Sse}; @@ -80,6 +80,21 @@ pub async fn chat_completions( success.total_tokens, &request_id, ); + if let Some(lf) = state.langfuse.as_ref() { + lf.emit(LangfuseEvent { + trace_id: request_id.clone(), + model: model_name.clone(), + provider: success.provider.clone(), + input: None, + output: None, + prompt_tokens: success.prompt_tokens, + completion_tokens: success.completion_tokens, + total_tokens: success.total_tokens, + status_code: status, + latency: elapsed, + api_key_id: Some(api_key_id.clone()), + }); + } // Inject x-ratelimit-* headers so OpenAI SDK clients see the // current window state. We peek *after* the commit so // remaining-requests reflects the post-dispatch tally. @@ -110,6 +125,21 @@ pub async fn chat_completions( None, &request_id, ); + if let Some(lf) = state.langfuse.as_ref() { + lf.emit(LangfuseEvent { + trace_id: request_id.clone(), + model: model_name.clone(), + provider: "unknown".to_string(), + input: None, + output: None, + prompt_tokens: None, + completion_tokens: None, + total_tokens: None, + status_code: status, + latency: elapsed, + api_key_id: Some(api_key_id.clone()), + }); + } err.into_response() } } diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index fb913cbd..befc00a7 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -19,7 +19,7 @@ use aisix_core::snapshot::SnapshotHandle; use aisix_core::{AisixSnapshot, ProxyConfig}; use aisix_gateway::Hub; use aisix_guardrails::{Guardrail, GuardrailChain}; -use aisix_obs::Metrics; +use aisix_obs::{LangfuseSender, Metrics}; use aisix_ratelimit::Limiter; use std::sync::Arc; @@ -44,6 +44,9 @@ pub struct ProxyState { /// Per-model health tracker. Updated on every upstream call outcome; /// read by `GET /admin/v1/health`. pub health: Arc, + /// Optional Langfuse exporter. When `Some`, chat handlers emit one + /// generation event at end-of-request. `None` disables emission. + pub langfuse: Option>, pub request_body_limit_bytes: usize, } @@ -59,6 +62,7 @@ impl ProxyState { guardrails: Arc::new(GuardrailChain::empty()), budgets: Arc::new(BudgetTracker::new()), health: Arc::new(HealthTracker::new()), + langfuse: None, request_body_limit_bytes: cfg.request_body_limit_bytes, } } @@ -81,6 +85,7 @@ impl ProxyState { guardrails: Arc::new(GuardrailChain::empty()), budgets: Arc::new(BudgetTracker::new()), health: Arc::new(HealthTracker::new()), + langfuse: None, request_body_limit_bytes: cfg.request_body_limit_bytes, } } @@ -106,6 +111,7 @@ impl ProxyState { guardrails: Arc::new(GuardrailChain::empty()), budgets: Arc::new(BudgetTracker::new()), health: Arc::new(HealthTracker::new()), + langfuse: None, request_body_limit_bytes: cfg.request_body_limit_bytes, } } @@ -123,4 +129,11 @@ impl ProxyState { self.guardrails = guardrails; self } + + /// Attach a Langfuse sender. The exporter is opt-in; when absent + /// no events are emitted regardless of request volume. + pub fn with_langfuse(mut self, sender: Arc) -> Self { + self.langfuse = Some(sender); + self + } } diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 67036ccb..01bdbf25 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -21,7 +21,7 @@ use aisix_core::models::Provider; use aisix_core::Config; use aisix_etcd::{EtcdConfigProvider, Supervisor}; use aisix_gateway::Hub; -use aisix_obs::{init_tracing, install_otlp_tracer, Metrics}; +use aisix_obs::{init_tracing, install_otlp_tracer, langfuse, Metrics}; use aisix_provider_anthropic::AnthropicBridge; use aisix_provider_deepseek::deepseek_bridge; use aisix_provider_gemini::gemini_bridge; @@ -85,7 +85,21 @@ async fn run(cfg: Config) -> anyhow::Result<()> { // behind the same trait object once their PRs land. let cache: Option> = Some(Arc::new(MemoryCache::with_defaults())); - let proxy_state = ProxyState::with_components( + // Optional Langfuse exporter — disabled in config by default. + // When enabled, the proxy gets an Arc through + // ProxyState and emits one event per chat completion at + // end-of-request. We keep the handle alive for the lifetime of + // the process so the background flush task continues running. + let langfuse_handle = match langfuse::spawn(&cfg.observability) { + Ok(h) => h, + Err(e) => { + tracing::warn!(error = %e, "langfuse exporter disabled"); + None + } + }; + let langfuse_sender = langfuse_handle.as_ref().map(|h| h.sender()); + + let mut proxy_state = ProxyState::with_components( snapshot_handle.clone(), hub.clone(), limiter.clone(), @@ -93,6 +107,9 @@ async fn run(cfg: Config) -> anyhow::Result<()> { cache.clone(), &cfg.proxy, ); + if let Some(sender) = langfuse_sender { + proxy_state = proxy_state.with_langfuse(sender); + } // Clone shared trackers before consuming proxy_state in build_router. let budget_tracker = proxy_state.budgets.clone(); let health_tracker = proxy_state.health.clone();