Merge pull request #782 from lmnr-ai/dev

optimized browser events format, fix to clickhouse span payloads, migrate some tables to clickhouse, AI SDK v5 support

Author: dinmukhamedm

app-server/.env.example

@@ -20,3 +20,5 @@ 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=

app-server/build.rs

@@ -8,6 +8,15 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .out_dir("./src/agent_manager/")
         .compile_protos(&[proto_file], &["proto"])?;
 
+    let query_engine_proto_file = "./proto/query_engine.proto";
+
+    tonic_build::configure()
+        .protoc_arg("--experimental_allow_proto3_optional") // for older systems
+        .build_client(true)
+        .build_server(false)
+        .out_dir("./src/query_engine/")
+        .compile_protos(&[query_engine_proto_file], &["proto"])?;
+
     tonic_build::configure()
         .protoc_arg("--experimental_allow_proto3_optional") // for older systems
         .build_client(false)

app-server/proto/query_engine.proto

@@ -0,0 +1,28 @@
+syntax = "proto3";
+
+package query_engine;
+
+message QueryRequest {
+    string query = 1;
+    string project_id = 2;
+}
+
+message QueryResponse {
+    oneof result {
+        SuccessResponse success = 1;
+        ErrorResponse error = 2;
+    }
+}
+
+message SuccessResponse {
+    string query = 1;
+    bool success = 2;
+}
+
+message ErrorResponse {
+    string error = 1;
+}
+
+service QueryEngineService {
+    rpc ValidateQuery(QueryRequest) returns (QueryResponse);
+}

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

@@ -1,19 +1,19 @@
 use std::sync::Arc;
 
-use actix_web::{post, web, HttpResponse};
+use actix_web::{HttpResponse, post, web};
 use futures_util::StreamExt;
 use serde::Deserialize;
 use uuid::Uuid;
 
+use crate::agent_manager::RunAgentParams;
 use crate::agent_manager::channel::AgentManagerWorkers;
 use crate::agent_manager::types::RunAgentResponseStreamChunk;
-use crate::agent_manager::worker::{run_agent_worker, RunAgentWorkerOptions};
-use crate::agent_manager::RunAgentParams;
-use crate::agent_manager::{types::ModelProvider, AgentManager, AgentManagerTrait};
+use crate::agent_manager::worker::{RunAgentWorkerOptions, run_agent_worker};
+use crate::agent_manager::{AgentManager, AgentManagerTrait, types::ModelProvider};
 use crate::cache::Cache;
 use crate::db::project_api_keys::ProjectApiKey;
 use crate::db::{self, DB};
-use crate::features::{is_feature_enabled, Feature};
+use crate::features::{Feature, is_feature_enabled};
 use crate::routes::types::ResponseResult;
 use crate::traces::limits::get_workspace_limit_exceeded_by_project_id;
 
@@ -62,6 +62,7 @@ pub async fn run_agent_manager(
     agent_manager: web::Data<Arc<AgentManager>>,
     worker_states: web::Data<Arc<AgentManagerWorkers>>,
     db: web::Data<DB>,
+    clickhouse: web::Data<clickhouse::Client>,
     project_api_key: ProjectApiKey,
     cache: web::Data<Cache>,
     request: web::Json<RunAgentRequest>,
@@ -76,6 +77,7 @@ pub async fn run_agent_manager(
     if is_feature_enabled(Feature::UsageLimit) {
         match get_workspace_limit_exceeded_by_project_id(
             db.clone(),
+            clickhouse.into_inner().as_ref().clone(),
             cache.clone(),
             project_api_key.project_id,
         )
@@ -128,7 +130,6 @@ pub async fn run_agent_manager(
             return_agent_state: request.return_agent_state,
             return_storage_state: request.return_storage_state,
         };
-        let pool = db.pool.clone();
         let worker_states_clone = worker_states.clone();
         let handle = tokio::spawn(async move {
             run_agent_worker(
@@ -149,13 +150,6 @@ pub async fn run_agent_manager(
             while let Some(message) = receiver.recv().await {
                 match message {
                     Ok(agent_chunk) => {
-                        if let Err(e) =
-                            db::stats::add_agent_steps_to_project_usage_stats(&pool, &project_api_key.project_id, 1)
-                                .await
-                        {
-                            log::error!("Error adding agent steps to project usage stats: {}", e);
-                        }
-
                         match agent_chunk {
                             RunAgentResponseStreamChunk::FinalOutput(_) => {
                                 yield anyhow::Ok(agent_chunk.into());
@@ -222,15 +216,6 @@ pub async fn run_agent_manager(
         match fut.await {
             Ok(response) => {
                 let response = response?;
-                if let Err(e) = db::stats::add_agent_steps_to_project_usage_stats(
-                    &db.pool,
-                    &project_api_key.project_id,
-                    response.step_count.unwrap_or(0) as i64,
-                )
-                .await
-                {
-                    log::error!("Error adding agent steps to project usage stats: {}", e);
-                }
                 Ok(HttpResponse::Ok().json(response))
             }
             Err(e) if e.is_cancelled() => Ok(HttpResponse::NoContent().finish()),

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

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
-use actix_web::{HttpResponse, options, post, web};
-use serde::{Deserialize, Serialize};
+use actix_web::{HttpResponse, post, web};
+use serde::{Deserialize, Deserializer, Serialize};
 use uuid::Uuid;
 
 use crate::{
@@ -15,18 +15,42 @@ pub const BROWSER_SESSIONS_QUEUE: &str = "browser_sessions_queue";
 pub const BROWSER_SESSIONS_EXCHANGE: &str = "browser_sessions_exchange";
 pub const BROWSER_SESSIONS_ROUTING_KEY: &str = "browser_sessions_routing_key";
 
+// Custom deserializer for the data field to support both Vec<u8> and base64 string
+fn deserialize_data<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    use serde::de::Error;
+
+    #[derive(Deserialize)]
+    #[serde(untagged)]
+    enum DataFormat {
+        Bytes(Vec<u8>),
+        Base64String(String),
+    }
+
+    match DataFormat::deserialize(deserializer)? {
+        DataFormat::Bytes(bytes) => Ok(bytes),
+        DataFormat::Base64String(s) => {
+            base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &s)
+                .map_err(|e| Error::custom(format!("Invalid base64: {}", e)))
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct RRWebEvent {
     #[serde(rename = "type")]
     pub event_type: u8,
-    pub timestamp: i64,
+    pub timestamp: f64, // milliseconds
+    #[serde(deserialize_with = "deserialize_data")]
     pub data: Vec<u8>,
 }
 
 impl RRWebEvent {
     pub fn estimate_size_bytes(&self) -> usize {
         // 1 byte for event_type, 8 bytes for timestamp
-        return 9 + self.data.len();
+        9 + self.data.len()
     }
 }
 
@@ -42,19 +66,6 @@ pub struct EventBatch {
     pub sdk_version: Option<String>,
 }
 
-#[options("events")]
-async fn options_handler() -> ResponseResult {
-    Ok(HttpResponse::Ok()
-        .insert_header(("Access-Control-Allow-Origin", "*"))
-        .insert_header(("Access-Control-Allow-Methods", "POST, OPTIONS"))
-        .insert_header((
-            "Access-Control-Allow-Headers",
-            "Authorization, Content-Type, Content-Encoding, Accept",
-        ))
-        .insert_header(("Access-Control-Max-Age", "86400"))
-        .finish())
-}
-
 #[post("events")]
 async fn create_session_event(
     batch: web::Json<EventBatch>,

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

@@ -1,5 +1,6 @@
-use actix_web::{HttpResponse, get, web};
+use actix_web::{HttpResponse, get, post, web};
 use serde::Deserialize;
+use uuid::Uuid;
 
 use crate::{
     ch::datapoints as ch_datapoints,
@@ -58,3 +59,73 @@ async fn get_datapoints(
 
     Ok(HttpResponse::Ok().json(response))
 }
+
+#[derive(Deserialize)]
+pub struct CreateDatapointsRequest {
+    pub dataset_name: String,
+    pub datapoints: Vec<CreateDatapointRequest>,
+}
+
+#[derive(Deserialize)]
+pub struct CreateDatapointRequest {
+    pub data: serde_json::Value,
+    pub target: Option<serde_json::Value>,
+    #[serde(default)]
+    pub metadata: std::collections::HashMap<String, serde_json::Value>,
+}
+
+/// Create datapoints in a dataset
+///
+/// Request body should contain:
+/// - dataset_name: The name of the dataset to add datapoints to
+/// - datapoints: Array of datapoint objects with data, optional target, and optional metadata
+#[post("/datasets/datapoints")]
+async fn create_datapoints(
+    req: web::Json<CreateDatapointsRequest>,
+    db: web::Data<DB>,
+    clickhouse: web::Data<clickhouse::Client>,
+    project_api_key: ProjectApiKey,
+) -> ResponseResult {
+    let project_id = project_api_key.project_id;
+    let db = db.into_inner();
+    let clickhouse = clickhouse.into_inner().as_ref().clone();
+    let request = req.into_inner();
+
+    // Validate that we have datapoints to insert
+    if request.datapoints.is_empty() {
+        return Ok(HttpResponse::BadRequest().json(serde_json::json!({
+            "error": "No datapoints provided"
+        })));
+    }
+
+    // Get dataset metadata from PostgreSQL
+    let dataset_id =
+        db::datasets::get_dataset_id_by_name(&db.pool, &request.dataset_name, project_id).await?;
+
+    // Convert request datapoints to Datapoint structs
+    let datapoints: Vec<Datapoint> = request
+        .datapoints
+        .into_iter()
+        .map(|dp_req| Datapoint {
+            id: Uuid::new_v4(),
+            dataset_id,
+            data: dp_req.data,
+            target: dp_req.target,
+            metadata: dp_req.metadata,
+        })
+        .collect();
+
+    // Convert to ClickHouse datapoints
+    let ch_datapoints: Vec<ch_datapoints::CHDatapoint> = datapoints
+        .iter()
+        .map(|dp| ch_datapoints::CHDatapoint::from_datapoint(dp, project_id))
+        .collect();
+
+    // Insert into ClickHouse
+    ch_datapoints::insert_datapoints(clickhouse, ch_datapoints).await?;
+
+    Ok(HttpResponse::Created().json(serde_json::json!({
+        "message": "Datapoints created successfully",
+        "count": datapoints.len()
+    })))
+}

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

@@ -108,7 +108,7 @@ pub async fn update_eval_datapoint(
             .await?;
 
     // Update database (PostgreSQL)
-    db::evaluations::update_evaluation_datapoint(
+    let trace_id = db::evaluations::update_evaluation_datapoint_and_get_trace_id(
         &db.pool,
         eval_id,
         datapoint_id,
@@ -124,6 +124,7 @@ pub async fn update_eval_datapoint(
         group_id,
         eval_id,
         datapoint_id,
+        trace_id,
         scores_clone,
     )
     .await?;

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

@@ -4,5 +4,6 @@ pub mod datasets;
 pub mod evals;
 pub mod evaluators;
 pub mod metrics;
+pub mod sql;
 pub mod tag;
 pub mod traces;

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

@@ -0,0 +1,47 @@
+use std::{collections::HashMap, sync::Arc};
+
+use actix_web::{HttpResponse, post, web};
+use serde::Deserialize;
+
+use crate::{
+    db::project_api_keys::ProjectApiKey,
+    query_engine::QueryEngine,
+    sql::{self, ClickhouseReadonlyClient},
+};
+
+use crate::routes::types::ResponseResult;
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SqlQueryRequest {
+    pub query: String,
+}
+
+#[post("sql/query")]
+pub async fn execute_sql_query(
+    req: web::Json<SqlQueryRequest>,
+    project_api_key: ProjectApiKey,
+    clickhouse_ro: web::Data<Option<Arc<ClickhouseReadonlyClient>>>,
+    query_engine: web::Data<Arc<QueryEngine>>,
+) -> ResponseResult {
+    let project_id = project_api_key.project_id;
+    let SqlQueryRequest { query } = req.into_inner();
+
+    match clickhouse_ro.as_ref() {
+        Some(ro_client) => {
+            match sql::execute_sql_query(
+                query,
+                project_id,
+                HashMap::new(),
+                ro_client.clone(),
+                query_engine.into_inner().as_ref().clone(),
+            )
+            .await
+            {
+                Ok(result_json) => Ok(HttpResponse::Ok().json(result_json)),
+                Err(e) => Err(e.into()),
+            }
+        }
+        None => Err(anyhow::anyhow!("ClickHouse read-only client is not configured.").into()),
+    }
+}

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

@@ -28,6 +28,7 @@ pub async fn process_traces(
     cache: web::Data<crate::cache::Cache>,
     spans_message_queue: web::Data<Arc<MessageQueue>>,
     db: web::Data<DB>,
+    clickhouse: web::Data<clickhouse::Client>,
 ) -> ResponseResult {
     let db = db.into_inner();
     let cache = cache.into_inner();
@@ -39,6 +40,7 @@ pub async fn process_traces(
     if is_feature_enabled(Feature::UsageLimit) {
         let limits_exceeded = get_workspace_limit_exceeded_by_project_id(
             db.clone(),
+            clickhouse.into_inner().as_ref().clone(),
             cache.clone(),
             project_api_key.project_id,
         )