Merge pull request #982 from lmnr-ai/dev

Fix CH migrations, fix console logs, optimize tag writing

Author: dinmukhamedm

app-server/src/ch/tags.rs

@@ -2,9 +2,10 @@ use anyhow::Result;
 use chrono::Utc;
 use clickhouse::Row;
 use serde::{Deserialize, Serialize};
+use tracing::instrument;
 use uuid::Uuid;
 
-use crate::db::tags::TagSource;
+use crate::db::tags::{SpanTag, TagSource};
 
 use super::utils::chrono_to_nanoseconds;
 
@@ -79,3 +80,44 @@ pub async fn insert_tag(
         }
     }
 }
+
+#[instrument(skip(client, tags))]
+pub async fn insert_tags_batch(client: clickhouse::Client, tags: &[SpanTag]) -> Result<()> {
+    if tags.is_empty() {
+        return Ok(());
+    }
+
+    let ch_insert = client.insert("tags");
+    match ch_insert {
+        Ok(mut ch_insert) => {
+            for span_tag in tags {
+                let id = Uuid::new_v4();
+                let tag = CHTag::new(
+                    span_tag.project_id,
+                    id,
+                    span_tag.name.clone(),
+                    span_tag.source.clone(),
+                    span_tag.span_id,
+                );
+                ch_insert.write(&tag).await?;
+            }
+
+            let ch_insert_end_res = ch_insert.end().await;
+            match ch_insert_end_res {
+                Ok(_) => Ok(()),
+                Err(e) => {
+                    return Err(anyhow::anyhow!(
+                        "Clickhouse batch tag insertion failed: {:?}",
+                        e
+                    ));
+                }
+            }
+        }
+        Err(e) => {
+            return Err(anyhow::anyhow!(
+                "Failed to insert tags batch into Clickhouse: {:?}",
+                e
+            ));
+        }
+    }
+}

app-server/src/db/prices.rs

@@ -11,7 +11,11 @@ pub struct DBPriceEntry {
     pub additional_prices: Value,
 }
 
-pub async fn get_price(pool: &PgPool, provider: &str, model: &str) -> anyhow::Result<DBPriceEntry> {
+pub async fn get_price(
+    pool: &PgPool,
+    provider: &str,
+    model: &str,
+) -> anyhow::Result<Option<DBPriceEntry>> {
     let price = sqlx::query_as::<_, DBPriceEntry>(
         "SELECT
             provider,
@@ -32,7 +36,7 @@ pub async fn get_price(pool: &PgPool, provider: &str, model: &str) -> anyhow::Re
     )
     .bind(provider)
     .bind(model)
-    .fetch_one(pool)
+    .fetch_optional(pool)
     .await?;
 
     Ok(price)

app-server/src/db/tags.rs

@@ -1,9 +1,30 @@
 use serde::{Deserialize, Serialize};
+use uuid::Uuid;
 
-#[derive(sqlx::Type, Serialize, Deserialize, Clone, PartialEq)]
+#[derive(Debug, sqlx::Type, Serialize, Deserialize, Clone, PartialEq)]
 #[sqlx(type_name = "tag_source")]
 pub enum TagSource {
     MANUAL,
     AUTO,
     CODE,
 }
+
+/// Structured representation of a span tag for batch operations
+#[derive(Debug, Clone)]
+pub struct SpanTag {
+    pub project_id: Uuid,
+    pub name: String,
+    pub source: TagSource,
+    pub span_id: Uuid,
+}
+
+impl SpanTag {
+    pub fn new(project_id: Uuid, name: String, source: TagSource, span_id: Uuid) -> Self {
+        Self {
+            project_id,
+            name,
+            source,
+            span_id,
+        }
+    }
+}

app-server/src/instrumentation/mod.rs

@@ -9,30 +9,35 @@ use opentelemetry::trace::TracerProvider;
 use opentelemetry_sdk::trace::SdkTracerProvider;
 use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
 
-pub fn setup_tracing() {
+pub fn setup_tracing(enable_otel: bool) {
     let env_filter = if std::env::var("RUST_LOG").is_ok_and(|s| !s.is_empty()) {
         EnvFilter::from_default_env()
     } else {
         EnvFilter::new("info")
     };
 
-    let mut provider_builder = SdkTracerProvider::builder();
+    let registry = tracing_subscriber::registry()
+        .with(env_filter)
+        .with(tracing_subscriber::fmt::layer());
 
-    if std::env::var("SENTRY_DSN").is_ok() {
-        provider_builder = provider_builder
-            .with_span_processor(sentry::integrations::opentelemetry::SentrySpanProcessor::new());
-    }
+    if enable_otel {
+        let mut provider_builder = SdkTracerProvider::builder();
 
-    let tracer_provider = provider_builder.build();
-    let tracer = tracer_provider.tracer("app-server");
+        if std::env::var("SENTRY_DSN").is_ok() {
+            provider_builder = provider_builder.with_span_processor(
+                sentry::integrations::opentelemetry::SentrySpanProcessor::new(),
+            );
+        }
 
-    global::set_tracer_provider(tracer_provider);
+        let tracer_provider = provider_builder.build();
+        let tracer = tracer_provider.tracer("app-server");
 
-    let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
+        global::set_tracer_provider(tracer_provider);
 
-    tracing_subscriber::registry()
-        .with(env_filter)
-        .with(otel_layer)
-        .with(tracing_subscriber::fmt::layer())
-        .init();
+        let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
+
+        registry.with(otel_layer).init();
+    } else {
+        registry.init();
+    }
 }

app-server/src/language_model/costs/mod.rs

@@ -13,6 +13,8 @@ use utils::calculate_cost;
 
 mod utils;
 
+const LLM_PRICES_CACHE_TTL_SECONDS: u64 = 60 * 60 * 24; // 24 hours
+
 #[derive(Clone, Deserialize, Serialize)]
 pub struct LLMPriceEntry {
     _provider: String,
@@ -44,15 +46,30 @@ pub async fn estimate_output_cost(
     num_tokens: i64,
 ) -> Option<f64> {
     let cache_key = format!("{LLM_PRICES_CACHE_KEY}:{provider}:{model}");
-    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await.ok()?;
+    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await;
 
     let price_per_million_tokens = match cache_res {
-        Some(price) => price.output_price_per_million,
-        None => {
-            let price = get_price(&db.pool, provider, model).await.ok()?;
+        Ok(Some(price)) => price.output_price_per_million,
+        Ok(None) | Err(_) => {
+            let price = get_price(&db.pool, provider, model)
+                .await
+                .map_err(|e| {
+                    log::error!(
+                        "Error getting price from DB for provider: {}, model: {}: {:?}",
+                        provider,
+                        model,
+                        e
+                    );
+                    e
+                })
+                .unwrap_or_default()?;
             let price = LLMPriceEntry::from(price);
             let _ = cache
-                .insert::<LLMPriceEntry>(&cache_key, price.clone())
+                .insert_with_ttl::<LLMPriceEntry>(
+                    &cache_key,
+                    price.clone(),
+                    LLM_PRICES_CACHE_TTL_SECONDS,
+                )
                 .await;
             price.output_price_per_million
         }
@@ -68,16 +85,30 @@ pub async fn estimate_input_cost(
     input_tokens: InputTokens,
 ) -> Option<f64> {
     let cache_key = format!("{LLM_PRICES_CACHE_KEY}:{provider}:{model}");
-    // let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await.ok()?;
-    let cache_res = None;
+    let cache_res = cache.get::<LLMPriceEntry>(&cache_key).await;
 
     let price = match cache_res {
-        Some(price) => price,
-        None => {
-            let price = get_price(&db.pool, provider, model).await.ok()?;
+        Ok(Some(price)) => price,
+        Ok(None) | Err(_) => {
+            let price = get_price(&db.pool, provider, model)
+                .await
+                .map_err(|e| {
+                    log::error!(
+                        "Error getting price from DB for provider: {}, model: {}: {:?}",
+                        provider,
+                        model,
+                        e
+                    );
+                    e
+                })
+                .unwrap_or_default()?;
             let price = LLMPriceEntry::from(price);
             let _ = cache
-                .insert::<LLMPriceEntry>(&cache_key, price.clone())
+                .insert_with_ttl::<LLMPriceEntry>(
+                    &cache_key,
+                    price.clone(),
+                    LLM_PRICES_CACHE_TTL_SECONDS,
+                )
                 .await;
             price
         }

app-server/src/main.rs

@@ -115,9 +115,7 @@ fn main() -> anyhow::Result<()> {
         drop(_sentry_guard);
     }
 
-    if is_feature_enabled(Feature::Tracing) {
-        instrumentation::setup_tracing();
-    }
+    instrumentation::setup_tracing(is_feature_enabled(Feature::Tracing));
 
     let http_payload_limit: usize = env::var("HTTP_PAYLOAD_LIMIT")
         .unwrap_or(String::from("5242880")) // default to 5MB

app-server/src/traces/consumer.rs

@@ -24,7 +24,13 @@ use crate::{
         spans::CHSpan,
         traces::{CHTrace, TraceAggregation, upsert_traces_batch},
     },
-    db::{DB, events::Event, spans::Span, trace::upsert_trace_statistics_batch},
+    db::{
+        DB,
+        events::Event,
+        spans::Span,
+        tags::{SpanTag, TagSource},
+        trace::upsert_trace_statistics_batch,
+    },
     evaluators::{get_evaluators_by_path, push_to_evaluators_queue},
     features::{Feature, is_feature_enabled},
     mq::{
@@ -38,7 +44,7 @@ use crate::{
         events::record_events,
         limits::update_workspace_limit_exceeded_by_project_id,
         provider::convert_span_to_provider_format,
-        utils::{get_llm_usage_for_span, prepare_span_for_recording, record_tags},
+        utils::{get_llm_usage_for_span, prepare_span_for_recording},
     },
 };
 
@@ -427,19 +433,22 @@ async fn process_batch(
         }
     }
 
-    for span in &stripped_spans {
-        if let Err(e) = record_tags(
-            clickhouse.clone(),
-            &span.tags,
-            &span.span_id,
-            &span.project_id,
-        )
-        .await
-        {
+    // Collect all tags from all spans for batch insertion
+    let tags_batch: Vec<SpanTag> = stripped_spans
+        .iter()
+        .flat_map(|span| {
+            span.tags.iter().map(move |tag| {
+                SpanTag::new(span.project_id, tag.clone(), TagSource::CODE, span.span_id)
+            })
+        })
+        .collect();
+
+    // Record all tags in a single batch
+    if !tags_batch.is_empty() {
+        if let Err(e) = crate::ch::tags::insert_tags_batch(clickhouse.clone(), &tags_batch).await {
             log::error!(
-                "Failed to record tags to DB. span_id [{}], project_id [{}]: {:?}",
-                span.span_id,
-                span.project_id,
+                "Failed to record tags to DB for batch of {} tags: {:?}",
+                tags_batch.len(),
                 e
             );
         }

app-server/src/traces/utils.rs

@@ -13,7 +13,7 @@ use crate::{
 
 use crate::{
     cache::Cache,
-    db::{DB, spans::Span, tags::TagSource},
+    db::{DB, spans::Span},
     language_model::costs::estimate_cost_by_provider_name,
 };
 
@@ -97,31 +97,6 @@ pub async fn get_llm_usage_for_span(
     }
 }
 
-#[instrument(skip(clickhouse, tags, span_id, project_id))]
-pub async fn record_tags(
-    clickhouse: clickhouse::Client,
-    tags: &[String],
-    span_id: &Uuid,
-    project_id: &Uuid,
-) -> anyhow::Result<()> {
-    if tags.is_empty() {
-        return Ok(());
-    }
-
-    for tag_name in tags {
-        crate::ch::tags::insert_tag(
-            clickhouse.clone(),
-            *project_id,
-            tag_name.clone(),
-            TagSource::CODE,
-            *span_id,
-        )
-        .await?;
-    }
-
-    Ok(())
-}
-
 pub fn skip_span_name(name: &str) -> bool {
     SKIP_SPAN_NAME_REGEX.is_match(name)
 }