diff --git a/Cargo.lock b/Cargo.lock index 81653fb..dabb7e6 100644 --- a/Cargo.lock +++ b/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", ] diff --git a/Cargo.toml b/Cargo.toml index 1a11930..236815e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,4 +39,4 @@ thiserror = "1.0.49" utoipa = "3.3" axum = "0.6" hyper = "0.14" -juniper = "0.15" \ No newline at end of file +juniper = "0.15" diff --git a/crates/tabby-common/Cargo.toml b/crates/tabby-common/Cargo.toml index ccaac95..c222370 100644 --- a/crates/tabby-common/Cargo.toml +++ b/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 diff --git a/crates/tabby-common/src/lib.rs b/crates/tabby-common/src/lib.rs index 2458fef..ec17a42 100644 --- a/crates/tabby-common/src/lib.rs +++ b/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; diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml index 1d6450a..66d9e04 100644 --- a/crates/tabby/Cargo.toml +++ b/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" diff --git a/crates/tabby/src/api/event.rs b/crates/tabby/src/api/event.rs index 81193c9..1688ffb 100644 --- a/crates/tabby/src/api/event.rs +++ b/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, + }, + Completion { + completion_id: &'a str, + language: &'a str, + prompt: &'a str, + segments: &'a Option, + choices: Vec>, + user: Option<&'a str>, + }, +} +#[derive(Serialize)] +pub struct Segments { + pub prefix: String, + pub suffix: Option, +} + +#[derive(Serialize)] +struct Log<'a> { + ts: u128, + event: &'a Event<'a>, +} + +pub trait EventLogger: Send + Sync { + fn log(&self, e: &Event); +} diff --git a/crates/tabby/src/routes/events.rs b/crates/tabby/src/routes/events.rs index e509d15..99061f3 100644 --- a/crates/tabby/src/routes/events.rs +++ b/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>, Query(params): Query>, Json(request): Json, ) -> 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 diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs index 0422764..3fb2b71 100644 --- a/crates/tabby/src/serve/mod.rs +++ b/crates/tabby/src/serve/mod.rs @@ -18,8 +18,9 @@ use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; use crate::{ - api, fatal, routes, - services::{chat, completions, health, model}, + api::{self}, + 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 = - completions::CompletionService::new(engine.clone(), code.clone(), prompt_template); + let state = completions::CompletionService::new( + 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()), diff --git a/crates/tabby/src/services/completions.rs b/crates/tabby/src/services/completions.rs index 80bf2b8..3f04c45 100644 --- a/crates/tabby/src/services/completions.rs +++ b/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, } -impl From for events::Segments { +impl From 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, + logger: Arc, prompt_builder: completions_prompt::PromptBuilder, } @@ -164,11 +165,13 @@ impl CompletionService { pub fn new( engine: Arc, code: Arc, + logger: Arc, prompt_template: Option, ) -> 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 diff --git a/crates/tabby-common/src/events.rs b/crates/tabby/src/services/event.rs similarity index 66% rename from crates/tabby-common/src/events.rs rename to crates/tabby/src/services/event.rs index dacec9b..bf910e5 100644 --- a/crates/tabby-common/src/events.rs +++ b/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 = { let (tx, mut rx) = unbounded_channel::(); 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)] -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, - }, - Completion { - completion_id: &'a str, - language: &'a str, - prompt: &'a str, - segments: &'a Option, - choices: Vec>, - user: Option<&'a str>, - }, -} -#[derive(Serialize)] -pub struct Segments { - pub prefix: String, - pub suffix: Option, -} +struct EventService; #[derive(Serialize)] struct Log<'a> { @@ -99,11 +64,11 @@ struct Log<'a> { event: &'a Event<'a>, } -impl Event<'_> { - pub fn log(&self) { +impl EventLogger for EventService { + 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 +} diff --git a/crates/tabby/src/services/mod.rs b/crates/tabby/src/services/mod.rs index 5329861..8a5ade9 100644 --- a/crates/tabby/src/services/mod.rs +++ b/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;