Merge pull request #906 from lmnr-ai/dev

Realtime updates, move traces and spans to query engine, frontend fixes

Author: dinmukhamedm

.env

@@ -8,17 +8,21 @@ RABBITMQ_DEFAULT_USER=admin
 RABBITMQ_DEFAULT_PASS=adminpasswd
 
 CLICKHOUSE_USER=ch_user
+CLICKHOUSE_RO_USER=ch_user
 CLICKHOUSE_PASSWORD=ch_passwd
+CLICKHOUSE_RO_PASSWORD=ch_passwd
 POSTGRES_PORT=5433
 # must be exactly 32 bytes (64 hex characters)
 AEAD_SECRET_KEY=0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef
+# could be any value
+NEXTAUTH_SECRET=0123456789abcdef0123456789abcdef
 
 
 # Uncomment and set these to override the default exposed ports.
-# APP_SERVER_HOST_PORT=8080
-# APP_SERVER_GRPC_HOST_PORT=8081
+# APP_SERVER_HOST_PORT=8000
+# APP_SERVER_GRPC_HOST_PORT=8001
 # FRONTEND_HOST_PORT=5667
-# AGENT_MANAGER_HOST_PORT=8901
+# QUERY_ENGINE_HOST_PORT=8903
 
 # Uncomment and set this to enable OpenAI support. Currently used in Next.js
 # to generate chat names.

.github/workflows/build-push.yml

@@ -20,6 +20,9 @@ jobs:
           - dockerfile: ./frontend/Dockerfile
             context: ./frontend
             image: ghcr.io/lmnr-ai/frontend
+          - dockerfile: ./query-engine/Dockerfile
+            context: ./query-engine
+            image: ghcr.io/lmnr-ai/query-engine
     permissions:
       contents: read
       packages: write

.github/workflows/query-engine-tests-check.yml

@@ -0,0 +1,32 @@
+name: Check Query Engine Tests
+
+on:
+  pull_request:
+    types:
+      - synchronize
+      - opened
+      - reopened
+    paths:
+      - 'query-engine/**' 
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.13
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install the project
+        working-directory: ./query-engine
+        run: uv sync --all-extras --dev
+      - name: Run tests
+        working-directory: ./query-engine
+        run: uv run pytest

app-server/.env.example

@@ -9,6 +9,8 @@ RABBITMQ_URL=amqp://admin:adminpasswd@localhost:5672/%2f
 CLICKHOUSE_URL=http://localhost:8123
 CLICKHOUSE_USER=ch_user
 CLICKHOUSE_PASSWORD=ch_passwd
+CLICKHOUSE_RO_USER=ch_user
+CLICKHOUSE_RO_PASSWORD=ch_passwd
 
 SHARED_SECRET_TOKEN=some_secret
 # must be exactly 32 bytes (64 hex characters)
@@ -18,5 +20,4 @@ ENVIRONMENT=FULL
 # Optional, if you want to use Redis for caching. Create a Redis/Valkey instance and set the URL here.
 # REDIS_URL=redis://localhost:6379
 
-
-QUERY_ENGINE_URL=
+QUERY_ENGINE_URL=http://localhost:8903

app-server/src/api/v1/evaluators.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix_web::{
     HttpResponse, post,
     web::{Data, Json},
@@ -12,9 +14,10 @@ use crate::{
         DB,
         evaluators::{EvaluatorScoreSource, insert_evaluator_score},
         project_api_keys::ProjectApiKey,
-        spans::{get_root_span_id, is_span_in_project},
     },
+    query_engine::QueryEngine,
     routes::types::ResponseResult,
+    sql::{self, ClickhouseReadonlyClient},
 };
 
 #[derive(Deserialize)]
@@ -55,9 +58,14 @@ pub async fn create_evaluator_score(
     req: Json<CreateEvaluatorScoreRequest>,
     db: Data<DB>,
     clickhouse: Data<clickhouse::Client>,
+    clickhouse_ro: Data<Option<Arc<ClickhouseReadonlyClient>>>,
+    query_engine: Data<Arc<QueryEngine>>,
     project_api_key: ProjectApiKey,
 ) -> ResponseResult {
     let req = req.into_inner();
+    let clickhouse_ro = clickhouse_ro.as_ref().clone().unwrap();
+    let query_engine = query_engine.as_ref().clone();
+    let clickhouse = clickhouse.as_ref().clone();
 
     // Extract common fields from both variants
     let (name, metadata, score, source) = match &req {
@@ -85,11 +93,21 @@ pub async fn create_evaluator_score(
 
     let span_id = match &req {
         CreateEvaluatorScoreRequest::WithTraceId(req) => {
-            get_root_span_id(&db.pool, &req.trace_id, &project_api_key.project_id).await?
+            sql::queries::get_top_span_id(
+                clickhouse_ro,
+                query_engine,
+                req.trace_id,
+                project_api_key.project_id,
+            )
+            .await?
         }
         CreateEvaluatorScoreRequest::WithSpanId(req) => {
-            let exists =
-                is_span_in_project(&db.pool, &req.span_id, &project_api_key.project_id).await?;
+            let exists = crate::ch::spans::is_span_in_project(
+                clickhouse.clone(),
+                req.span_id,
+                project_api_key.project_id,
+            )
+            .await?;
             if !exists {
                 return Ok(HttpResponse::NotFound().body("No matching spans found"));
             }
@@ -118,7 +136,7 @@ pub async fn create_evaluator_score(
     .await?;
 
     let _ = insert_evaluator_score_ch(
-        clickhouse.into_inner().as_ref().clone(),
+        clickhouse.clone(),
         score_id,
         project_id,
         name,

app-server/src/api/v1/tag.rs

@@ -1,8 +1,11 @@
+use std::sync::Arc;
+
 use crate::{
-    ch::spans::append_tags_to_span,
-    db::{self, DB, project_api_keys::ProjectApiKey, tags::TagSource},
+    ch::{spans::append_tags_to_span, tags::insert_tag},
+    db::{project_api_keys::ProjectApiKey, tags::TagSource},
+    query_engine::QueryEngine,
     routes::types::ResponseResult,
-    tags::insert_or_update_tag,
+    sql::{self, ClickhouseReadonlyClient},
 };
 use actix_web::{
     HttpResponse, post,
@@ -37,8 +40,9 @@ pub enum TagRequest {
 #[post("tag")]
 pub async fn tag_trace(
     req: Json<TagRequest>,
-    db: web::Data<DB>,
     clickhouse: web::Data<clickhouse::Client>,
+    clickhouse_ro: web::Data<Option<Arc<ClickhouseReadonlyClient>>>,
+    query_engine: web::Data<Arc<QueryEngine>>,
     project_api_key: ProjectApiKey,
 ) -> ResponseResult {
     let req = req.into_inner();
@@ -49,15 +53,27 @@ pub async fn tag_trace(
     if names.is_empty() {
         return Ok(HttpResponse::BadRequest().body("No names provided"));
     }
+    let clickhouse_ro = clickhouse_ro.as_ref().clone().unwrap();
+    let query_engine = query_engine.as_ref().clone();
+    let clickhouse = clickhouse.as_ref().clone();
+
     let span_id = match &req {
         TagRequest::WithTraceId(req) => {
-            db::spans::get_root_span_id(&db.pool, &req.trace_id, &project_api_key.project_id)
-                .await?
+            sql::queries::get_top_span_id(
+                clickhouse_ro,
+                query_engine,
+                req.trace_id,
+                project_api_key.project_id,
+            )
+            .await?
         }
         TagRequest::WithSpanId(req) => {
-            let exists =
-                db::spans::is_span_in_project(&db.pool, &req.span_id, &project_api_key.project_id)
-                    .await?;
+            let exists = crate::ch::spans::is_span_in_project(
+                clickhouse.clone(),
+                req.span_id,
+                project_api_key.project_id,
+            )
+            .await?;
             if !exists {
                 return Ok(HttpResponse::NotFound().body("No matching spans found"));
             }
@@ -69,26 +85,20 @@ pub async fn tag_trace(
         return Ok(HttpResponse::NotFound().body("No matching spans found"));
     };
 
-    let clickhouse = clickhouse.as_ref().clone();
-
     let futures = names
         .iter()
         .map(|name| {
-            insert_or_update_tag(
-                &db.pool,
+            insert_tag(
                 clickhouse.clone(),
                 project_api_key.project_id,
-                Uuid::new_v4(),
-                span_id,
-                None,
-                None,
                 name.clone(),
                 TagSource::CODE,
+                span_id,
             )
         })
         .collect::<Vec<_>>();
 
-    let tags = futures_util::future::try_join_all(futures).await?;
+    let tag_ids = futures_util::future::try_join_all(futures).await?;
 
     append_tags_to_span(
         clickhouse.clone(),
@@ -98,14 +108,12 @@ pub async fn tag_trace(
     )
     .await?;
 
-    let response = tags
+    let response = tag_ids
         .iter()
-        .map(|tag| {
+        .map(|id| {
             serde_json::json!({
-                "id": tag.id,
-                "spanId": tag.span_id,
-                "createdAt": tag.created_at,
-                "updatedAt": tag.updated_at,
+                "id": id,
+                "spanId": span_id,
             })
         })
         .collect::<Vec<_>>();

app-server/src/ch/spans.rs

@@ -99,17 +99,23 @@ impl CHSpan {
         let user_id = span.attributes.user_id();
         let path = span.attributes.flat_path();
 
-        let span_input_string = span
-            .input
-            .as_ref()
-            .map(|input| sanitize_string(&input.to_string()))
-            .unwrap_or(String::new());
-
-        let span_output_string = span
-            .output
-            .as_ref()
-            .map(|output| sanitize_string(&output.to_string()))
-            .unwrap_or(String::new());
+        let span_input_string = if let Some(input_url) = &span.input_url {
+            format!("<lmnr_payload_url>{}</lmnr_payload_url>", input_url)
+        } else {
+            span.input
+                .as_ref()
+                .map(|input| sanitize_string(&input.to_string()))
+                .unwrap_or(String::new())
+        };
+
+        let span_output_string = if let Some(output_url) = &span.output_url {
+            format!("<lmnr_payload_url>{}</lmnr_payload_url>", output_url)
+        } else {
+            span.output
+                .as_ref()
+                .map(|output| sanitize_string(&output.to_string()))
+                .unwrap_or(String::new())
+        };
 
         let trace_metadata = span.attributes.metadata().map_or(String::new(), |m| {
             serde_json::to_string(&m).unwrap_or_default()
@@ -135,18 +141,15 @@ impl CHSpan {
                 .unwrap_or(String::from("")),
             request_model: usage.request_model.clone().unwrap_or(String::from("")),
             response_model: usage.response_model.clone().unwrap_or(String::from("")),
-            session_id: session_id.unwrap_or(String::from("<null>")),
+            session_id: session_id.unwrap_or(String::from("")),
             project_id: project_id,
             trace_id: span.trace_id,
-            provider: usage
-                .provider_name
-                .clone()
-                .unwrap_or(String::from("<null>")),
-            user_id: user_id.unwrap_or(String::from("<null>")),
-            path: path.unwrap_or(String::from("<null>")),
+            provider: usage.provider_name.clone().unwrap_or(String::from("")),
+            user_id: user_id.unwrap_or(String::from("")),
+            path: path.unwrap_or(String::from("")),
             input: span_input_string,
             output: span_output_string,
-            status: span.status.clone().unwrap_or(String::from("<null>")),
+            status: span.status.clone().unwrap_or(String::from("")),
             size_bytes: size_bytes as u64,
             attributes: span.attributes.to_string(),
             trace_metadata,
@@ -215,3 +218,18 @@ pub async fn append_tags_to_span(
 
     Ok(())
 }
+
+pub async fn is_span_in_project(
+    clickhouse: clickhouse::Client,
+    span_id: Uuid,
+    project_id: Uuid,
+) -> Result<bool> {
+    let result = clickhouse
+        .query("SELECT count(*) FROM spans WHERE span_id = ? AND project_id = ?")
+        .bind(span_id)
+        .bind(project_id)
+        .fetch_one::<u64>()
+        .await?;
+
+    Ok(result > 0)
+}

app-server/src/ch/tags.rs

@@ -36,17 +36,11 @@ pub struct CHTag {
 }
 
 impl CHTag {
-    pub fn new(
-        project_id: Uuid,
-        class_id: Uuid,
-        id: Uuid,
-        name: String,
-        source: TagSource,
-        span_id: Uuid,
-    ) -> Self {
+    pub fn new(project_id: Uuid, id: Uuid, name: String, source: TagSource, span_id: Uuid) -> Self {
         Self {
             project_id,
-            class_id,
+            // TODO: Remove this once we drop the class_id column
+            class_id: Uuid::nil(),
             created_at: chrono_to_nanoseconds(Utc::now()),
             id,
             name,
@@ -59,20 +53,19 @@ impl CHTag {
 pub async fn insert_tag(
     client: clickhouse::Client,
     project_id: Uuid,
-    class_id: Uuid,
-    id: Uuid,
     name: String,
     source: TagSource,
     span_id: Uuid,
-) -> Result<()> {
-    let tag = CHTag::new(project_id, class_id, id, name, source, span_id);
+) -> Result<Uuid> {
+    let id = Uuid::new_v4();
+    let tag = CHTag::new(project_id, id, name, source, span_id);
     let ch_insert = client.insert("tags");
     match ch_insert {
         Ok(mut ch_insert) => {
             ch_insert.write(&tag).await?;
             let ch_insert_end_res = ch_insert.end().await;
             match ch_insert_end_res {
-                Ok(_) => Ok(()),
+                Ok(_) => Ok(id),
                 Err(e) => {
                     return Err(anyhow::anyhow!("Clickhouse tag insertion failed: {:?}", e));
                 }