add thread_id and repo_ref to analytics

Author: oppiliappan

server/bleep/src/analytics.rs

@@ -1,6 +1,9 @@
 use std::sync::Arc;
 
-use crate::state::{PersistedState, StateSource};
+use crate::{
+    repo::RepoRef,
+    state::{PersistedState, StateSource},
+};
 
 use rudderanalytics::{
     client::RudderAnalytics,
@@ -13,6 +16,8 @@ use tracing::{info, warn};
 #[derive(Debug, Clone)]
 pub struct QueryEvent {
     pub query_id: uuid::Uuid,
+    pub thread_id: uuid::Uuid,
+    pub repo_ref: Option<RepoRef>,
     pub data: EventData,
 }
 
@@ -137,6 +142,8 @@ impl RudderHub {
                         event: "openai query".to_owned(),
                         properties: Some(json!({
                             "query_id": ev.query_id,
+                            "thread_id": ev.thread_id,
+                            "repo_ref": ev.repo_ref.as_ref().map(ToString::to_string),
                             "data": ev.data,
                             "package_metadata": options.package_metadata,
                         })),

server/bleep/src/webserver/answer.rs

@@ -62,7 +62,7 @@ pub(super) async fn handle(
 ) -> super::Result<impl IntoResponse> {
     let query_id = uuid::Uuid::new_v4();
     let response = _handle(
-        Query(params),
+        Query(params.clone()),
         Extension(app.clone()),
         Extension(user.clone()),
         query_id,
@@ -74,6 +74,8 @@ pub(super) async fn handle(
             &user,
             &QueryEvent {
                 query_id,
+                thread_id: params.thread_id,
+                repo_ref: Some(params.repo_ref.clone()),
                 data: EventData::output_stage("error")
                     .with_payload("status", err.status.as_u16())
                     .with_payload("message", err.message()),
@@ -108,6 +110,7 @@ pub(super) async fn _handle(
     let ctx = AppContext::new(app, user)
         .map_err(|e| super::Error::user(e).with_status(StatusCode::UNAUTHORIZED))?
         .with_query_id(query_id)
+        .with_thread_id(params.thread_id)
         .with_repo_ref(params.repo_ref.clone());
 
     // confirm client compatibility with answer-api
@@ -1137,8 +1140,9 @@ impl Action {
 struct AppContext {
     app: Application,
     llm_gateway: llm_gateway::Client,
-    query_id: uuid::Uuid,
     user: User,
+    query_id: uuid::Uuid,
+    thread_id: uuid::Uuid,
     repo_ref: Option<RepoRef>,
 }
 
@@ -1151,8 +1155,9 @@ impl AppContext {
         Ok(Self {
             app,
             llm_gateway,
-            query_id: uuid::Uuid::nil(),
             user,
+            query_id: uuid::Uuid::nil(),
+            thread_id: uuid::Uuid::nil(),
             repo_ref: None,
         })
     }
@@ -1162,6 +1167,11 @@ impl AppContext {
         self
     }
 
+    fn with_thread_id(mut self, thread_id: uuid::Uuid) -> Self {
+        self.thread_id = thread_id;
+        self
+    }
+
     fn with_repo_ref(mut self, repo_ref: RepoRef) -> Self {
         self.repo_ref = Some(repo_ref);
         self
@@ -1180,6 +1190,8 @@ impl AppContext {
     fn track_query(&self, data: EventData) {
         let event = QueryEvent {
             query_id: self.query_id,
+            thread_id: self.thread_id,
+            repo_ref: self.repo_ref.clone(),
             data,
         };
         self.app.track_query(&self.user, &event);