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 * updateadd-dagster-data-pipeline
parent
981133d6c8
commit
99d1bf34bb
|
|
@ -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);
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue