Merge pull request #680 from lmnr-ai/dev

CTEs, improved span processing, update models, support human evaluator, fix queue bug

Author: dinmukhamedm

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

@@ -14,11 +14,13 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use uuid::Uuid;
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Default)]
 #[serde(rename_all = "camelCase")]
 pub struct InitEvalRequest {
     pub name: Option<String>,
     pub group_name: Option<String>,
+    #[serde(default)]
+    pub metadata: Option<Value>,
 }
 
 #[post("/evals")]
@@ -31,15 +33,16 @@ pub async fn init_eval(
     let req = req.into_inner();
     let group_name = req.group_name.unwrap_or("default".to_string());
     let project_id = project_api_key.project_id;
-
+    let metadata = req.metadata;
     let name = if let Some(name) = req.name {
         name
     } else {
         name_generator.next().await
     };
 
     let evaluation =
-        db::evaluations::create_evaluation(&db.pool, &name, project_id, &group_name).await?;
+        db::evaluations::create_evaluation(&db.pool, &name, project_id, &group_name, &metadata)
+            .await?;
 
     Ok(HttpResponse::Ok().json(evaluation))
 }

app-server/src/db/evaluations.rs

@@ -18,6 +18,7 @@ pub struct Evaluation {
     ///
     /// Conceptually, evaluations with different group ids are used to test different features.
     pub group_id: String,
+    pub metadata: Option<Value>,
 }
 
 #[derive(Serialize, FromRow)]
@@ -52,20 +53,23 @@ pub async fn create_evaluation(
     name: &String,
     project_id: Uuid,
     group_id: &str,
+    metadata: &Option<Value>,
 ) -> Result<Evaluation> {
     let evaluation = sqlx::query_as::<_, Evaluation>(
-        "INSERT INTO evaluations (name, project_id, group_id)
-        VALUES ($1, $2, $3)
+        "INSERT INTO evaluations (name, project_id, group_id, metadata)
+        VALUES ($1, $2, $3, $4)
         RETURNING
             id,
             created_at,
             name,
             project_id,
-            group_id",
+            group_id,
+            metadata",
     )
     .bind(name)
     .bind(project_id)
     .bind(group_id)
+    .bind(metadata)
     .fetch_one(pool)
     .await?;
 

app-server/src/main.rs

@@ -169,11 +169,11 @@ fn main() -> anyhow::Result<()> {
 
     let connection_for_health = connection.clone(); // Clone before moving into HttpServer
 
-    // ==== 3.1 Spans message queue ====
-    let spans_message_queue: Arc<MessageQueue> = if let Some(connection) = connection.as_ref() {
+    let queue: Arc<MessageQueue> = if let Some(connection) = connection.as_ref() {
         runtime_handle.block_on(async {
             let channel = connection.create_channel().await.unwrap();
-
+            // Register queues
+            // ==== 3.1 Spans message queue ====
             channel
                 .exchange_declare(
                     OBSERVATIONS_EXCHANGE,
@@ -193,28 +193,7 @@ fn main() -> anyhow::Result<()> {
                 .await
                 .unwrap();
 
-            let max_channel_pool_size = env::var("RABBITMQ_MAX_CHANNEL_POOL_SIZE")
-                .ok()
-                .and_then(|v| v.parse().ok())
-                .unwrap_or(64);
-
-            log::info!("RabbitMQ span channels: {}", max_channel_pool_size);
-
-            let rabbit_mq = mq::rabbit::RabbitMQ::new(connection.clone(), max_channel_pool_size);
-            Arc::new(rabbit_mq.into())
-        })
-    } else {
-        log::info!("Using tokio mpsc span queue");
-        Arc::new(mq::tokio_mpsc::TokioMpscQueue::new().into())
-    };
-
-    // ==== 3.2 Browser events message queue ====
-    let browser_events_message_queue: Arc<MessageQueue> = if let Some(connection) =
-        connection.as_ref()
-    {
-        runtime_handle.block_on(async {
-            let channel = connection.create_channel().await.unwrap();
-
+            // ==== 3.2 Browser events message queue ====
             channel
                 .exchange_declare(
                     BROWSER_SESSIONS_EXCHANGE,
@@ -234,30 +213,7 @@ fn main() -> anyhow::Result<()> {
                 .await
                 .unwrap();
 
-            let max_channel_pool_size = env::var("RABBITMQ_MAX_CHANNEL_POOL_SIZE")
-                .ok()
-                .and_then(|v| v.parse().ok())
-                .unwrap_or(64);
-
-            log::info!(
-                "RabbitMQ browser events channels: {}",
-                max_channel_pool_size
-            );
-
-            let rabbit_mq = mq::rabbit::RabbitMQ::new(connection.clone(), max_channel_pool_size);
-            Arc::new(rabbit_mq.into())
-        })
-    } else {
-        log::info!("Using tokio mpsc browser events queue");
-        Arc::new(mq::tokio_mpsc::TokioMpscQueue::new().into())
-    };
-
-    // ==== 3.3 Evaluators message queue ====
-    let evaluators_message_queue: Arc<MessageQueue> = if let Some(connection) = connection.as_ref()
-    {
-        runtime_handle.block_on(async {
-            let channel = connection.create_channel().await.unwrap();
-
+            // ==== 3.3 Evaluators message queue ====
             channel
                 .exchange_declare(
                     EVALUATORS_EXCHANGE,
@@ -282,14 +238,22 @@ fn main() -> anyhow::Result<()> {
                 .and_then(|v| v.parse().ok())
                 .unwrap_or(64);
 
-            log::info!("RabbitMQ evaluators channels: {}", max_channel_pool_size);
+            log::info!("RabbitMQ channels: {}", max_channel_pool_size);
 
             let rabbit_mq = mq::rabbit::RabbitMQ::new(connection.clone(), max_channel_pool_size);
             Arc::new(rabbit_mq.into())
         })
     } else {
-        log::info!("Using tokio mpsc evaluators queue");
-        Arc::new(mq::tokio_mpsc::TokioMpscQueue::new().into())
+        let queue = mq::tokio_mpsc::TokioMpscQueue::new();
+        // register queues
+        // ==== 3.1 Spans message queue ====
+        queue.register_queue(OBSERVATIONS_EXCHANGE, OBSERVATIONS_QUEUE);
+        // ==== 3.2 Browser events message queue ====
+        queue.register_queue(BROWSER_SESSIONS_EXCHANGE, BROWSER_SESSIONS_QUEUE);
+        // ==== 3.3 Evaluators message queue ====
+        queue.register_queue(EVALUATORS_EXCHANGE, EVALUATORS_QUEUE);
+        log::info!("Using tokio mpsc queue");
+        Arc::new(queue.into())
     };
 
     // ==== 3.4 Agent worker message queue ====
@@ -298,7 +262,7 @@ fn main() -> anyhow::Result<()> {
     let runtime_handle_for_http = runtime_handle.clone();
     let db_for_http = db.clone();
     let cache_for_http = cache.clone();
-    let spans_mq_for_http = spans_message_queue.clone();
+    let mq_for_http = queue.clone();
 
     // == HTTP server and listener workers ==
     let http_server_handle = thread::Builder::new()
@@ -456,8 +420,7 @@ fn main() -> anyhow::Result<()> {
                         tokio::spawn(process_queue_spans(
                             db_for_http.clone(),
                             cache_for_http.clone(),
-                            spans_mq_for_http.clone(),
-                            evaluators_message_queue.clone(),
+                            mq_for_http.clone(),
                             clickhouse.clone(),
                             storage.clone(),
                         ));
@@ -467,15 +430,15 @@ fn main() -> anyhow::Result<()> {
                         tokio::spawn(process_browser_events(
                             db_for_http.clone(),
                             clickhouse.clone(),
-                            browser_events_message_queue.clone(),
+                            mq_for_http.clone(),
                         ));
                     }
 
                     for _ in 0..num_evaluators_workers_per_thread {
                         tokio::spawn(process_evaluators(
                             db_for_http.clone(),
                             clickhouse.clone(),
-                            evaluators_message_queue.clone(),
+                            mq_for_http.clone(),
                             evaluator_client.clone(),
                             python_online_evaluator_url.clone(),
                         ));
@@ -488,13 +451,11 @@ fn main() -> anyhow::Result<()> {
                         .app_data(PayloadConfig::new(http_payload_limit))
                         .app_data(web::Data::from(cache_for_http.clone()))
                         .app_data(web::Data::from(db_for_http.clone()))
-                        .app_data(web::Data::new(spans_mq_for_http.clone()))
+                        .app_data(web::Data::new(mq_for_http.clone()))
                         .app_data(web::Data::new(clickhouse.clone()))
                         .app_data(web::Data::new(name_generator.clone()))
                         .app_data(web::Data::new(storage.clone()))
                         .app_data(web::Data::new(machine_manager.clone()))
-                        .app_data(web::Data::new(browser_events_message_queue.clone()))
-                        .app_data(web::Data::new(evaluators_message_queue.clone()))
                         .app_data(web::Data::new(agent_manager_workers.clone()))
                         .app_data(web::Data::new(connection_for_health.clone()))
                         .app_data(web::Data::new(browser_agent.clone()))
@@ -570,7 +531,7 @@ fn main() -> anyhow::Result<()> {
                 let process_traces_service = ProcessTracesService::new(
                     db.clone(),
                     cache.clone(),
-                    spans_message_queue.clone(),
+                    queue.clone(),
                 );
 
                 Server::builder()

app-server/src/mq/rabbit.rs

@@ -1,10 +1,10 @@
 use deadpool::managed::{Manager, Pool, PoolError, RecycleError};
 use futures_util::StreamExt;
 use lapin::{
+    BasicProperties, Channel, Connection, Consumer,
     acker::Acker,
     options::{BasicConsumeOptions, BasicPublishOptions, QueueBindOptions},
     types::FieldTable,
-    BasicProperties, Channel, Connection, Consumer,
 };
 use std::sync::Arc;
 

app-server/src/mq/tokio_mpsc.rs

@@ -5,8 +5,8 @@ use super::{
 use dashmap::DashMap;
 use std::sync::Arc;
 use tokio::sync::{
-    mpsc::{self, Receiver, Sender},
     Mutex,
+    mpsc::{self, Receiver, Sender},
 };
 
 const CHANNEL_CAPACITY: usize = 100;
@@ -53,6 +53,11 @@ impl TokioMpscQueue {
     fn key(&self, exchange: &str, routing_key: &str) -> String {
         format!("{}:-:{}", exchange, routing_key)
     }
+
+    pub fn register_queue(&self, exchange: &str, routing_key: &str) {
+        let key = self.key(exchange, routing_key);
+        self.senders.entry(key).or_default();
+    }
 }
 
 impl MessageQueueTrait for TokioMpscQueue {

app-server/src/traces/consumer.rs

@@ -21,7 +21,6 @@ pub async fn process_queue_spans(
     db: Arc<DB>,
     cache: Arc<Cache>,
     queue: Arc<MessageQueue>,
-    evaluators_queue: Arc<MessageQueue>,
     clickhouse: clickhouse::Client,
     storage: Arc<Storage>,
 ) {
@@ -30,7 +29,6 @@ pub async fn process_queue_spans(
             db.clone(),
             cache.clone(),
             queue.clone(),
-            evaluators_queue.clone(),
             clickhouse.clone(),
             storage.clone(),
         )
@@ -43,7 +41,6 @@ async fn inner_process_queue_spans(
     db: Arc<DB>,
     cache: Arc<Cache>,
     queue: Arc<MessageQueue>,
-    evaluators_queue: Arc<MessageQueue>,
     clickhouse: clickhouse::Client,
     storage: Arc<Storage>,
 ) {
@@ -142,7 +139,7 @@ async fn inner_process_queue_spans(
             clickhouse.clone(),
             cache.clone(),
             acker,
-            evaluators_queue.clone(),
+            queue.clone(),
         )
         .await;
     }

app-server/src/traces/spans.rs

@@ -938,11 +938,26 @@ fn output_message_from_completion_content(
 
     if tool_calls.is_empty() {
         if let Some(Value::String(s)) = msg_content {
-            Some(ChatMessage {
-                role: msg_role,
-                content: ChatMessageContent::Text(s.clone()),
-                tool_call_id: None,
-            })
+            if let Ok(content) =
+                serde_json::from_str::<Vec<InstrumentationChatMessageContentPart>>(&s)
+            {
+                Some(ChatMessage {
+                    role: msg_role,
+                    content: ChatMessageContent::ContentPartList(
+                        content
+                            .into_iter()
+                            .map(ChatMessageContentPart::from_instrumentation_content_part)
+                            .collect(),
+                    ),
+                    tool_call_id: None,
+                })
+            } else {
+                Some(ChatMessage {
+                    role: msg_role,
+                    content: ChatMessageContent::Text(s.clone()),
+                    tool_call_id: None,
+                })
+            }
         } else {
             None
         }