chore: log v1/chat/completion event (#4042)

Author: zwpaper

crates/tabby/src/routes/chat.rs

@@ -9,7 +9,10 @@ use axum::{
 use axum_extra::TypedHeader;
 use futures::{Stream, StreamExt};
 use hyper::StatusCode;
-use tabby_common::axum::MaybeUser;
+use tabby_common::{
+    api::event::{Event as LoggerEvent, EventLogger},
+    axum::MaybeUser,
+};
 use tabby_inference::ChatCompletionStream;
 use tracing::{error, instrument, warn};
 
@@ -32,17 +35,23 @@ pub async fn chat_completions_utoipa(_request: Json<serde_json::Value>) -> Statu
     unimplemented!()
 }
 
+pub struct ChatState {
+    pub chat_completion: Arc<dyn ChatCompletionStream>,
+    pub logger: Arc<dyn EventLogger>,
+}
+
 #[instrument(skip(state, request))]
 pub async fn chat_completions(
-    State(state): State<Arc<dyn ChatCompletionStream>>,
+    State(state): State<Arc<ChatState>>,
     TypedHeader(MaybeUser(user)): TypedHeader<MaybeUser>,
     Json(mut request): Json<async_openai_alt::types::CreateChatCompletionRequest>,
 ) -> Result<Sse<impl Stream<Item = Result<Event, anyhow::Error>>>, StatusCode> {
     if let Some(user) = user {
         request.user.replace(user);
     }
+    let user = request.user.clone();
 
-    let s = match state.chat_stream(request).await {
+    let s = match state.chat_completion.chat_stream(request).await {
         Ok(s) => s,
         Err(err) => {
             warn!("Error happens during chat completion: {}", err);
@@ -71,5 +80,7 @@ pub async fn chat_completions(
         }
     };
 
+    state.logger.log(user, LoggerEvent::ChatCompletion {});
+
     Ok(Sse::new(s).keep_alive(KeepAlive::default()))
 }

crates/tabby/src/serve.rs

@@ -11,7 +11,6 @@ use tabby_common::{
     usage,
 };
 use tabby_download::ModelKind;
-use tabby_inference::ChatCompletionStream;
 use tokio::{sync::oneshot::Sender, time::sleep};
 use tower_http::timeout::TimeoutLayer;
 use tracing::{debug, warn};
@@ -22,7 +21,7 @@ use utoipa::{
 use utoipa_swagger_ui::SwaggerUi;
 
 use crate::{
-    routes::{self, run_app},
+    routes::{self, run_app, ChatState},
     services::{
         self,
         code::create_code_search,
@@ -160,13 +159,19 @@ pub async fn main(config: &Config, args: &ServeArgs) {
     )
     .await;
 
+    let chat_state = chat.as_ref().map(|c| {
+        Arc::new(ChatState {
+            chat_completion: c.clone(),
+            logger: logger.clone(),
+        })
+    });
     let mut api = api_router(
         args,
         &config,
         logger.clone(),
         code.clone(),
         completion,
-        chat.clone(),
+        chat_state,
         webserver,
     )
     .await;
@@ -220,7 +225,7 @@ async fn api_router(
     logger: Arc<dyn EventLogger>,
     _code: Arc<dyn CodeSearch>,
     completion_state: Option<CompletionService>,
-    chat_state: Option<Arc<dyn ChatCompletionStream>>,
+    chat_state: Option<Arc<ChatState>>,
     webserver: Option<bool>,
 ) -> Router {
     let mut routers = vec![];