diff --git a/Cargo.lock b/Cargo.lock index 21de5d3..ea28ac6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4125,6 +4125,7 @@ dependencies = [ "tantivy", "thiserror", "tokio", + "tracing", "utoipa", "uuid 1.4.1", ] diff --git a/crates/tabby-common/Cargo.toml b/crates/tabby-common/Cargo.toml index 238ef9d..2de7ad5 100644 --- a/crates/tabby-common/Cargo.toml +++ b/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 = [] diff --git a/crates/tabby-common/src/api/mod.rs b/crates/tabby-common/src/api/mod.rs deleted file mode 100644 index 9de50d4..0000000 --- a/crates/tabby-common/src/api/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod code; diff --git a/crates/tabby-common/src/api/code.rs b/crates/tabby-common/src/code/api.rs similarity index 96% rename from crates/tabby-common/src/api/code.rs rename to crates/tabby-common/src/code/api.rs index 2491a99..13464ed 100644 --- a/crates/tabby-common/src/api/code.rs +++ b/crates/tabby-common/src/code/api.rs @@ -54,3 +54,5 @@ pub trait CodeSearch: Send + Sync { offset: usize, ) -> Result; } + +pub type BoxCodeSearch = Box; diff --git a/crates/tabby/src/search.rs b/crates/tabby-common/src/code/imp.rs similarity index 96% rename from crates/tabby/src/search.rs rename to crates/tabby-common/src/code/imp.rs index ed123e5..1ecccea 100644 --- a/crates/tabby/src/search.rs +++ b/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>>, } impl CodeSearchService { - pub fn new() -> Self { + pub(crate) fn new() -> Self { let search = Arc::new(Mutex::new(None)); let ret = Self { diff --git a/crates/tabby-common/src/code/mod.rs b/crates/tabby-common/src/code/mod.rs new file mode 100644 index 0000000..bdd24e3 --- /dev/null +++ b/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()) +} diff --git a/crates/tabby-common/src/lib.rs b/crates/tabby-common/src/lib.rs index cc548cc..6cb2a2e 100644 --- a/crates/tabby-common/src/lib.rs +++ b/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; diff --git a/crates/tabby/src/main.rs b/crates/tabby/src/main.rs index e12c979..a3b5859 100644 --- a/crates/tabby/src/main.rs +++ b/crates/tabby/src/main.rs @@ -1,6 +1,5 @@ mod chat; mod download; -mod search; mod serve; use clap::{Parser, Subcommand}; diff --git a/crates/tabby/src/serve/completions.rs b/crates/tabby/src/serve/completions.rs index 1032ee4..89a83ac 100644 --- a/crates/tabby/src/serve/completions.rs +++ b/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>, - code: Arc, + code: Arc, prompt_template: Option, ) -> Self { Self { diff --git a/crates/tabby/src/serve/completions/prompt.rs b/crates/tabby/src/serve/completions/prompt.rs index a9a2f5a..48a7d73 100644 --- a/crates/tabby/src/serve/completions/prompt.rs +++ b/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, - code: Option>, + code: Option>, } impl PromptBuilder { - pub fn new(prompt_template: Option, code: Option>) -> Self { + pub fn new(prompt_template: Option, code: Option>) -> Self { PromptBuilder { schema: CodeSearchSchema::new(), prompt_template, @@ -44,7 +43,13 @@ impl PromptBuilder { pub async fn collect(&self, language: &str, segments: &Segments) -> Vec { 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 { diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs index 62c5ea3..64fde30 100644 --- a/crates/tabby/src/serve/mod.rs +++ b/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, diff --git a/crates/tabby/src/serve/search.rs b/crates/tabby/src/serve/search.rs index bdc742f..78805c7 100644 --- a/crates/tabby/src/serve/search.rs +++ b/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>, + State(state): State>, query: Query, ) -> Result, StatusCode> { match state