fix: add flag to disable retrieval augmented code completion and set … (#580)

* fix: add flag to disable retrieval augmented code completion and set it in golden test

* update
add-dagster-data-pipeline
Meng Zhang 2023-10-17 15:26:23 -07:00 committed by GitHub
parent 981133d6c8
commit 99d1bf34bb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 1 deletions

View File

@ -43,7 +43,16 @@ pub struct CompletionRequest {
#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
pub struct DebugOptions { pub struct DebugOptions {
/// When true, returns debug_data in completion response. /// When true, returns debug_data in completion response.
#[serde(default = "default_false")]
enabled: bool, enabled: bool,
/// When true, disable retrieval augmented code completion.
#[serde(default = "default_false")]
disable_retrieval_augmented_code_completion: bool,
}
fn default_false() -> bool {
false
} }
#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@ -126,7 +135,15 @@ pub async fn completions(
}; };
debug!("PREFIX: {}, SUFFIX: {:?}", segments.prefix, segments.suffix); debug!("PREFIX: {}, SUFFIX: {:?}", segments.prefix, segments.suffix);
let snippets = state.prompt_builder.collect(&language, &segments); let snippets = if !request
.debug_options
.as_ref()
.is_some_and(|x| x.disable_retrieval_augmented_code_completion)
{
state.prompt_builder.collect(&language, &segments)
} else {
vec![]
};
let prompt = state let prompt = state
.prompt_builder .prompt_builder
.build(&language, segments.clone(), &snippets); .build(&language, segments.clone(), &snippets);

View File

@ -3,6 +3,7 @@ use std::path::PathBuf;
use assert_json_diff::assert_json_include; use assert_json_diff::assert_json_include;
use lazy_static::lazy_static; use lazy_static::lazy_static;
use serde::Deserialize; use serde::Deserialize;
use serde_json::json;
use tokio::{ use tokio::{
process::Command, process::Command,
time::{sleep, Duration}, time::{sleep, Duration},
@ -68,6 +69,14 @@ async fn wait_for_server() {
} }
async fn golden_test(body: serde_json::Value, expected: serde_json::Value) { async fn golden_test(body: serde_json::Value, expected: serde_json::Value) {
let mut body = body.clone();
body.as_object_mut().unwrap().insert(
"debug_options".to_owned(),
json!({
"disable_retrieval_augmented_code_completion": true
}),
);
let actual: serde_json::Value = CLIENT let actual: serde_json::Value = CLIENT
.post("http://localhost:9090/v1/completions") .post("http://localhost:9090/v1/completions")
.json(&body) .json(&body)