1use crate::{PredictionProvider, PromptFormat};
2use anyhow::{Context as _, Result};
3use collections::HashMap;
4use edit_prediction::example_spec::ExampleSpec;
5use edit_prediction::udiff::OpenedBuffers;
6use gpui::Entity;
7use http_client::Url;
8use language::{Anchor, Buffer};
9use project::Project;
10use serde::{Deserialize, Serialize};
11use std::sync::Arc;
12use std::{
13 borrow::Cow,
14 io::{Read, Write},
15 path::{Path, PathBuf},
16};
17use zeta_prompt::RelatedFile;
18
19#[derive(Clone, Debug, Serialize, Deserialize)]
20pub struct Example {
21 #[serde(flatten)]
22 pub spec: ExampleSpec,
23
24 /// The full content of the file where an edit is being predicted, and the
25 /// actual cursor offset.
26 #[serde(skip_serializing_if = "Option::is_none")]
27 pub buffer: Option<ExampleBuffer>,
28
29 /// The context retrieved for the prediction. This requires the worktree to
30 /// be loaded and the language server to be started.
31 #[serde(skip_serializing_if = "Option::is_none")]
32 pub context: Option<ExampleContext>,
33
34 /// The input and expected output from the edit prediction model.
35 #[serde(skip_serializing_if = "Option::is_none")]
36 pub prompt: Option<ExamplePrompt>,
37
38 /// The actual predictions from the model.
39 #[serde(default, skip_serializing_if = "Vec::is_empty")]
40 pub predictions: Vec<ExamplePrediction>,
41
42 /// The scores, for how well the actual predictions match the expected
43 /// predictions.
44 #[serde(default, skip_serializing_if = "Vec::is_empty")]
45 pub score: Vec<ExampleScore>,
46
47 /// The application state used to process this example.
48 #[serde(skip)]
49 pub state: Option<ExampleState>,
50}
51
52#[derive(Clone, Debug)]
53pub struct ExampleState {
54 pub project: Entity<Project>,
55 pub buffer: Entity<Buffer>,
56 pub cursor_position: Anchor,
57 pub _open_buffers: OpenedBuffers,
58}
59
60#[derive(Clone, Debug, Serialize, Deserialize)]
61pub struct ExampleContext {
62 pub files: Arc<[RelatedFile]>,
63}
64
65#[derive(Clone, Debug, Serialize, Deserialize)]
66pub struct ExampleBuffer {
67 pub content: String,
68 pub cursor_row: u32,
69 pub cursor_column: u32,
70 pub cursor_offset: usize,
71}
72
73#[derive(Clone, Debug, Serialize, Deserialize)]
74pub struct ExamplePrompt {
75 pub input: String,
76 pub expected_output: String,
77 pub format: PromptFormat,
78}
79
80#[derive(Clone, Debug, Serialize, Deserialize)]
81pub struct ExamplePrediction {
82 pub actual_patch: String,
83 pub actual_output: String,
84 pub provider: PredictionProvider,
85}
86
87#[derive(Clone, Debug, Serialize, Deserialize)]
88pub struct ExampleScore {
89 pub delta_chr_f: f32,
90}
91
92impl Example {
93 pub fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
94 // git@github.com:owner/repo.git
95 if self.spec.repository_url.contains('@') {
96 let (owner, repo) = self
97 .spec
98 .repository_url
99 .split_once(':')
100 .context("expected : in git url")?
101 .1
102 .split_once('/')
103 .context("expected / in git url")?;
104 Ok((
105 Cow::Borrowed(owner),
106 Cow::Borrowed(repo.trim_end_matches(".git")),
107 ))
108 // http://github.com/owner/repo.git
109 } else {
110 let url = Url::parse(&self.spec.repository_url)?;
111 let mut segments = url.path_segments().context("empty http url")?;
112 let owner = segments
113 .next()
114 .context("expected owner path segment")?
115 .to_string();
116 let repo = segments
117 .next()
118 .context("expected repo path segment")?
119 .trim_end_matches(".git")
120 .to_string();
121 assert!(segments.next().is_none());
122
123 Ok((owner.into(), repo.into()))
124 }
125 }
126}
127
128pub fn read_example_files(inputs: &[PathBuf]) -> Vec<Example> {
129 let mut examples = Vec::new();
130
131 for path in inputs {
132 let is_stdin = path.as_path() == Path::new("-");
133 let content = if is_stdin {
134 let mut buffer = String::new();
135 std::io::stdin()
136 .read_to_string(&mut buffer)
137 .expect("Failed to read from stdin");
138 buffer
139 } else {
140 std::fs::read_to_string(path)
141 .unwrap_or_else(|_| panic!("Failed to read path: {:?}", &path))
142 };
143 let filename = path.file_stem().unwrap().to_string_lossy().to_string();
144 let ext = if !is_stdin {
145 path.extension()
146 .map(|ext| ext.to_string_lossy().to_string())
147 .unwrap_or_else(|| panic!("{} should have an extension", path.display()))
148 } else {
149 "jsonl".to_string()
150 };
151
152 match ext.as_ref() {
153 "json" => {
154 let mut example =
155 serde_json::from_str::<Example>(&content).unwrap_or_else(|error| {
156 panic!("Failed to parse example file: {}\n{error}", path.display())
157 });
158 if example.spec.name.is_empty() {
159 example.spec.name = filename;
160 }
161 examples.push(example);
162 }
163 "jsonl" => examples.extend(
164 content
165 .lines()
166 .enumerate()
167 .map(|(line_ix, line)| {
168 let mut example =
169 serde_json::from_str::<Example>(line).unwrap_or_else(|error| {
170 panic!(
171 "Failed to parse example on {}:{}\n{error}",
172 path.display(),
173 line_ix + 1
174 )
175 });
176 if example.spec.name.is_empty() {
177 example.spec.name = format!("{filename}-{line_ix}")
178 }
179 example
180 })
181 .collect::<Vec<Example>>(),
182 ),
183 "md" => {
184 let mut example = parse_markdown_example(&content).unwrap();
185 if example.spec.name.is_empty() {
186 example.spec.name = filename;
187 }
188 examples.push(example);
189 }
190 ext => {
191 panic!("{} has invalid example extension `{ext}`", path.display())
192 }
193 }
194 }
195
196 examples
197}
198
199pub fn write_examples(examples: &[Example], output_path: Option<&PathBuf>) {
200 let mut content = String::new();
201 for example in examples {
202 let line = serde_json::to_string(example).unwrap();
203 content.push_str(&line);
204 content.push('\n');
205 }
206 if let Some(output_path) = output_path {
207 std::fs::write(output_path, content).expect("Failed to write examples");
208 } else {
209 std::io::stdout().write_all(&content.as_bytes()).unwrap();
210 }
211}
212
213pub fn sort_examples_by_repo_and_rev(examples: &mut [Example]) {
214 examples.sort_by(|a, b| {
215 a.spec
216 .repository_url
217 .cmp(&b.spec.repository_url)
218 .then(b.spec.revision.cmp(&a.spec.revision))
219 });
220}
221
222pub fn group_examples_by_repo(examples: &mut [Example]) -> Vec<Vec<&mut Example>> {
223 let mut examples_by_repo = HashMap::default();
224 for example in examples.iter_mut() {
225 examples_by_repo
226 .entry(example.spec.repository_url.clone())
227 .or_insert_with(Vec::new)
228 .push(example);
229 }
230 examples_by_repo.into_values().collect()
231}
232
233fn parse_markdown_example(input: &str) -> Result<Example> {
234 let spec = ExampleSpec::from_markdown(input)?;
235 Ok(Example {
236 spec,
237 buffer: None,
238 context: None,
239 prompt: None,
240 predictions: Vec::new(),
241 score: Vec::new(),
242 state: None,
243 })
244}