refactor: extract tabby-common/src/code

Cargo.lock

@@ -4125,6 +4125,7 @@ dependencies = [
  "tantivy",
  "thiserror",
  "tokio",
+ "tracing",
  "utoipa",
  "uuid 1.4.1",
 ]

crates/tabby-common/Cargo.toml

@@ -18,6 +18,7 @@ anyhow.workspace = true
 async-trait.workspace = true
 thiserror.workspace = true
 utoipa = { workspace = true, features = ["axum_extras", "preserve_order"] }
+tracing.workspace = true
 
 [features]
 testutils = []

crates/tabby-common/src/api/mod.rs

@@ -1 +0,0 @@
-pub mod code;

crates/tabby-common/src/code/api.rs

@@ -54,3 +54,5 @@ pub trait CodeSearch: Send + Sync {
         offset: usize,
     ) -> Result<SearchResponse, CodeSearchError>;
 }
+
+pub type BoxCodeSearch = Box<dyn CodeSearch>;

crates/tabby-common/src/code/imp.rs

@@ -2,11 +2,6 @@ use std::{sync::Arc, time::Duration};
 
 use anyhow::Result;
 use async_trait::async_trait;
-use tabby_common::{
-    api::code::{CodeSearch, CodeSearchError, Hit, HitDocument, SearchResponse},
-    index::{self, register_tokenizers, CodeSearchSchema},
-    path,
-};
 use tantivy::{
     collector::{Count, TopDocs},
     query::QueryParser,
@@ -16,6 +11,12 @@ use tantivy::{
 use tokio::{sync::Mutex, time::sleep};
 use tracing::{debug, log::info};
 
+use super::api::{CodeSearch, CodeSearchError, Hit, HitDocument, SearchResponse};
+use crate::{
+    index::{self, register_tokenizers, CodeSearchSchema},
+    path,
+};
+
 struct CodeSearchImpl {
     reader: IndexReader,
     query_parser: QueryParser,
@@ -118,12 +119,12 @@ fn get_field(doc: &Document, field: Field) -> String {
         .to_owned()
 }
 
-pub struct CodeSearchService {
+pub(crate) struct CodeSearchService {
     search: Arc<Mutex<Option<CodeSearchImpl>>>,
 }
 
 impl CodeSearchService {
-    pub fn new() -> Self {
+    pub(crate) fn new() -> Self {
         let search = Arc::new(Mutex::new(None));
 
         let ret = Self {

crates/tabby-common/src/code/mod.rs

@@ -0,0 +1,8 @@
+mod api;
+mod imp;
+
+pub use api::*;
+
+pub fn create_local() -> BoxCodeSearch {
+    Box::new(imp::CodeSearchService::new())
+}

crates/tabby-common/src/lib.rs

@@ -1,4 +1,4 @@
-pub mod api;
+pub mod code;
 pub mod config;
 pub mod events;
 pub mod index;

crates/tabby/src/main.rs

@@ -1,6 +1,5 @@
 mod chat;
 mod download;
-mod search;
 mod serve;
 
 use clap::{Parser, Subcommand};

crates/tabby/src/serve/completions.rs

@@ -5,13 +5,11 @@ use std::sync::Arc;
 use axum::{extract::State, Json};
 use hyper::StatusCode;
 use serde::{Deserialize, Serialize};
-use tabby_common::{events, languages::get_language};
+use tabby_common::{code::BoxCodeSearch, events, languages::get_language};
 use tabby_inference::{TextGeneration, TextGenerationOptionsBuilder};
 use tracing::{debug, instrument};
 use utoipa::ToSchema;
 
-use crate::search::CodeSearchService;
-
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
 #[schema(example=json!({
     "language": "python",
@@ -211,7 +209,7 @@ pub struct CompletionState {
 impl CompletionState {
     pub fn new(
         engine: Arc<Box<dyn TextGeneration>>,
-        code: Arc<CodeSearchService>,
+        code: Arc<BoxCodeSearch>,
         prompt_template: Option<String>,
     ) -> Self {
         Self {

crates/tabby/src/serve/completions/prompt.rs

@@ -4,7 +4,7 @@ use lazy_static::lazy_static;
 use regex::Regex;
 use strfmt::strfmt;
 use tabby_common::{
-    api::code::{CodeSearch, CodeSearchError},
+    code::{BoxCodeSearch, CodeSearch, CodeSearchError},
     index::CodeSearchSchema,
     languages::get_language,
 };
@@ -13,7 +13,6 @@ use textdistance::Algorithm;
 use tracing::warn;
 
 use super::{Segments, Snippet};
-use crate::search::CodeSearchService;
 
 static MAX_SNIPPETS_TO_FETCH: usize = 20;
 static MAX_SNIPPET_CHARS_IN_PROMPT: usize = 768;
@@ -22,11 +21,11 @@ static MAX_SIMILARITY_THRESHOLD: f32 = 0.9;
 pub struct PromptBuilder {
     schema: CodeSearchSchema,
     prompt_template: Option<String>,
-    code: Option<Arc<CodeSearchService>>,
+    code: Option<Arc<BoxCodeSearch>>,
 }
 
 impl PromptBuilder {
-    pub fn new(prompt_template: Option<String>, code: Option<Arc<CodeSearchService>>) -> Self {
+    pub fn new(prompt_template: Option<String>, code: Option<Arc<BoxCodeSearch>>) -> Self {
         PromptBuilder {
             schema: CodeSearchSchema::new(),
             prompt_template,
@@ -44,7 +43,13 @@ impl PromptBuilder {
 
     pub async fn collect(&self, language: &str, segments: &Segments) -> Vec<Snippet> {
         if let Some(code) = &self.code {
-            collect_snippets(&self.schema, code, language, &segments.prefix).await
+            collect_snippets(
+                &self.schema,
+                code.as_ref().as_ref(),
+                language,
+                &segments.prefix,
+            )
+            .await
         } else {
             vec![]
         }
@@ -113,7 +118,7 @@ fn build_prefix(language: &str, prefix: &str, snippets: &[Snippet]) -> String {
 
 async fn collect_snippets(
     schema: &CodeSearchSchema,
-    code: &CodeSearchService,
+    code: &dyn CodeSearch,
     language: &str,
     text: &str,
 ) -> Vec<Snippet> {

crates/tabby/src/serve/mod.rs

@@ -17,7 +17,7 @@ use axum::{routing, Router, Server};
 use axum_tracing_opentelemetry::opentelemetry_tracing_layer;
 use clap::Args;
 use tabby_common::{
-    api::code::{Hit, HitDocument, SearchResponse},
+    code::{create_local, Hit, HitDocument, SearchResponse},
     config::Config,
     usage,
 };
@@ -32,7 +32,7 @@ use self::{
     engine::{create_engine, EngineInfo},
     health::HealthState,
 };
-use crate::{chat::ChatService, fatal, search::CodeSearchService};
+use crate::{chat::ChatService, fatal};
 
 #[derive(OpenApi)]
 #[openapi(
@@ -173,7 +173,7 @@ pub async fn main(config: &Config, args: &ServeArgs) {
 }
 
 async fn api_router(args: &ServeArgs, config: &Config) -> Router {
-    let code = Arc::new(CodeSearchService::new());
+    let code = Arc::new(create_local());
     let completion_state = {
         let (
             engine,

crates/tabby/src/serve/search.rs

@@ -7,12 +7,10 @@ use axum::{
 };
 use hyper::StatusCode;
 use serde::Deserialize;
-use tabby_common::api::code::{CodeSearch, CodeSearchError, SearchResponse};
+use tabby_common::code::{BoxCodeSearch, CodeSearchError, SearchResponse};
 use tracing::{instrument, warn};
 use utoipa::IntoParams;
 
-use crate::search::CodeSearchService;
-
 #[derive(Deserialize, IntoParams)]
 pub struct SearchQuery {
     #[param(default = "get")]
@@ -38,7 +36,7 @@ pub struct SearchQuery {
     )]
 #[instrument(skip(state, query))]
 pub async fn search(
-    State(state): State<Arc<CodeSearchService>>,
+    State(state): State<Arc<BoxCodeSearch>>,
     query: Query<SearchQuery>,
 ) -> Result<Json<SearchResponse>, StatusCode> {
     match state