From 5038167f31c442eaccf18364a3be33a2f0a32150 Mon Sep 17 00:00:00 2001
From: boxbeam <instantiable@gmail.com>
Date: Thu, 30 May 2024 18:32:40 -0400
Subject: [PATCH] feat(tabby): populate user field for chat completion logging
 (#2302)

---
 crates/tabby/src/routes/chat.rs   | 9 ++++++++-
 crates/tabby/src/services/chat.rs | 6 ++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/crates/tabby/src/routes/chat.rs b/crates/tabby/src/routes/chat.rs
index bf64bdf991f7..c14b515d8297 100644
--- 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
index c7f303db51b2..b8d31ecbcd72 100644
--- a/crates/tabby/src/services/chat.rs
+++ b/crates/tabby/src/services/chat.rs
@@ -27,6 +27,9 @@ use super::model;
     ]
 }))]
 pub struct ChatCompletionRequest {
+    #[builder(default = "None")]
+    pub(crate) user: Option<String>,
+
     messages: Vec<Message>,
 
     #[builder(default = "None")]
@@ -135,8 +138,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)