Skip to main content

lc_observability/
mongo.rs

1// lc-observability/src/mongo.rs
2//! MongoDB sink (feature `mongodb`).
3
4use lc_core::observability::{MetricsSink, ObsError, ObsEvent};
5
6/// Inserts each event as one document into the given collection. Records keep the
7/// serde tag (`"kind"`), so a single collection can hold token-usage and agent
8/// metrics side by side.
9#[derive(Clone)]
10pub struct MongoSink {
11    collection: mongodb::Collection<serde_json::Value>,
12}
13
14impl MongoSink {
15    /// Builds a sink from a MongoDB connection string.
16    pub async fn new(uri: &str, database: &str, collection: &str) -> Result<Self, ObsError> {
17        let client = mongodb::Client::with_uri_str(uri)
18            .await
19            .map_err(|e| ObsError::Transport(format!("mongo: connect {uri}: {e}")))?;
20        Ok(Self::with_client(&client, database, collection))
21    }
22
23    /// Builds a sink from an existing client (shares its connection pool).
24    pub fn with_client(client: &mongodb::Client, database: &str, collection: &str) -> Self {
25        Self {
26            collection: client
27                .database(database)
28                .collection::<serde_json::Value>(collection),
29        }
30    }
31}
32
33#[async_trait::async_trait]
34impl MetricsSink for MongoSink {
35    async fn export(&self, event: &ObsEvent) -> Result<(), ObsError> {
36        let doc = serde_json::to_value(event).map_err(|e| ObsError::Encode(e.to_string()))?;
37        self.collection
38            .insert_one(doc, None)
39            .await
40            .map_err(|e| ObsError::Transport(format!("mongo: insert: {e}")))?;
41        Ok(())
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use lc_core::language_models::TokenUsage;
49
50    #[tokio::test]
51    #[ignore = "requires a running MongoDB (set MONGODB_URI)"]
52    async fn mongo_inserts_document() {
53        let uri =
54            std::env::var("MONGODB_URI").unwrap_or_else(|_| "mongodb://localhost:27017".into());
55        let sink = MongoSink::new(&uri, "lc_obs_test", "events").await.unwrap();
56        sink.export(&ObsEvent::TokenUsage(TokenUsage {
57            prompt_tokens: 1,
58            completion_tokens: 1,
59            total_tokens: 2,
60        }))
61        .await
62        .unwrap();
63    }
64}