extract EventLogger service

Cargo.lock

@@ -4282,6 +4282,7 @@ dependencies = [
  "axum",
  "axum-streams",
  "axum-tracing-opentelemetry",
+ "chrono",
  "clap 4.4.7",
  "futures",
  "http-api-bindings",
@@ -4324,7 +4325,6 @@ name = "tabby-common"
 version = "0.6.0-dev"
 dependencies = [
  "anyhow",
- "chrono",
  "filenamify",
  "lazy_static",
  "reqwest",
@@ -4332,7 +4332,6 @@ dependencies = [
  "serde-jsonlines",
  "serdeconv",
  "tantivy",
- "tokio",
  "uuid 1.4.1",
 ]
 

crates/tabby-common/Cargo.toml

@@ -4,14 +4,12 @@ version = "0.6.0-dev"
 edition = "2021"
 
 [dependencies]
-chrono = "0.4.26"
 filenamify = "0.1.0"
 lazy_static = { workspace = true }
 serde = { workspace = true }
 serdeconv = { workspace = true }
 serde-jsonlines = { workspace = true }
 reqwest = { workspace = true, features = [ "json" ] }
-tokio = { workspace = true, features = ["rt", "macros"] }
 uuid = { version = "1.4.1", features = ["v4"] }
 tantivy.workspace = true
 anyhow.workspace = true

crates/tabby-common/src/lib.rs

@@ -1,5 +1,4 @@
 pub mod config;
-pub mod events;
 pub mod index;
 pub mod languages;
 pub mod path;

crates/tabby/Cargo.toml

@@ -46,6 +46,7 @@ futures.workspace = true
 async-trait.workspace = true
 tabby-webserver = { path = "../../ee/tabby-webserver" }
 thiserror.workspace = true
+chrono = "0.4.31"
 
 [dependencies.uuid]
 version = "1.3.3"

crates/tabby/src/api/event.rs

@@ -1,9 +1,4 @@
-
-
-
-
 use serde::{Deserialize, Serialize};
-
 use utoipa::ToSchema;
 
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@@ -17,3 +12,53 @@ pub struct LogEventRequest {
 
     pub choice_index: u32,
 }
+
+#[derive(Serialize)]
+pub struct Choice<'a> {
+    pub index: u32,
+    pub text: &'a str,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum SelectKind {
+    Line,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Event<'a> {
+    View {
+        completion_id: &'a str,
+        choice_index: u32,
+    },
+    Select {
+        completion_id: &'a str,
+        choice_index: u32,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        kind: Option<SelectKind>,
+    },
+    Completion {
+        completion_id: &'a str,
+        language: &'a str,
+        prompt: &'a str,
+        segments: &'a Option<Segments>,
+        choices: Vec<Choice<'a>>,
+        user: Option<&'a str>,
+    },
+}
+#[derive(Serialize)]
+pub struct Segments {
+    pub prefix: String,
+    pub suffix: Option<String>,
+}
+
+#[derive(Serialize)]
+struct Log<'a> {
+    ts: u128,
+    event: &'a Event<'a>,
+}
+
+pub trait EventLogger: Send + Sync {
+    fn log(&self, e: &Event);
+}

crates/tabby/src/routes/events.rs

@@ -1,12 +1,13 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
-use axum::{extract::Query, Json};
+use axum::{
+    extract::{Query, State},
+    Json,
+};
 use hyper::StatusCode;
-
-use tabby_common::events::{self, SelectKind};
 use utoipa::ToSchema;
 
-use crate::api::LogEventRequest;
+use crate::api::{Event, EventLogger, LogEventRequest, SelectKind};
 
 #[utoipa::path(
     post,
@@ -20,22 +21,22 @@ use crate::api::LogEventRequest;
     )
 )]
 pub async fn log_event(
+    State(logger): State<Arc<dyn EventLogger>>,
     Query(params): Query<HashMap<String, String>>,
     Json(request): Json<LogEventRequest>,
 ) -> StatusCode {
     if request.event_type == "view" {
-        events::Event::View {
+        logger.log(&Event::View {
             completion_id: &request.completion_id,
             choice_index: request.choice_index,
-        }
+        });
-        .log();
         StatusCode::OK
     } else if request.event_type == "select" {
         let is_line = params
             .get("select_kind")
             .map(|x| x == "line")
             .unwrap_or(false);
-        events::Event::Select {
+        logger.log(&Event::Select {
             completion_id: &request.completion_id,
             choice_index: request.choice_index,
             kind: if is_line {
@@ -43,8 +44,7 @@ pub async fn log_event(
             } else {
                 None
             },
-        }
+        });
-        .log();
         StatusCode::OK
     } else {
         StatusCode::BAD_REQUEST

crates/tabby/src/serve/mod.rs

@@ -18,8 +18,9 @@ use utoipa::OpenApi;
 use utoipa_swagger_ui::SwaggerUi;
 
 use crate::{
-    api, fatal, routes,
+    api::{self},
-    services::{chat, completions, health, model},
+    fatal, routes,
+    services::{chat, completions, event::create_event_logger, health, model},
 };
 
 #[derive(OpenApi)]
@@ -162,6 +163,7 @@ async fn load_model(args: &ServeArgs) {
 }
 
 async fn api_router(args: &ServeArgs, config: &Config) -> Router {
+    let logger = Arc::new(create_event_logger());
     let code = Arc::new(crate::services::code::create_code_search());
     let completion_state = {
         let (
@@ -170,8 +172,12 @@ async fn api_router(args: &ServeArgs, config: &Config) -> Router {
                 prompt_template, ..
             },
         ) = model::load_text_generation(&args.model, &args.device, args.parallelism).await;
-        let state =
+        let state = completions::CompletionService::new(
-            completions::CompletionService::new(engine.clone(), code.clone(), prompt_template);
+            engine.clone(),
+            code.clone(),
+            logger.clone(),
+            prompt_template,
+        );
         Arc::new(state)
     };
 
@@ -196,7 +202,10 @@ async fn api_router(args: &ServeArgs, config: &Config) -> Router {
     ));
     routers.push({
         Router::new()
-            .route("/v1/events", routing::post(routes::log_event))
+            .route(
+                "/v1/events",
+                routing::post(routes::log_event).with_state(logger),
+            )
             .route(
                 "/v1/health",
                 routing::post(routes::health).with_state(health_state.clone()),

crates/tabby/src/services/completions.rs

@@ -3,13 +3,13 @@ mod completions_prompt;
 use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
-use tabby_common::{events, languages::get_language};
+use tabby_common::languages::get_language;
 use tabby_inference::{TextGeneration, TextGenerationOptions, TextGenerationOptionsBuilder};
 use thiserror::Error;
 use tracing::debug;
 use utoipa::ToSchema;
 
-use crate::api::CodeSearch;
+use crate::api::{self, CodeSearch, Event, EventLogger};
 
 #[derive(Error, Debug)]
 pub enum CompletionError {
@@ -95,9 +95,9 @@ pub struct Segments {
     suffix: Option<String>,
 }
 
-impl From<Segments> for events::Segments {
+impl From<Segments> for api::Segments {
     fn from(val: Segments) -> Self {
-        events::Segments {
+        api::Segments {
             prefix: val.prefix,
             suffix: val.suffix,
         }
@@ -157,6 +157,7 @@ pub struct DebugData {
 
 pub struct CompletionService {
     engine: Arc<dyn TextGeneration>,
+    logger: Arc<dyn EventLogger>,
     prompt_builder: completions_prompt::PromptBuilder,
 }
 
@@ -164,11 +165,13 @@ impl CompletionService {
     pub fn new(
         engine: Arc<dyn TextGeneration>,
         code: Arc<dyn CodeSearch>,
+        logger: Arc<dyn EventLogger>,
         prompt_template: Option<String>,
     ) -> Self {
         Self {
             engine,
             prompt_builder: completions_prompt::PromptBuilder::new(prompt_template, Some(code)),
+            logger,
         }
     }
 
@@ -226,18 +229,17 @@ impl CompletionService {
         let text = self.engine.generate(&prompt, options).await;
         let segments = segments.map(|s| s.into());
 
-        events::Event::Completion {
+        self.logger.log(&Event::Completion {
             completion_id: &completion_id,
             language: &language,
             prompt: &prompt,
             segments: &segments,
-            choices: vec![events::Choice {
+            choices: vec![api::Choice {
                 index: 0,
                 text: &text,
             }],
             user: request.user.as_deref(),
-        }
+        });
-        .log();
 
         let debug_data = request
             .debug_options

crates/tabby/src/services/event.rs

@@ -7,17 +7,20 @@ use std::{
 use chrono::Utc;
 use lazy_static::lazy_static;
 use serde::Serialize;
+use tabby_common::path;
 use tokio::{
     sync::mpsc::{unbounded_channel, UnboundedSender},
     time::{self},
 };
 
+use crate::api::{Event, EventLogger};
+
 lazy_static! {
     static ref WRITER: UnboundedSender<String> = {
         let (tx, mut rx) = unbounded_channel::<String>();
 
         tokio::spawn(async move {
-            let events_dir = crate::path::events_dir();
+            let events_dir = path::events_dir();
             std::fs::create_dir_all(events_dir.as_path()).ok();
 
             let now = Utc::now();
@@ -53,45 +56,7 @@ lazy_static! {
     };
 }
 
-#[derive(Serialize)]
+struct EventService;
-pub struct Choice<'a> {
-    pub index: u32,
-    pub text: &'a str,
-}
-
-#[derive(Serialize)]
-#[serde(rename_all = "snake_case")]
-pub enum SelectKind {
-    Line,
-}
-
-#[derive(Serialize)]
-#[serde(rename_all = "snake_case")]
-pub enum Event<'a> {
-    View {
-        completion_id: &'a str,
-        choice_index: u32,
-    },
-    Select {
-        completion_id: &'a str,
-        choice_index: u32,
-        #[serde(skip_serializing_if = "Option::is_none")]
-        kind: Option<SelectKind>,
-    },
-    Completion {
-        completion_id: &'a str,
-        language: &'a str,
-        prompt: &'a str,
-        segments: &'a Option<Segments>,
-        choices: Vec<Choice<'a>>,
-        user: Option<&'a str>,
-    },
-}
-#[derive(Serialize)]
-pub struct Segments {
-    pub prefix: String,
-    pub suffix: Option<String>,
-}
 
 #[derive(Serialize)]
 struct Log<'a> {
@@ -99,11 +64,11 @@ struct Log<'a> {
     event: &'a Event<'a>,
 }
 
-impl Event<'_> {
+impl EventLogger for EventService {
-    pub fn log(&self) {
+    fn log(&self, e: &Event) {
         let content = serdeconv::to_json_string(&Log {
             ts: timestamp(),
-            event: self,
+            event: e,
         })
         .unwrap();
 
@@ -119,3 +84,7 @@ fn timestamp() -> u128 {
         .expect("Time went backwards")
         .as_millis()
 }
+
+pub fn create_event_logger() -> impl EventLogger {
+    EventService
+}

crates/tabby/src/services/mod.rs

@@ -1,5 +1,6 @@
 pub mod chat;
 pub mod code;
 pub mod completions;
+pub mod event;
 pub mod health;
 pub mod model;