Feat/bg summary (#925)

olzhik11 · web-flow · commit 7c5f971cbec3 · 2025-10-03T14:44:20.000+01:00
* feat: summary WIP * feat: remove eligibility check from FE, rename vars * feat: remove test * feat: remove cache * Revert "feat: remove cache" This reverts commit 42bc3b2.
diff --git a/app-server/src/db/mod.rs b/app-server/src/db/mod.rs
@@ -8,6 +8,7 @@ pub mod evaluators;
 pub mod events;
 pub mod prices;
 pub mod project_api_keys;
+pub mod project_settings;
 pub mod projects;
 pub mod provider_api_keys;
 pub mod spans;
diff --git a/app-server/src/db/project_settings.rs b/app-server/src/db/project_settings.rs
@@ -0,0 +1,46 @@
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use sqlx::{FromRow, PgPool};
+use uuid::Uuid;
+
+#[derive(Debug, Clone, FromRow, Serialize, Deserialize)]
+pub struct ProjectSetting {
+    pub id: Uuid,
+    pub name: Option<String>,
+    pub value: Option<String>,
+    pub project_id: Uuid,
+}
+
+/// Get a specific project setting by name
+pub async fn get_project_setting(
+    pool: &PgPool,
+    project_id: &Uuid,
+    setting_name: &str,
+) -> Result<Option<ProjectSetting>> {
+    let setting = sqlx::query_as::<_, ProjectSetting>(
+        "SELECT id, name, value, project_id
+         FROM project_settings
+         WHERE project_id = $1 AND name = $2
+         LIMIT 1",
+    )
+    .bind(project_id)
+    .bind(setting_name)
+    .fetch_optional(pool)
+    .await?;
+
+    Ok(setting)
+}
+
+/// Check if a project setting is enabled (value = 'true')
+pub async fn is_project_setting_enabled(
+    pool: &PgPool,
+    project_id: &Uuid,
+    setting_name: &str,
+) -> Result<bool> {
+    let setting = get_project_setting(pool, project_id, setting_name).await?;
+
+    Ok(setting
+        .and_then(|s| s.value)
+        .map(|v| v.trim().to_lowercase() == "true")
+        .unwrap_or(false))
+}
diff --git a/app-server/src/main.rs b/app-server/src/main.rs
@@ -588,6 +588,8 @@ fn main() -> anyhow::Result<()> {
 
                     for _ in 0..num_trace_summary_workers_per_thread {
                         tokio::spawn(process_trace_summaries(
+                            db_for_http.clone(),
+                            cache_for_http.clone(),
                             mq_for_http.clone(),
                         ));
                     }
diff --git a/app-server/src/traces/eligibility.rs b/app-server/src/traces/eligibility.rs
@@ -0,0 +1,81 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use uuid::Uuid;
+
+use crate::{
+    cache::{Cache, CacheTrait, keys::PROJECT_CACHE_KEY},
+    db::{DB, project_settings::is_project_setting_enabled, projects::get_project_and_workspace_billing_info},
+};
+
+#[derive(Debug, Clone)]
+pub struct TraceEligibilityResult {
+    pub is_eligible: bool,
+    pub reason: Option<String>,
+    pub tier_name: Option<String>,
+    pub has_trace_analysis: bool,
+}
+
+pub async fn check_trace_eligibility(
+    db: Arc<DB>,
+    cache: Arc<Cache>,
+    project_id: Uuid,
+) -> Result<TraceEligibilityResult> {
+    let cache_key = format!("{}:{}", PROJECT_CACHE_KEY, project_id);
+    let project_info = cache.get::<serde_json::Value>(&cache_key).await;
+
+    let tier_name = match project_info {
+        Ok(Some(info)) => {
+            info.get("tier_name")
+                .and_then(|v| v.as_str())
+                .map(|s| s.to_string())
+        }
+        _ => {
+            // Fallback: query database if not in cache
+            match get_project_and_workspace_billing_info(&db.pool, &project_id).await {
+                Ok(project) => Some(project.tier_name),
+                Err(_) => {
+                    return Ok(TraceEligibilityResult {
+                        is_eligible: false,
+                        reason: Some("project not found".to_string()),
+                        tier_name: None,
+                        has_trace_analysis: false,
+                    });
+                }
+            }
+        }
+    };
+
+    let is_paid_tier = tier_name
+        .as_ref()
+        .map(|name| name.trim().to_lowercase() != "free")
+        .unwrap_or(false);
+
+    if !is_paid_tier {
+        return Ok(TraceEligibilityResult {
+            is_eligible: false,
+            reason: Some("workspace is on free tier".to_string()),
+            tier_name,
+            has_trace_analysis: false,
+        });
+    }
+
+    let is_trace_analysis_enabled =
+        is_project_setting_enabled(&db.pool, &project_id, "enable_trace_analysis").await?;
+
+    if !is_trace_analysis_enabled {
+        return Ok(TraceEligibilityResult {
+            is_eligible: false,
+            reason: Some("trace analysis not enabled in project settings".to_string()),
+            tier_name,
+            has_trace_analysis: false,
+        });
+    }
+
+    Ok(TraceEligibilityResult {
+        is_eligible: true,
+        reason: None,
+        tier_name,
+        has_trace_analysis: true,
+    })
+}
diff --git a/app-server/src/traces/mod.rs b/app-server/src/traces/mod.rs
@@ -1,4 +1,5 @@
 pub mod consumer;
+pub mod eligibility;
 pub mod events;
 pub mod grpc_service;
 pub mod limits;
diff --git a/app-server/src/traces/summary.rs b/app-server/src/traces/summary.rs
@@ -1,17 +1,21 @@
-//! This module handles trace summary generation
-//! It reads trace completion messages from RabbitMQ and generates summaries via internal API
-
 use std::env;
 use std::sync::Arc;
 
 use backoff::ExponentialBackoffBuilder;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
-use super::{TRACE_SUMMARY_EXCHANGE, TRACE_SUMMARY_QUEUE, TRACE_SUMMARY_ROUTING_KEY};
-use crate::mq::{
-    MessageQueue, MessageQueueAcker, MessageQueueDeliveryTrait, MessageQueueReceiverTrait,
-    MessageQueueTrait,
+use super::{
+    TRACE_SUMMARY_EXCHANGE, TRACE_SUMMARY_QUEUE, TRACE_SUMMARY_ROUTING_KEY,
+    eligibility::check_trace_eligibility,
+};
+use crate::{
+    cache::Cache,
+    db::DB,
+    mq::{
+        MessageQueue, MessageQueueAcker, MessageQueueDeliveryTrait, MessageQueueReceiverTrait,
+        MessageQueueTrait,
+    },
 };
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -51,14 +55,14 @@ pub async fn push_to_trace_summary_queue(
 }
 
 /// Main worker function to process trace summary messages
-pub async fn process_trace_summaries(queue: Arc<MessageQueue>) {
+pub async fn process_trace_summaries(db: Arc<DB>, cache: Arc<Cache>, queue: Arc<MessageQueue>) {
     loop {
-        inner_process_trace_summaries(queue.clone()).await;
+        inner_process_trace_summaries(db.clone(), cache.clone(), queue.clone()).await;
         log::warn!("Trace summary listener exited. Rebinding queue connection...");
     }
 }
 
-async fn inner_process_trace_summaries(queue: Arc<MessageQueue>) {
+async fn inner_process_trace_summaries(db: Arc<DB>, cache: Arc<Cache>, queue: Arc<MessageQueue>) {
     // Add retry logic with exponential backoff for connection failures
     let get_receiver = || async {
         queue
@@ -119,7 +123,15 @@ async fn inner_process_trace_summaries(queue: Arc<MessageQueue>) {
             };
 
         // Process the trace summary generation
-        if let Err(e) = process_single_trace_summary(&client, trace_summary_message, acker).await {
+        if let Err(e) = process_single_trace_summary(
+            &client,
+            db.clone(),
+            cache.clone(),
+            trace_summary_message,
+            acker,
+        )
+        .await
+        {
             log::error!("Failed to process trace summary: {:?}", e);
         }
     }
@@ -129,43 +141,70 @@ async fn inner_process_trace_summaries(queue: Arc<MessageQueue>) {
 
 async fn process_single_trace_summary(
     client: &reqwest::Client,
+    db: Arc<DB>,
+    cache: Arc<Cache>,
     message: TraceSummaryMessage,
     acker: MessageQueueAcker,
 ) -> anyhow::Result<()> {
-    // Get the internal API base URL - this should be the internal service URL
-    let internal_api_base_url =
-        env::var("NEXT_BACKEND_URL").unwrap_or_else(|_| "http://localhost:3000".to_string());
+    let eligibility_result = check_trace_eligibility(db, cache, message.project_id).await?;
+
+    if !eligibility_result.is_eligible {
+        log::info!(
+            "Skipping trace summary generation: trace_id={}, project_id={}, reason={}",
+            message.trace_id,
+            message.project_id,
+            eligibility_result.reason.unwrap_or_default()
+        );
+        if let Err(e) = acker.ack().await {
+            log::error!("Failed to ack trace summary message: {:?}", e);
+        }
+        return Ok(());
+    }
+
+    let summarizer_service_url = env::var("TRACE_SUMMARIZER_URL")
+        .map_err(|_| anyhow::anyhow!("TRACE_SUMMARIZER_URL environment variable not set"))?;
 
-    let url = format!("{}/api/traces/summary", internal_api_base_url);
+    let auth_token = env::var("TRACE_SUMMARIZER_SECRET_KEY")
+        .map_err(|_| anyhow::anyhow!("TRACE_SUMMARIZER_SECRET_KEY environment variable not set"))?;
 
     let request_body = serde_json::json!({
         "projectId": message.project_id.to_string(),
         "traceId": message.trace_id.to_string(),
+        "maxRetries": 5
     });
 
-    let call_internal_api = || async {
+    let call_summarizer_service = || async {
         let response = client
-            .post(&url)
+            .post(&summarizer_service_url)
+            .header("Authorization", format!("Bearer {}", auth_token))
+            .header("Content-Type", "application/json")
             .json(&request_body)
             .send()
             .await
             .map_err(|e| {
-                log::warn!("Failed to call internal API for trace summary: {:?}", e);
+                log::warn!("Failed to call summarizer service for trace summary: {:?}", e);
                 backoff::Error::transient(anyhow::Error::from(e))
             })?;
 
         if response.status().is_success() {
-            Ok(response)
+            let response_text = response.text().await.unwrap_or_default();
+            log::debug!(
+                "Summarizer service response for trace_id={}, project_id={}: {}",
+                message.trace_id,
+                message.project_id,
+                response_text
+            );
+            Ok(())
         } else {
             let status = response.status();
             let response_text = response.text().await.unwrap_or_default();
             log::warn!(
-                "Internal API returned error status for trace summary: {}, Response: {}",
+                "Summarizer service returned error status for trace summary: {}, Response: {}",
                 status,
                 response_text
             );
             Err(backoff::Error::transient(anyhow::anyhow!(
-                "Internal API error: {}, Response: {}",
+                "Summarizer service error: {}, Response: {}",
                 status,
                 response_text
             )))
@@ -178,8 +217,13 @@ async fn process_single_trace_summary(
         .with_max_elapsed_time(Some(std::time::Duration::from_secs(60 * 5))) // 5 minutes max
         .build();
 
-    match backoff::future::retry(backoff, call_internal_api).await {
-        Ok(_response) => {
+    match backoff::future::retry(backoff, call_summarizer_service).await {
+        Ok(_) => {
+            log::info!(
+                "Successfully generated trace summary: trace_id={}, project_id={}",
+                message.trace_id,
+                message.project_id
+            );
             if let Err(e) = acker.ack().await {
                 log::error!("Failed to ack trace summary message: {:?}", e);
             }
diff --git a/frontend/app/api/traces/summary/route.ts b/frontend/app/api/traces/summary/route.ts
@@ -1,10 +1,9 @@
-import { observe } from '@lmnr-ai/lmnr';
-import { prettifyError } from 'zod/v4';
+import { observe } from "@lmnr-ai/lmnr";
+import { prettifyError } from "zod/v4";
 
-import { checkTraceEligibility } from '@/lib/actions/project/trace-eligibility';
-import { executeQuery } from '@/lib/actions/sql';
-import { generateTraceSummary } from '@/lib/actions/trace/agent';
-import { GenerateTraceSummaryRequestSchema } from '@/lib/actions/trace/agent/summary';
+import { executeQuery } from "@/lib/actions/sql";
+import { generateTraceSummary } from "@/lib/actions/trace/agent";
+import { GenerateTraceSummaryRequestSchema } from "@/lib/actions/trace/agent/summary";
 
 /**
  * Internal endpoint for trace summary generation.
@@ -16,7 +15,7 @@ export async function POST(req: Request) {
   const traceSummaryResult = GenerateTraceSummaryRequestSchema.safeParse(body);
 
   if (!traceSummaryResult.success) {
-    console.error('Validation error for trace summary request:', prettifyError(traceSummaryResult.error));
+    console.error("Validation error for trace summary request:", prettifyError(traceSummaryResult.error));
     return Response.json({ error: prettifyError(traceSummaryResult.error) }, { status: 400 });
   }
 
@@ -26,16 +25,6 @@ export async function POST(req: Request) {
   const { projectId, traceId } = traceSummaryResult.data;
 
   try {
-    // Check if project is eligible for trace summary generation
-    const eligibilityResult = await checkTraceEligibility({ projectId });
-
-    if (!eligibilityResult.isEligible) {
-      return Response.json({
-        success: true,
-        message: `Skipped - ${eligibilityResult.reason}`
-      });
-    }
-
     // check if the trace contains at least one LLM span
     const llmSpanCheckQuery = `
      SELECT COUNT(*) as llm_span_count
@@ -50,28 +39,32 @@ export async function POST(req: Request) {
       query: llmSpanCheckQuery,
       parameters: {
         traceId,
-      }
+      },
     });
 
     const hasLlmSpans = llmSpanResult.length > 0 && llmSpanResult[0].llm_span_count > 0;
 
     if (!hasLlmSpans) {
       return Response.json({
         success: true,
-        message: "Skipped - trace contains no LLM spans"
+        message: "Skipped - trace contains no LLM spans",
       });
     }
 
     // Generate the trace summary since all requirements are met
     // Disable retries for this call since we want to fail fast if the summary generation fails
-    await observe({ name: "generateTraceSummaryIfNeeded" }, async () => await generateTraceSummary({
-      ...traceSummaryResult.data,
-      maxRetries: 0,
-    }));
+    await observe(
+      { name: "generateTraceSummaryIfNeeded" },
+      async () =>
+        await generateTraceSummary({
+          ...traceSummaryResult.data,
+          maxRetries: 0,
+        })
+    );
 
     return Response.json({ success: true });
   } catch (error) {
-    console.error('Failed to generate trace summary:', error);
+    console.error("Failed to generate trace summary:", error);
     return Response.json(
       { error: error instanceof Error ? error.message : "Failed to generate trace summary." },
       { status: 500 }
diff --git a/frontend/lib/actions/project/trace-eligibility.ts b/frontend/lib/actions/project/trace-eligibility.ts

Original file line number	Diff line number	Diff line change
`@@ -588,6 +588,8 @@ fn main() -> anyhow::Result<()> {`
`588`	`588`
`589`	`589`	`for _ in 0..num_trace_summary_workers_per_thread {`
`590`	`590`	`tokio::spawn(process_trace_summaries(`
	`591`	`+ db_for_http.clone(),`
	`592`	`+ cache_for_http.clone(),`
`591`	`593`	`mq_for_http.clone(),`
`592`	`594`	`));`
`593`	`595`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`pub mod consumer;`
	`2`	`+pub mod eligibility;`
`2`	`3`	`pub mod events;`
`3`	`4`	`pub mod grpc_service;`
`4`	`5`	`pub mod limits;`