feat: add debug request / response to visualize prompting with source…

… code index (TabbyML#544)

* feat: logs segments in completion log

* feat: tune prompt format and improve testing

* add debug options for easier of visualizing the prompt

* update

Cargo.toml

@@ -36,3 +36,4 @@ derive_builder = "0.12.0"
 tokenizers = "0.13.4-rc3"
 futures = "0.3.28"
 async-stream = "0.3.5"
+regex = "1.10.0"

crates/tabby-common/src/events.rs

@@ -56,10 +56,16 @@ pub enum Event<'a> {
         completion_id: &'a str,
         language: &'a str,
         prompt: &'a str,
+        segments: &'a Segments<'a>,
         choices: Vec<Choice<'a>>,
         user: Option<&'a str>,
     },
 }
+#[derive(Serialize)]
+pub struct Segments<'a> {
+    pub prefix: &'a str,
+    pub suffix: Option<&'a str>,
+}
 
 #[derive(Serialize)]
 struct Log<'a> {

crates/tabby-inference/Cargo.toml

@@ -11,5 +11,5 @@ async-trait = { workspace = true }
 dashmap = "5.5.3"
 derive_builder = "0.12.0"
 futures = { workspace = true }
-regex = "1.9.5"
+regex.workspace = true
 tokenizers.workspace = true

crates/tabby/Cargo.toml

@@ -40,6 +40,8 @@ futures = { workspace = true }
 async-stream = { workspace = true }
 axum-streams = { version = "0.9.1", features = ["json"] }
 minijinja = { version = "1.0.8", features = ["loader"] }
+textdistance = "1.0.2"
+regex.workspace = true
 
 [target.'cfg(all(target_os="macos", target_arch="aarch64"))'.dependencies]
 llama-cpp-bindings = { path = "../llama-cpp-bindings" }

crates/tabby/src/serve/completions.rs

@@ -23,6 +23,7 @@ use super::search::IndexServer;
     }
 }))]
 pub struct CompletionRequest {
+    #[deprecated]
     #[schema(example = "def fib(n):")]
     prompt: Option<String>,
 
@@ -37,6 +38,14 @@ pub struct CompletionRequest {
     // A unique identifier representing your end-user, which can help Tabby to monitor & generating
     // reports.
     user: Option<String>,
+
+    debug: Option<DebugRequest>,
+}
+
+#[derive(Serialize, ToSchema, Deserialize, Clone, Debug)]
+pub struct DebugRequest {
+    // When true, returns debug_data in completion response.
+    enabled: bool,
 }
 
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@@ -55,9 +64,31 @@ pub struct Choice {
 }
 
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
+pub struct Snippet {
+    filepath: String,
+    body: String,
+    score: f32,
+}
+
+#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
+#[schema(example=json!({
+    "id": "string",
+    "choices": [ { "index": 0, "text": "string" } ]
+}))]
 pub struct CompletionResponse {
     id: String,
     choices: Vec<Choice>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    debug_data: Option<DebugData>,
+}
+
+#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
+pub struct DebugData {
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    snippets: Vec<Snippet>,
+
+    prompt: String,
 }
 
 #[utoipa::path(
@@ -97,7 +128,10 @@ pub async fn completions(
     };
 
     debug!("PREFIX: {}, SUFFIX: {:?}", segments.prefix, segments.suffix);
-    let prompt = state.prompt_builder.build(&language, segments);
+    let snippets = state.prompt_builder.collect(&language, &segments);
+    let prompt = state
+        .prompt_builder
+        .build(&language, segments.clone(), &snippets);
     debug!("PROMPT: {}", prompt);
     let completion_id = format!("cmpl-{}", uuid::Uuid::new_v4());
     let text = state.engine.generate(&prompt, options).await;
@@ -106,6 +140,10 @@ pub async fn completions(
         completion_id: &completion_id,
         language: &language,
         prompt: &prompt,
+        segments: &tabby_common::events::Segments {
+            prefix: &segments.prefix,
+            suffix: segments.suffix.as_deref(),
+        },
         choices: vec![events::Choice {
             index: 0,
             text: &text,
@@ -114,9 +152,16 @@ pub async fn completions(
     }
     .log();
 
+    let debug_data = DebugData { snippets, prompt };
+
     Ok(Json(CompletionResponse {
         id: completion_id,
         choices: vec![Choice { index: 0, text }],
+        debug_data: if request.debug.is_some_and(|x| x.enabled) {
+            Some(debug_data)
+        } else {
+            None
+        },
     }))
 }