refactor(webserver): move answer routes to webserver (#2704)

* refactor(webserver): move answer routes to webserver

* fix

crates/tabby-common/Cargo.toml

@@ -23,6 +23,8 @@ derive_builder.workspace = true
 hash-ids.workspace = true
 tracing.workspace = true
 chrono.workspace = true
+axum.workspace = true
+axum-extra = { workspace = true, features = ["typed-header"] }
 
 [dev-dependencies]
 temp_testdir = { workspace = true }

crates/tabby-common/src/axum.rs

@@ -0,0 +1,31 @@
+use axum::http::HeaderName;
+use axum_extra::headers::Header;
+
+use crate::constants::USER_HEADER_FIELD_NAME;
+
+#[derive(Debug)]
+pub struct MaybeUser(pub Option<String>);
+
+pub static USER_HEADER: HeaderName = HeaderName::from_static(USER_HEADER_FIELD_NAME);
+
+impl Header for MaybeUser {
+    fn name() -> &'static axum::http::HeaderName {
+        &USER_HEADER
+    }
+
+    fn decode<'i, I>(values: &mut I) -> Result<Self, axum_extra::headers::Error>
+    where
+        Self: Sized,
+        I: Iterator<Item = &'i axum::http::HeaderValue>,
+    {
+        let Some(value) = values.next() else {
+            return Ok(MaybeUser(None));
+        };
+        let str = value.to_str().expect("User email is always a valid string");
+        Ok(MaybeUser(Some(str.to_string())))
+    }
+
+    fn encode<E: Extend<axum::http::HeaderValue>>(&self, _values: &mut E) {
+        todo!()
+    }
+}

crates/tabby-common/src/lib.rs

@@ -1,6 +1,7 @@
 //! Common tabby types and utilities.
 //! Defines common types and utilities used across multiple tabby subprojects, especially serialization and deserialization targets.
 pub mod api;
+pub mod axum;
 pub mod config;
 pub mod constants;
 pub mod index;

crates/tabby/src/routes/chat.rs

@@ -8,11 +8,10 @@ use axum::{
 use axum_extra::TypedHeader;
 use futures::{Stream, StreamExt};
 use hyper::StatusCode;
+use tabby_common::axum::MaybeUser;
 use tabby_inference::ChatCompletionStream;
 use tracing::{instrument, warn};
 
-use super::MaybeUser;
-
 #[utoipa::path(
     post,
     path = "/v1/chat/completions",

crates/tabby/src/routes/completions.rs

@@ -3,9 +3,9 @@ use std::sync::Arc;
 use axum::{extract::State, Json};
 use axum_extra::TypedHeader;
 use hyper::StatusCode;
+use tabby_common::axum::MaybeUser;
 use tracing::{instrument, warn};
 
-use super::MaybeUser;
 use crate::services::completion::{CompletionRequest, CompletionResponse, CompletionService};
 
 #[utoipa::path(

crates/tabby/src/routes/events.rs

@@ -6,9 +6,10 @@ use axum::{
 };
 use axum_extra::TypedHeader;
 use hyper::StatusCode;
-use tabby_common::api::event::{Event, EventLogger, LogEventRequest, SelectKind};
-
-use super::MaybeUser;
+use tabby_common::{
+    api::event::{Event, EventLogger, LogEventRequest, SelectKind},
+    axum::MaybeUser,
+};
 
 #[utoipa::path(
     post,

crates/tabby/src/routes/mod.rs

@@ -5,10 +5,8 @@ use std::{
     sync::Arc,
 };
 
-use axum::{http::HeaderName, routing, Router};
-use axum_extra::headers::Header;
+use axum::{routing, Router};
 use axum_prometheus::PrometheusMetricLayer;
-use tabby_common::constants::USER_HEADER_FIELD_NAME;
 use tower_http::cors::CorsLayer;
 
 use crate::fatal;
@@ -54,41 +52,12 @@ pub async fn run_app(api: Router, ui: Option<Router>, host: IpAddr, port: u16) {
     .unwrap_or_else(|err| fatal!("Error happens during serving: {}", err))
 }
 
-#[derive(Debug)]
-pub(crate) struct MaybeUser(pub Option<String>);
-
-pub(crate) static USER_HEADER: HeaderName = HeaderName::from_static(USER_HEADER_FIELD_NAME);
-
-impl Header for MaybeUser {
-    fn name() -> &'static axum::http::HeaderName {
-        &USER_HEADER
-    }
-
-    fn decode<'i, I>(values: &mut I) -> Result<Self, axum_extra::headers::Error>
-    where
-        Self: Sized,
-        I: Iterator<Item = &'i axum::http::HeaderValue>,
-    {
-        let Some(value) = values.next() else {
-            return Ok(MaybeUser(None));
-        };
-        let str = value.to_str().expect("User email is always a valid string");
-        Ok(MaybeUser(Some(str.to_string())))
-    }
-
-    fn encode<E: Extend<axum::http::HeaderValue>>(&self, _values: &mut E) {
-        todo!()
-    }
-}
-
-mod answer;
 mod chat;
 mod completions;
 mod events;
 mod health;
 mod server_setting;
 
-pub use answer::*;
 pub use chat::*;
 pub use completions::*;
 pub use events::*;

crates/tabby/src/serve.rs

@@ -9,7 +9,7 @@ use tabby_common::{
     config::{Config, ConfigAccess, ModelConfig, StaticConfigAccess},
     usage,
 };
-use tabby_inference::Embedding;
+use tabby_inference::ChatCompletionStream;
 use tokio::{sync::oneshot::Sender, time::sleep};
 use tower_http::timeout::TimeoutLayer;
 use tracing::{debug, warn};
@@ -22,7 +22,7 @@ use utoipa_swagger_ui::SwaggerUi;
 use crate::{
     routes::{self, run_app},
     services::{
-        self, answer,
+        self,
         code::create_code_search,
         completion::{self, create_completion_service},
         embedding,
@@ -51,7 +51,7 @@ Install following IDE / Editor extensions to get started with [Tabby](https://gi
     servers(
         (url = "/", description = "Server"),
     ),
-    paths(routes::log_event, routes::completions, routes::chat_completions_utoipa, routes::health, routes::answer, routes::setting),
+    paths(routes::log_event, routes::completions, routes::chat_completions_utoipa, routes::health, routes::setting),
     components(schemas(
         api::event::LogEventRequest,
         completion::CompletionRequest,
@@ -67,8 +67,6 @@ Install following IDE / Editor extensions to get started with [Tabby](https://gi
         api::code::CodeSearchDocument,
         api::code::CodeSearchQuery,
         api::doc::DocSearchDocument,
-        answer::AnswerRequest,
-        answer::AnswerResponseChunk,
         api::server_setting::ServerSetting
     )),
     modifiers(&SecurityAddon),
@@ -156,19 +154,29 @@ pub async fn main(config: &Config, args: &ServeArgs) {
     }
 
     let index_reader_provider = Arc::new(IndexReaderProvider::default());
+    let docsearch = Arc::new(services::doc::create(
+        embedding.clone(),
+        index_reader_provider.clone(),
+    ));
 
     let code = Arc::new(create_code_search(
         config_access,
         embedding.clone(),
         index_reader_provider.clone(),
     ));
+
+    let chat = if let Some(chat) = &config.model.chat {
+        Some(model::load_chat_completion(chat).await)
+    } else {
+        None
+    };
+
     let mut api = api_router(
         args,
         &config,
         logger.clone(),
         code.clone(),
-        embedding,
-        index_reader_provider,
+        chat.clone(),
         webserver,
     )
     .await;
@@ -178,7 +186,11 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 
     #[cfg(feature = "ee")]
     if let Some(ws) = &ws {
-        let (new_api, new_ui) = ws.attach(api, ui, code, config.model.chat.is_some()).await;
+        let (new_api, new_ui) = ws
+            .attach(api, ui, code, chat, docsearch, |x| {
+                Box::new(services::doc::create_serper(x))
+            })
+            .await;
         api = new_api;
         ui = new_ui;
     };
@@ -210,8 +222,7 @@ async fn api_router(
     config: &Config,
     logger: Arc<dyn EventLogger>,
     code: Arc<dyn CodeSearch>,
-    embedding: Arc<dyn Embedding>,
-    index_reader_provider: Arc<IndexReaderProvider>,
+    chat_state: Option<Arc<dyn ChatCompletionStream>>,
     webserver: Option<bool>,
 ) -> Router {
     let model = &config.model;
@@ -224,22 +235,6 @@ async fn api_router(
         None
     };
 
-    let chat_state = if let Some(chat) = &model.chat {
-        Some(model::load_chat_completion(chat).await)
-    } else {
-        None
-    };
-
-    let docsearch_state = Arc::new(services::doc::create(embedding, index_reader_provider));
-
-    let answer_state = chat_state.as_ref().map(|chat| {
-        Arc::new(services::answer::create(
-            chat.clone(),
-            code.clone(),
-            docsearch_state.clone(),
-        ))
-    });
-
     let mut routers = vec![];
 
     let health_state = Arc::new(health::HealthState::new(
@@ -318,19 +313,6 @@ async fn api_router(
         });
     }
 
-    if let Some(answer_state) = answer_state {
-        routers.push({
-            Router::new().route(
-                "/v1beta/answer",
-                routing::post(routes::answer).with_state(answer_state),
-            )
-        });
-    } else {
-        routers.push({
-            Router::new().route("/v1beta/answer", routing::post(StatusCode::NOT_IMPLEMENTED))
-        });
-    }
-
     let server_setting_router =
         Router::new().route("/v1beta/server_setting", routing::get(routes::setting));
 

crates/tabby/src/services/mod.rs

@@ -1,4 +1,3 @@
-pub mod answer;
 pub mod code;
 pub mod completion;
 pub mod doc;

ee/tabby-webserver/Cargo.toml

@@ -53,6 +53,8 @@ strum.workspace = true
 cron = "0.12.1"
 async-stream.workspace = true
 logkit.workspace = true
+async-openai.workspace = true
+utoipa.workspace = true
 
 [dev-dependencies]
 assert_matches.workspace = true

crates/tabby/src/routes/answer.rs → ee/tabby-webserver/src/routes/answer.rs

@@ -8,10 +8,10 @@ use axum::{
 };
 use axum_extra::TypedHeader;
 use futures::{Stream, StreamExt};
+use tabby_common::axum::MaybeUser;
 use tracing::instrument;
 
-use super::MaybeUser;
-use crate::services::answer::{AnswerRequest, AnswerService};
+use crate::service::answer::{AnswerRequest, AnswerService};
 
 #[utoipa::path(
     post,

ee/tabby-webserver/src/routes/mod.rs

@@ -1,3 +1,4 @@
+mod answer;
 mod hub;
 mod oauth;
 mod repositories;
@@ -23,16 +24,32 @@ use self::hub::HubState;
 use crate::{
     axum::{extract::AuthBearer, graphql, FromAuth},
     jwt::validate_jwt,
+    service::answer::AnswerService,
 };
 
-pub fn create(ctx: Arc<dyn ServiceLocator>, api: Router, ui: Router) -> (Router, Router) {
+pub fn create(
+    ctx: Arc<dyn ServiceLocator>,
+    api: Router,
+    ui: Router,
+    answer: Option<Arc<AnswerService>>,
+) -> (Router, Router) {
     let schema = Arc::new(create_schema());
 
-    let api = api
-        .route(
-            "/v1beta/server_setting",
-            routing::get(server_setting).with_state(ctx.clone()),
+    let api = api.route(
+        "/v1beta/server_setting",
+        routing::get(server_setting).with_state(ctx.clone()),
+    );
+
+    let api = if let Some(answer) = answer {
+        api.route(
+            "/v1beta/answer",
+            routing::post(answer::answer).with_state(answer),
         )
+    } else {
+        api.route("/v1beta/answer", routing::post(StatusCode::NOT_IMPLEMENTED))
+    };
+
+    let api = api
         // Routes before `distributed_tabby_layer` are protected by authentication middleware for following routes:
         // 1. /v1/*
         // 2. /v1beta/*

crates/tabby/src/services/answer.rs → ee/tabby-webserver/src/service/answer.rs

@@ -64,11 +64,12 @@ impl AnswerService {
         chat: Arc<dyn ChatCompletionStream>,
         code: Arc<dyn CodeSearch>,
         doc: Arc<dyn DocSearch>,
+        serper_factory_fn: impl Fn(&str) -> Box<dyn DocSearch>,
     ) -> Self {
         let serper: Option<Box<dyn DocSearch>> =
             if let Ok(api_key) = std::env::var("SERPER_API_KEY") {
                 debug!("Serper API key found, enabling serper...");
-                Some(Box::new(super::doc::create_serper(api_key.as_str())))
+                Some(serper_factory_fn(&api_key))
             } else {
                 None
             };
@@ -363,8 +364,9 @@ pub fn create(
     chat: Arc<dyn ChatCompletionStream>,
     code: Arc<dyn CodeSearch>,
     doc: Arc<dyn DocSearch>,
+    serper_factory_fn: impl Fn(&str) -> Box<dyn DocSearch>,
 ) -> AnswerService {
-    AnswerService::new(chat, code, doc)
+    AnswerService::new(chat, code, doc, serper_factory_fn)
 }
 
 fn get_content(message: &ChatCompletionRequestMessage) -> &str {