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