lc_observability/
mongo.rs1use lc_core::observability::{MetricsSink, ObsError, ObsEvent};
5
6#[derive(Clone)]
10pub struct MongoSink {
11 collection: mongodb::Collection<serde_json::Value>,
12}
13
14impl MongoSink {
15 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 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}