feat(tabby): populate user field for chat completion logging

TabbyML · May 30, 2024 · 95cc554 · 95cc554
1 parent 95feb55
commit 95cc554
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/crates/tabby/src/routes/chat.rs b/crates/tabby/src/routes/chat.rs
@@ -5,9 +5,11 @@ use axum::{
     response::sse::{Event, KeepAlive, Sse},
     Json,
 };
+use axum_extra::TypedHeader;
 use futures::{Stream, StreamExt};
 use tracing::instrument;
 
+use super::MaybeUser;
 use crate::services::chat::{ChatCompletionRequest, ChatService};
 
 #[utoipa::path(
@@ -28,8 +30,13 @@ use crate::services::chat::{ChatCompletionRequest, ChatService};
 #[instrument(skip(state, request))]
 pub async fn chat_completions(
     State(state): State<Arc<ChatService>>,
-    Json(request): Json<ChatCompletionRequest>,
+    TypedHeader(MaybeUser(user)): TypedHeader<MaybeUser>,
+    Json(mut request): Json<ChatCompletionRequest>,
 ) -> Sse<impl Stream<Item = Result<Event, serde_json::Error>>> {
+    if let Some(user) = user {
+        request.user.replace(user);
+    }
+
     let stream = state.generate(request).await;
     Sse::new(stream.map(|chunk| match serde_json::to_string(&chunk) {
         Ok(s) => Ok(Event::default().data(s)),

diff --git a/crates/tabby/src/services/chat.rs b/crates/tabby/src/services/chat.rs
@@ -27,6 +27,8 @@ use super::model;
     ]
 }))]
 pub struct ChatCompletionRequest {
+    pub(crate) user: Option<String>,
+
     messages: Vec<Message>,
 
     #[builder(default = "None")]
@@ -135,8 +137,7 @@ impl ChatService {
             }
             yield ChatCompletionChunk::new(String::default(), completion_id.clone(), created, true);
 
-            // FIXME(boxbeam): Log user for chat completion events
-            self.logger.log(None, Event::ChatCompletion {
+            self.logger.log(request.user, Event::ChatCompletion {
                 completion_id,
                 input: convert_messages(&request.messages),
                 output: create_assistant_message(output)