oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add experimental mercury edit prediction provider (#44256)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Ben Kunkle <ben@zed.dev>
Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Agus Zubiaga
2025-12-06 10:08:44 +00:00
committed by GitHub
co-authored by Ben Kunkle Max Brunsfeld
parent 51b7d06a27
commit f08fd732a7
18 changed files with 808 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/edit_prediction/src/cursor_excerpt.rs
+78
View File
@@ -0,0 +1,78 @@
use language::{BufferSnapshot, Point};
use std::ops::Range;
pub fn editable_and_context_ranges_for_cursor_position(
position: Point,
snapshot: &BufferSnapshot,
editable_region_token_limit: usize,
context_token_limit: usize,
) -> (Range<Point>, Range<Point>) {
let mut scope_range = position..position;
let mut remaining_edit_tokens = editable_region_token_limit;
while let Some(parent) = snapshot.syntax_ancestor(scope_range.clone()) {
let parent_tokens = guess_token_count(parent.byte_range().len());
let parent_point_range = Point::new(
parent.start_position().row as u32,
parent.start_position().column as u32,
)
..Point::new(
parent.end_position().row as u32,
parent.end_position().column as u32,
);
if parent_point_range == scope_range {
break;
} else if parent_tokens <= editable_region_token_limit {
scope_range = parent_point_range;
remaining_edit_tokens = editable_region_token_limit - parent_tokens;
} else {
break;
}
}
let editable_range = expand_range(snapshot, scope_range, remaining_edit_tokens);
let context_range = expand_range(snapshot, editable_range.clone(), context_token_limit);
(editable_range, context_range)
}
fn expand_range(
snapshot: &BufferSnapshot,
range: Range<Point>,
mut remaining_tokens: usize,
) -> Range<Point> {
let mut expanded_range = range;
expanded_range.start.column = 0;
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
loop {
let mut expanded = false;
if remaining_tokens > 0 && expanded_range.start.row > 0 {
expanded_range.start.row -= 1;
let line_tokens =
guess_token_count(snapshot.line_len(expanded_range.start.row) as usize);
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
expanded = true;
}
if remaining_tokens > 0 && expanded_range.end.row < snapshot.max_point().row {
expanded_range.end.row += 1;
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
let line_tokens = guess_token_count(expanded_range.end.column as usize);
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
expanded = true;
}
if !expanded {
break;
}
}
expanded_range
}
/// Typical number of string bytes per token for the purposes of limiting model input. This is
/// intentionally low to err on the side of underestimating limits.
pub(crate) const BYTES_PER_TOKEN_GUESS: usize = 3;
pub fn guess_token_count(bytes: usize) -> usize {
bytes / BYTES_PER_TOKEN_GUESS
}
crates/edit_prediction/src/edit_prediction.rs
+35 -2
View File
@@ -51,8 +51,11 @@ use thiserror::Error;
use util::{RangeExt as _, ResultExt as _};
use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification};
mod cursor_excerpt;
mod license_detection;
pub mod mercury;
mod onboarding_modal;
pub mod open_ai_response;
mod prediction;
pub mod sweep_ai;
pub mod udiff;
@@ -65,6 +68,7 @@ pub mod zeta2;
mod edit_prediction_tests;
use crate::license_detection::LicenseDetectionWatcher;
use crate::mercury::Mercury;
use crate::onboarding_modal::ZedPredictModal;
pub use crate::prediction::EditPrediction;
pub use crate::prediction::EditPredictionId;
@@ -96,6 +100,12 @@ impl FeatureFlag for SweepFeatureFlag {
const NAME: &str = "sweep-ai";
}
pub struct MercuryFeatureFlag;
impl FeatureFlag for MercuryFeatureFlag {
const NAME: &str = "mercury";
}
pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions {
context: EditPredictionExcerptOptions {
max_bytes: 512,
@@ -157,6 +167,7 @@ pub struct EditPredictionStore {
eval_cache: Option<Arc<dyn EvalCache>>,
edit_prediction_model: EditPredictionModel,
pub sweep_ai: SweepAi,
pub mercury: Mercury,
data_collection_choice: DataCollectionChoice,
reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejection>,
shown_predictions: VecDeque<EditPrediction>,
@@ -169,6 +180,7 @@ pub enum EditPredictionModel {
Zeta1,
Zeta2,
Sweep,
Mercury,
}
#[derive(Debug, Clone, PartialEq)]
@@ -474,6 +486,7 @@ impl EditPredictionStore {
eval_cache: None,
edit_prediction_model: EditPredictionModel::Zeta2,
sweep_ai: SweepAi::new(cx),
mercury: Mercury::new(cx),
data_collection_choice,
reject_predictions_tx: reject_tx,
rated_predictions: Default::default(),
@@ -509,6 +522,15 @@ impl EditPredictionStore {
.is_some()
}
pub fn has_mercury_api_token(&self) -> bool {
self.mercury
.api_token
.clone()
.now_or_never()
.flatten()
.is_some()
}
#[cfg(feature = "eval-support")]
pub fn with_eval_cache(&mut self, cache: Arc<dyn EvalCache>) {
self.eval_cache = Some(cache);
@@ -868,7 +890,7 @@ impl EditPredictionStore {
fn accept_current_prediction(&mut self, project: &Entity<Project>, cx: &mut Context<Self>) {
match self.edit_prediction_model {
EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 => {}
EditPredictionModel::Sweep => return,
EditPredictionModel::Sweep | EditPredictionModel::Mercury => return,
}
let Some(project_state) = self.projects.get_mut(&project.entity_id()) else {
@@ -1013,7 +1035,7 @@ impl EditPredictionStore {
) {
match self.edit_prediction_model {
EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 => {}
EditPredictionModel::Sweep => return,
EditPredictionModel::Sweep | EditPredictionModel::Mercury => return,
}
self.reject_predictions_tx
@@ -1373,6 +1395,17 @@ impl EditPredictionStore {
diagnostic_search_range.clone(),
cx,
),
EditPredictionModel::Mercury => self.mercury.request_prediction(
&project,
&active_buffer,
snapshot.clone(),
position,
events,
&project_state.recent_paths,
related_files,
diagnostic_search_range.clone(),
cx,
),
};
cx.spawn(async move |this, cx| {
crates/edit_prediction/src/edit_prediction_tests.rs
+1 -1
View File
@@ -1620,7 +1620,7 @@ async fn test_no_data_collection_for_events_in_uncollectable_buffers(cx: &mut Te
buffer.edit(
[(
0..0,
" ".repeat(MAX_EVENT_TOKENS * zeta1::BYTES_PER_TOKEN_GUESS),
" ".repeat(MAX_EVENT_TOKENS * cursor_excerpt::BYTES_PER_TOKEN_GUESS),
)],
None,
cx,
crates/edit_prediction/src/mercury.rs
+340
View File
@@ -0,0 +1,340 @@
use anyhow::{Context as _, Result};
use cloud_llm_client::predict_edits_v3::Event;
use credentials_provider::CredentialsProvider;
use edit_prediction_context::RelatedFile;
use futures::{AsyncReadExt as _, FutureExt, future::Shared};
use gpui::{
App, AppContext as _, Entity, Task,
http_client::{self, AsyncBody, Method},
};
use language::{Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToPoint as _};
use project::{Project, ProjectPath};
use std::{
collections::VecDeque, fmt::Write as _, mem, ops::Range, path::Path, sync::Arc, time::Instant,
};
use crate::{
EditPredictionId, EditPredictionInputs, open_ai_response::text_from_response,
prediction::EditPredictionResult,
};
const MERCURY_API_URL: &str = "https://api.inceptionlabs.ai/v1/edit/completions";
const MAX_CONTEXT_TOKENS: usize = 150;
const MAX_REWRITE_TOKENS: usize = 350;
pub struct Mercury {
pub api_token: Shared<Task<Option<String>>>,
}
impl Mercury {
pub fn new(cx: &App) -> Self {
Mercury {
api_token: load_api_token(cx).shared(),
}
}
pub fn set_api_token(&mut self, api_token: Option<String>, cx: &mut App) -> Task<Result<()>> {
self.api_token = Task::ready(api_token.clone()).shared();
store_api_token_in_keychain(api_token, cx)
}
pub fn request_prediction(
&self,
_project: &Entity<Project>,
active_buffer: &Entity<Buffer>,
snapshot: BufferSnapshot,
position: language::Anchor,
events: Vec<Arc<Event>>,
_recent_paths: &VecDeque<ProjectPath>,
related_files: Vec<RelatedFile>,
_diagnostic_search_range: Range<Point>,
cx: &mut App,
) -> Task<Result<Option<EditPredictionResult>>> {
let Some(api_token) = self.api_token.clone().now_or_never().flatten() else {
return Task::ready(Ok(None));
};
let full_path: Arc<Path> = snapshot
.file()
.map(|file| file.full_path(cx))
.unwrap_or_else(|| "untitled".into())
.into();
let http_client = cx.http_client();
let cursor_point = position.to_point(&snapshot);
let buffer_snapshotted_at = Instant::now();
let result = cx.background_spawn(async move {
let (editable_range, context_range) =
crate::cursor_excerpt::editable_and_context_ranges_for_cursor_position(
cursor_point,
&snapshot,
MAX_CONTEXT_TOKENS,
MAX_REWRITE_TOKENS,
);
let offset_range = editable_range.to_offset(&snapshot);
let prompt = build_prompt(
&events,
&related_files,
&snapshot,
full_path.as_ref(),
cursor_point,
editable_range,
context_range.clone(),
);
let inputs = EditPredictionInputs {
events: events,
included_files: vec![cloud_llm_client::predict_edits_v3::RelatedFile {
path: full_path.clone(),
max_row: cloud_llm_client::predict_edits_v3::Line(snapshot.max_point().row),
excerpts: vec![cloud_llm_client::predict_edits_v3::Excerpt {
start_line: cloud_llm_client::predict_edits_v3::Line(
context_range.start.row,
),
text: snapshot
.text_for_range(context_range.clone())
.collect::<String>()
.into(),
}],
}],
cursor_point: cloud_llm_client::predict_edits_v3::Point {
column: cursor_point.column,
line: cloud_llm_client::predict_edits_v3::Line(cursor_point.row),
},
cursor_path: full_path.clone(),
};
let request_body = open_ai::Request {
model: "mercury-coder".into(),
messages: vec![open_ai::RequestMessage::User {
content: open_ai::MessageContent::Plain(prompt),
}],
stream: false,
max_completion_tokens: None,
stop: vec![],
temperature: None,
tool_choice: None,
parallel_tool_calls: None,
tools: vec![],
prompt_cache_key: None,
reasoning_effort: None,
};
let buf = serde_json::to_vec(&request_body)?;
let body: AsyncBody = buf.into();
let request = http_client::Request::builder()
.uri(MERCURY_API_URL)
.header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {}", api_token))
.header("Connection", "keep-alive")
.method(Method::POST)
.body(body)
.context("Failed to create request")?;
let mut response = http_client
.send(request)
.await
.context("Failed to send request")?;
let mut body: Vec<u8> = Vec::new();
response
.body_mut()
.read_to_end(&mut body)
.await
.context("Failed to read response body")?;
let response_received_at = Instant::now();
if !response.status().is_success() {
anyhow::bail!(
"Request failed with status: {:?}\nBody: {}",
response.status(),
String::from_utf8_lossy(&body),
);
};
let mut response: open_ai::Response =
serde_json::from_slice(&body).context("Failed to parse response")?;
let id = mem::take(&mut response.id);
let response_str = text_from_response(response).unwrap_or_default();
let response_str = response_str.strip_prefix("```\n").unwrap_or(&response_str);
let response_str = response_str.strip_suffix("\n```").unwrap_or(&response_str);
let mut edits = Vec::new();
const NO_PREDICTION_OUTPUT: &str = "None";
if response_str != NO_PREDICTION_OUTPUT {
let old_text = snapshot
.text_for_range(offset_range.clone())
.collect::<String>();
edits.extend(
language::text_diff(&old_text, &response_str)
.into_iter()
.map(|(range, text)| {
(
snapshot.anchor_after(offset_range.start + range.start)
..snapshot.anchor_before(offset_range.start + range.end),
text,
)
}),
);
}
anyhow::Ok((id, edits, snapshot, response_received_at, inputs))
});
let buffer = active_buffer.clone();
cx.spawn(async move |cx| {
let (id, edits, old_snapshot, response_received_at, inputs) =
result.await.context("Mercury edit prediction failed")?;
anyhow::Ok(Some(
EditPredictionResult::new(
EditPredictionId(id.into()),
&buffer,
&old_snapshot,
edits.into(),
buffer_snapshotted_at,
response_received_at,
inputs,
cx,
)
.await,
))
})
}
}
fn build_prompt(
events: &[Arc<Event>],
related_files: &[RelatedFile],
cursor_buffer: &BufferSnapshot,
cursor_buffer_path: &Path,
cursor_point: Point,
editable_range: Range<Point>,
context_range: Range<Point>,
) -> String {
const RECENTLY_VIEWED_SNIPPETS_START: &str = "<|recently_viewed_code_snippets|>\n";
const RECENTLY_VIEWED_SNIPPETS_END: &str = "<|/recently_viewed_code_snippets|>\n";
const RECENTLY_VIEWED_SNIPPET_START: &str = "<|recently_viewed_code_snippet|>\n";
const RECENTLY_VIEWED_SNIPPET_END: &str = "<|/recently_viewed_code_snippet|>\n";
const CURRENT_FILE_CONTENT_START: &str = "<|current_file_content|>\n";
const CURRENT_FILE_CONTENT_END: &str = "<|/current_file_content|>\n";
const CODE_TO_EDIT_START: &str = "<|code_to_edit|>\n";
const CODE_TO_EDIT_END: &str = "<|/code_to_edit|>\n";
const EDIT_DIFF_HISTORY_START: &str = "<|edit_diff_history|>\n";
const EDIT_DIFF_HISTORY_END: &str = "<|/edit_diff_history|>\n";
const CURSOR_TAG: &str = "<|cursor|>";
const CODE_SNIPPET_FILE_PATH_PREFIX: &str = "code_snippet_file_path: ";
const CURRENT_FILE_PATH_PREFIX: &str = "current_file_path: ";
let mut prompt = String::new();
push_delimited(
&mut prompt,
RECENTLY_VIEWED_SNIPPETS_START..RECENTLY_VIEWED_SNIPPETS_END,
|prompt| {
for related_file in related_files {
for related_excerpt in &related_file.excerpts {
push_delimited(
prompt,
RECENTLY_VIEWED_SNIPPET_START..RECENTLY_VIEWED_SNIPPET_END,
|prompt| {
prompt.push_str(CODE_SNIPPET_FILE_PATH_PREFIX);
prompt.push_str(related_file.path.path.as_unix_str());
prompt.push('\n');
prompt.push_str(&related_excerpt.text.to_string());
},
);
}
}
},
);
push_delimited(
&mut prompt,
CURRENT_FILE_CONTENT_START..CURRENT_FILE_CONTENT_END,
|prompt| {
prompt.push_str(CURRENT_FILE_PATH_PREFIX);
prompt.push_str(cursor_buffer_path.as_os_str().to_string_lossy().as_ref());
prompt.push('\n');
let prefix_range = context_range.start..editable_range.start;
let suffix_range = editable_range.end..context_range.end;
prompt.extend(cursor_buffer.text_for_range(prefix_range));
push_delimited(prompt, CODE_TO_EDIT_START..CODE_TO_EDIT_END, |prompt| {
let range_before_cursor = editable_range.start..cursor_point;
let range_after_cursor = cursor_point..editable_range.end;
prompt.extend(cursor_buffer.text_for_range(range_before_cursor));
prompt.push_str(CURSOR_TAG);
prompt.extend(cursor_buffer.text_for_range(range_after_cursor));
});
prompt.extend(cursor_buffer.text_for_range(suffix_range));
},
);
push_delimited(
&mut prompt,
EDIT_DIFF_HISTORY_START..EDIT_DIFF_HISTORY_END,
|prompt| {
for event in events {
writeln!(prompt, "{event}").unwrap();
}
},
);
prompt
}
fn push_delimited(prompt: &mut String, delimiters: Range<&str>, cb: impl FnOnce(&mut String)) {
prompt.push_str(delimiters.start);
cb(prompt);
prompt.push_str(delimiters.end);
}
pub const MERCURY_CREDENTIALS_URL: &str = "https://api.inceptionlabs.ai/v1/edit/completions";
pub const MERCURY_CREDENTIALS_USERNAME: &str = "mercury-api-token";
pub fn load_api_token(cx: &App) -> Task<Option<String>> {
if let Some(api_token) = std::env::var("MERCURY_AI_TOKEN")
.ok()
.filter(|value| !value.is_empty())
{
return Task::ready(Some(api_token));
}
let credentials_provider = <dyn CredentialsProvider>::global(cx);
cx.spawn(async move |cx| {
let (_, credentials) = credentials_provider
.read_credentials(MERCURY_CREDENTIALS_URL, &cx)
.await
.ok()??;
String::from_utf8(credentials).ok()
})
}
fn store_api_token_in_keychain(api_token: Option<String>, cx: &App) -> Task<Result<()>> {
let credentials_provider = <dyn CredentialsProvider>::global(cx);
cx.spawn(async move |cx| {
if let Some(api_token) = api_token {
credentials_provider
.write_credentials(
MERCURY_CREDENTIALS_URL,
MERCURY_CREDENTIALS_USERNAME,
api_token.as_bytes(),
cx,
)
.await
.context("Failed to save Mercury API token to system keychain")
} else {
credentials_provider
.delete_credentials(MERCURY_CREDENTIALS_URL, cx)
.await
.context("Failed to delete Mercury API token from system keychain")
}
})
}
crates/edit_prediction/src/open_ai_response.rs
+31
View File
@@ -0,0 +1,31 @@
pub fn text_from_response(mut res: open_ai::Response) -> Option<String> {
let choice = res.choices.pop()?;
let output_text = match choice.message {
open_ai::RequestMessage::Assistant {
content: Some(open_ai::MessageContent::Plain(content)),
..
} => content,
open_ai::RequestMessage::Assistant {
content: Some(open_ai::MessageContent::Multipart(mut content)),
..
} => {
if content.is_empty() {
log::error!("No output from Baseten completion response");
return None;
}
match content.remove(0) {
open_ai::MessagePart::Text { text } => text,
open_ai::MessagePart::Image { .. } => {
log::error!("Expected text, got an image");
return None;
}
}
}
_ => {
log::error!("Invalid response message: {:?}", choice.message);
return None;
}
};
Some(output_text)
}
crates/edit_prediction/src/zeta1.rs
+170 -8
View File
@@ -1,9 +1,8 @@
mod input_excerpt;
use std::{fmt::Write, ops::Range, path::Path, sync::Arc, time::Instant};
use crate::{
EditPredictionId, EditPredictionStore, ZedUpdateRequiredError,
cursor_excerpt::{editable_and_context_ranges_for_cursor_position, guess_token_count},
prediction::{EditPredictionInputs, EditPredictionResult},
};
use anyhow::{Context as _, Result};
@@ -12,7 +11,6 @@ use cloud_llm_client::{
predict_edits_v3::Event,
};
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task};
use input_excerpt::excerpt_for_cursor_position;
use language::{
Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToPoint as _, text_diff,
};
@@ -495,10 +493,174 @@ pub fn format_event(event: &Event) -> String {
}
}
/// Typical number of string bytes per token for the purposes of limiting model input. This is
/// intentionally low to err on the side of underestimating limits.
pub(crate) const BYTES_PER_TOKEN_GUESS: usize = 3;
#[derive(Debug)]
pub struct InputExcerpt {
pub context_range: Range<Point>,
pub editable_range: Range<Point>,
pub prompt: String,
}
fn guess_token_count(bytes: usize) -> usize {
bytes / BYTES_PER_TOKEN_GUESS
pub fn excerpt_for_cursor_position(
position: Point,
path: &str,
snapshot: &BufferSnapshot,
editable_region_token_limit: usize,
context_token_limit: usize,
) -> InputExcerpt {
let (editable_range, context_range) = editable_and_context_ranges_for_cursor_position(
position,
snapshot,
editable_region_token_limit,
context_token_limit,
);
let mut prompt = String::new();
writeln!(&mut prompt, "```{path}").unwrap();
if context_range.start == Point::zero() {
writeln!(&mut prompt, "{START_OF_FILE_MARKER}").unwrap();
}
for chunk in snapshot.chunks(context_range.start..editable_range.start, false) {
prompt.push_str(chunk.text);
}
push_editable_range(position, snapshot, editable_range.clone(), &mut prompt);
for chunk in snapshot.chunks(editable_range.end..context_range.end, false) {
prompt.push_str(chunk.text);
}
write!(prompt, "\n```").unwrap();
InputExcerpt {
context_range,
editable_range,
prompt,
}
}
fn push_editable_range(
cursor_position: Point,
snapshot: &BufferSnapshot,
editable_range: Range<Point>,
prompt: &mut String,
) {
writeln!(prompt, "{EDITABLE_REGION_START_MARKER}").unwrap();
for chunk in snapshot.chunks(editable_range.start..cursor_position, false) {
prompt.push_str(chunk.text);
}
prompt.push_str(CURSOR_MARKER);
for chunk in snapshot.chunks(cursor_position..editable_range.end, false) {
prompt.push_str(chunk.text);
}
write!(prompt, "\n{EDITABLE_REGION_END_MARKER}").unwrap();
}
#[cfg(test)]
mod tests {
use super::*;
use gpui::{App, AppContext};
use indoc::indoc;
use language::{Buffer, Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
use std::sync::Arc;
#[gpui::test]
fn test_excerpt_for_cursor_position(cx: &mut App) {
let text = indoc! {r#"
fn foo() {
let x = 42;
println!("Hello, world!");
}
fn bar() {
let x = 42;
let mut sum = 0;
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
return sum;
}
fn generate_random_numbers() -> Vec<i32> {
let mut rng = rand::thread_rng();
let mut numbers = Vec::new();
for _ in 0..5 {
numbers.push(rng.random_range(1..101));
}
numbers
}
"#};
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
let snapshot = buffer.read(cx).snapshot();
// Ensure we try to fit the largest possible syntax scope, resorting to line-based expansion
// when a larger scope doesn't fit the editable region.
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 50, 32);
assert_eq!(
excerpt.prompt,
indoc! {r#"
```main.rs
let x = 42;
println!("Hello, world!");
<|editable_region_start|>
}
fn bar() {
let x = 42;
let mut sum = 0;
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
r<|user_cursor_is_here|>eturn sum;
}
fn generate_random_numbers() -> Vec<i32> {
<|editable_region_end|>
let mut rng = rand::thread_rng();
let mut numbers = Vec::new();
```"#}
);
// The `bar` function won't fit within the editable region, so we resort to line-based expansion.
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 40, 32);
assert_eq!(
excerpt.prompt,
indoc! {r#"
```main.rs
fn bar() {
let x = 42;
let mut sum = 0;
<|editable_region_start|>
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
r<|user_cursor_is_here|>eturn sum;
}
fn generate_random_numbers() -> Vec<i32> {
let mut rng = rand::thread_rng();
<|editable_region_end|>
let mut numbers = Vec::new();
for _ in 0..5 {
numbers.push(rng.random_range(1..101));
```"#}
);
}
fn rust_lang() -> Language {
Language::new(
LanguageConfig {
name: "Rust".into(),
matcher: LanguageMatcher {
path_suffixes: vec!["rs".to_string()],
..Default::default()
},
..Default::default()
},
Some(tree_sitter_rust::LANGUAGE.into()),
)
}
}
crates/edit_prediction/src/zeta1/input_excerpt.rs
-231
View File
@@ -1,231 +0,0 @@
use super::{
CURSOR_MARKER, EDITABLE_REGION_END_MARKER, EDITABLE_REGION_START_MARKER, START_OF_FILE_MARKER,
guess_token_count,
};
use language::{BufferSnapshot, Point};
use std::{fmt::Write, ops::Range};
#[derive(Debug)]
pub struct InputExcerpt {
pub context_range: Range<Point>,
pub editable_range: Range<Point>,
pub prompt: String,
}
pub fn excerpt_for_cursor_position(
position: Point,
path: &str,
snapshot: &BufferSnapshot,
editable_region_token_limit: usize,
context_token_limit: usize,
) -> InputExcerpt {
let mut scope_range = position..position;
let mut remaining_edit_tokens = editable_region_token_limit;
while let Some(parent) = snapshot.syntax_ancestor(scope_range.clone()) {
let parent_tokens = guess_token_count(parent.byte_range().len());
let parent_point_range = Point::new(
parent.start_position().row as u32,
parent.start_position().column as u32,
)
..Point::new(
parent.end_position().row as u32,
parent.end_position().column as u32,
);
if parent_point_range == scope_range {
break;
} else if parent_tokens <= editable_region_token_limit {
scope_range = parent_point_range;
remaining_edit_tokens = editable_region_token_limit - parent_tokens;
} else {
break;
}
}
let editable_range = expand_range(snapshot, scope_range, remaining_edit_tokens);
let context_range = expand_range(snapshot, editable_range.clone(), context_token_limit);
let mut prompt = String::new();
writeln!(&mut prompt, "```{path}").unwrap();
if context_range.start == Point::zero() {
writeln!(&mut prompt, "{START_OF_FILE_MARKER}").unwrap();
}
for chunk in snapshot.chunks(context_range.start..editable_range.start, false) {
prompt.push_str(chunk.text);
}
push_editable_range(position, snapshot, editable_range.clone(), &mut prompt);
for chunk in snapshot.chunks(editable_range.end..context_range.end, false) {
prompt.push_str(chunk.text);
}
write!(prompt, "\n```").unwrap();
InputExcerpt {
context_range,
editable_range,
prompt,
}
}
fn push_editable_range(
cursor_position: Point,
snapshot: &BufferSnapshot,
editable_range: Range<Point>,
prompt: &mut String,
) {
writeln!(prompt, "{EDITABLE_REGION_START_MARKER}").unwrap();
for chunk in snapshot.chunks(editable_range.start..cursor_position, false) {
prompt.push_str(chunk.text);
}
prompt.push_str(CURSOR_MARKER);
for chunk in snapshot.chunks(cursor_position..editable_range.end, false) {
prompt.push_str(chunk.text);
}
write!(prompt, "\n{EDITABLE_REGION_END_MARKER}").unwrap();
}
fn expand_range(
snapshot: &BufferSnapshot,
range: Range<Point>,
mut remaining_tokens: usize,
) -> Range<Point> {
let mut expanded_range = range;
expanded_range.start.column = 0;
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
loop {
let mut expanded = false;
if remaining_tokens > 0 && expanded_range.start.row > 0 {
expanded_range.start.row -= 1;
let line_tokens =
guess_token_count(snapshot.line_len(expanded_range.start.row) as usize);
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
expanded = true;
}
if remaining_tokens > 0 && expanded_range.end.row < snapshot.max_point().row {
expanded_range.end.row += 1;
expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
let line_tokens = guess_token_count(expanded_range.end.column as usize);
remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
expanded = true;
}
if !expanded {
break;
}
}
expanded_range
}
#[cfg(test)]
mod tests {
use super::*;
use gpui::{App, AppContext};
use indoc::indoc;
use language::{Buffer, Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
use std::sync::Arc;
#[gpui::test]
fn test_excerpt_for_cursor_position(cx: &mut App) {
let text = indoc! {r#"
fn foo() {
let x = 42;
println!("Hello, world!");
}
fn bar() {
let x = 42;
let mut sum = 0;
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
return sum;
}
fn generate_random_numbers() -> Vec<i32> {
let mut rng = rand::thread_rng();
let mut numbers = Vec::new();
for _ in 0..5 {
numbers.push(rng.random_range(1..101));
}
numbers
}
"#};
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
let snapshot = buffer.read(cx).snapshot();
// Ensure we try to fit the largest possible syntax scope, resorting to line-based expansion
// when a larger scope doesn't fit the editable region.
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 50, 32);
assert_eq!(
excerpt.prompt,
indoc! {r#"
```main.rs
let x = 42;
println!("Hello, world!");
<|editable_region_start|>
}
fn bar() {
let x = 42;
let mut sum = 0;
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
r<|user_cursor_is_here|>eturn sum;
}
fn generate_random_numbers() -> Vec<i32> {
<|editable_region_end|>
let mut rng = rand::thread_rng();
let mut numbers = Vec::new();
```"#}
);
// The `bar` function won't fit within the editable region, so we resort to line-based expansion.
let excerpt = excerpt_for_cursor_position(Point::new(12, 5), "main.rs", &snapshot, 40, 32);
assert_eq!(
excerpt.prompt,
indoc! {r#"
```main.rs
fn bar() {
let x = 42;
let mut sum = 0;
<|editable_region_start|>
for i in 0..x {
sum += i;
}
println!("Sum: {}", sum);
r<|user_cursor_is_here|>eturn sum;
}
fn generate_random_numbers() -> Vec<i32> {
let mut rng = rand::thread_rng();
<|editable_region_end|>
let mut numbers = Vec::new();
for _ in 0..5 {
numbers.push(rng.random_range(1..101));
```"#}
);
}
fn rust_lang() -> Language {
Language::new(
LanguageConfig {
name: "Rust".into(),
matcher: LanguageMatcher {
path_suffixes: vec!["rs".to_string()],
..Default::default()
},
..Default::default()
},
Some(tree_sitter_rust::LANGUAGE.into()),
)
}
}
crates/edit_prediction/src/zeta2.rs
+2 -33
View File
@@ -1,5 +1,6 @@
#[cfg(feature = "eval-support")]
use crate::EvalCacheEntryKind;
use crate::open_ai_response::text_from_response;
use crate::prediction::EditPredictionResult;
use crate::{
DebugEvent, EDIT_PREDICTIONS_MODEL_ID, EditPredictionId, EditPredictionInputs,
@@ -199,7 +200,7 @@ pub fn request_prediction_with_zeta2(
stream: false,
max_completion_tokens: None,
stop: generation_params.stop.unwrap_or_default(),
temperature: generation_params.temperature.unwrap_or(0.7),
temperature: generation_params.temperature.or(Some(0.7)),
tool_choice: None,
parallel_tool_calls: None,
tools: vec![],
@@ -324,35 +325,3 @@ pub fn request_prediction_with_zeta2(
))
})
}
pub fn text_from_response(mut res: open_ai::Response) -> Option<String> {
let choice = res.choices.pop()?;
let output_text = match choice.message {
open_ai::RequestMessage::Assistant {
content: Some(open_ai::MessageContent::Plain(content)),
..
} => content,
open_ai::RequestMessage::Assistant {
content: Some(open_ai::MessageContent::Multipart(mut content)),
..
} => {
if content.is_empty() {
log::error!("No output from Baseten completion response");
return None;
}
match content.remove(0) {
open_ai::MessagePart::Text { text } => text,
open_ai::MessagePart::Image { .. } => {
log::error!("Expected text, got an image");
return None;
}
}
}
_ => {
log::error!("Invalid response message: {:?}", choice.message);
return None;
}
};
Some(output_text)
}