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)]
|
||||
pub struct DebugOptions {
|
||||
/// When true, returns debug_data in completion response.
|
||||
#[serde(default = "default_false")]
|
||||
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)]
|
||||
|
|
@ -126,7 +135,15 @@ pub async fn completions(
|
|||
};
|
||||
|
||||
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
|
||||
.prompt_builder
|
||||
.build(&language, segments.clone(), &snippets);
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ use std::path::PathBuf;
|
|||
use assert_json_diff::assert_json_include;
|
||||
use lazy_static::lazy_static;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
use tokio::{
|
||||
process::Command,
|
||||
time::{sleep, Duration},
|
||||
|
|
@ -68,6 +69,14 @@ async fn wait_for_server() {
|
|||
}
|
||||
|
||||
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
|
||||
.post("http://localhost:9090/v1/completions")
|
||||
.json(&body)
|
||||
|
|
|
|||
Loading…
Reference in New Issue