tenant.openagents/omega

1use crate::PredictionProvider;
2use crate::paths::WORKTREES_DIR;
3use crate::qa::QaResult;
4use anyhow::{Context as _, Result};
5use collections::HashMap;
6use edit_prediction::example_spec::ExampleSpec;
7use edit_prediction::udiff::OpenedBuffers;
8use gpui::Entity;
9use http_client::Url;
10use language::{Anchor, Buffer};
11use project::Project;
12use serde::{Deserialize, Serialize};
13use std::{
14    borrow::Cow,
15    collections::VecDeque,
16    io::Read,
17    path::{Path, PathBuf},
18};
19use zeta_prompt::Zeta2PromptInput;
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<Zeta2PromptInput>,
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    /// QA evaluation results for each prediction (indexed parallel to `predictions`).
45    #[serde(default, skip_serializing_if = "Vec::is_empty")]
46    pub qa: Vec<Option<QaResult>>,
47
48    /// The Zed version used to generate this example.
49    pub zed_version: Option<String>,
50
51    /// The application state used to process this example.
52    #[serde(skip)]
53    pub state: Option<ExampleState>,
54}
55
56#[derive(Clone, Debug)]
57pub struct ExampleState {
58    pub project: Entity<Project>,
59    pub buffer: Entity<Buffer>,
60    pub cursor_position: Anchor,
61    pub _open_buffers: OpenedBuffers,
62}
63
64#[derive(Clone, Debug, Serialize, Deserialize)]
65pub struct ExamplePrompt {
66    pub input: String,
67    #[serde(default)]
68    pub expected_output: Option<String>,
69    pub rejected_output: Option<String>, // For DPO
70    #[serde(default)]
71    pub prefill: Option<String>,
72    pub provider: PredictionProvider,
73}
74
75#[derive(Clone, Debug, Serialize, Deserialize)]
76pub struct ExamplePrediction {
77    #[serde(default, skip_serializing_if = "Option::is_none")]
78    pub actual_patch: Option<String>,
79    #[serde(deserialize_with = "deserialize_null_as_empty_string")]
80    pub actual_output: String,
81    #[serde(default, skip_serializing_if = "Option::is_none")]
82    pub actual_cursor: Option<ActualCursor>,
83    #[serde(default, skip_serializing_if = "Option::is_none")]
84    pub error: Option<String>,
85    pub provider: PredictionProvider,
86    #[serde(default, skip_serializing_if = "Option::is_none")]
87    pub cumulative_logprob: Option<f64>,
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub avg_logprob: Option<f64>,
90}
91
92#[derive(Clone, Debug, Serialize, Deserialize)]
93pub struct ActualCursor {
94    pub path: String,
95    pub row: u32,
96    pub column: u32,
97    pub offset: usize,
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub editable_region_offset: Option<usize>,
100}
101
102impl ActualCursor {
103    /// Construct an `ActualCursor` from a cursor offset within the new editable region.
104    ///
105    /// - `path`: file path the cursor is in
106    /// - `editable_region_cursor_offset`: byte offset of the cursor within the new editable region text
107    /// - `new_editable_region`: the full new editable region text (after marker removal)
108    /// - `content`: the full file content (before the edit)
109    /// - `editable_region_byte_offset`: byte offset where the editable region starts in `content`
110    /// - `editable_region_start_line`: 0-based line number where the editable region starts in `content`
111    pub fn from_editable_region(
112        path: &std::path::Path,
113        editable_region_cursor_offset: usize,
114        new_editable_region: &str,
115        content: &str,
116        editable_region_byte_offset: usize,
117        editable_region_start_line: usize,
118    ) -> Self {
119        // Defensive: a malformed/edge-case cursor offset must never panic and
120        // abort an (expensive) batch run. Clamp into range and snap down to a
121        // char boundary before slicing.
122        let mut editable_region_cursor_offset =
123            editable_region_cursor_offset.min(new_editable_region.len());
124        while editable_region_cursor_offset > 0
125            && !new_editable_region.is_char_boundary(editable_region_cursor_offset)
126        {
127            editable_region_cursor_offset -= 1;
128        }
129        let global_offset = editable_region_byte_offset + editable_region_cursor_offset;
130        let new_region_prefix = &new_editable_region[..editable_region_cursor_offset];
131        let row = (editable_region_start_line + new_region_prefix.matches('\n').count()) as u32;
132        let column = match new_region_prefix.rfind('\n') {
133            Some(pos) => (editable_region_cursor_offset - pos - 1) as u32,
134            None => {
135                let content_prefix = &content[..editable_region_byte_offset];
136                let content_column = match content_prefix.rfind('\n') {
137                    Some(pos) => editable_region_byte_offset - pos - 1,
138                    None => editable_region_byte_offset,
139                };
140                (content_column + editable_region_cursor_offset) as u32
141            }
142        };
143        ActualCursor {
144            path: path.to_string_lossy().to_string(),
145            row,
146            column,
147            offset: global_offset,
148            editable_region_offset: Some(editable_region_cursor_offset),
149        }
150    }
151}
152
153fn deserialize_null_as_empty_string<'de, D>(deserializer: D) -> Result<String, D::Error>
154where
155    D: serde::Deserializer<'de>,
156{
157    let opt = Option::<String>::deserialize(deserializer)?;
158    Ok(opt.unwrap_or_default())
159}
160
161pub type ExampleScore = edit_prediction_metrics::PredictionScore;
162
163impl Example {
164    pub fn repo_name(&self) -> Result<RepoName<'_>> {
165        // git@github.com:owner/repo.git
166        if self.spec.repository_url.contains('@') {
167            let (owner, repo) = self
168                .spec
169                .repository_url
170                .split_once(':')
171                .context("expected : in git url")?
172                .1
173                .split_once('/')
174                .context("expected / in git url")?;
175            Ok(RepoName {
176                owner: Cow::Borrowed(owner),
177                name: Cow::Borrowed(repo.trim_end_matches(".git")),
178            })
179        // http://github.com/owner/repo.git
180        } else {
181            let url = Url::parse(&self.spec.repository_url)?;
182            let mut segments = url.path_segments().context("empty http url")?;
183            let owner = segments
184                .next()
185                .context("expected owner path segment")?
186                .to_string();
187            let repo = segments
188                .next()
189                .context("expected repo path segment")?
190                .trim_end_matches(".git")
191                .to_string();
192            assert!(segments.next().is_none());
193
194            Ok(RepoName {
195                owner: Cow::Owned(owner),
196                name: Cow::Owned(repo),
197            })
198        }
199    }
200}
201
202pub struct RepoName<'a> {
203    pub owner: Cow<'a, str>,
204    pub name: Cow<'a, str>,
205}
206
207impl RepoName<'_> {
208    pub fn worktree_path(&self) -> PathBuf {
209        WORKTREES_DIR
210            .join(self.owner.as_ref())
211            .join(self.name.as_ref())
212    }
213}
214
215pub fn read_example_files(inputs: &[PathBuf]) -> Vec<Example> {
216    let mut examples = Vec::new();
217
218    for path in inputs {
219        let is_stdin = path.as_path() == Path::new("-");
220        let content = if is_stdin {
221            let mut buffer = String::new();
222            std::io::stdin()
223                .read_to_string(&mut buffer)
224                .expect("Failed to read from stdin");
225            buffer
226        } else {
227            std::fs::read_to_string(path)
228                .unwrap_or_else(|_| panic!("Failed to read path: {:?}", &path))
229        };
230        let filename = path.file_stem().unwrap().to_string_lossy().to_string();
231        let ext = if !is_stdin {
232            path.extension()
233                .map(|ext| ext.to_string_lossy().to_string())
234                .unwrap_or_else(|| panic!("{} should have an extension", path.display()))
235        } else {
236            "jsonl".to_string()
237        };
238
239        match ext.as_ref() {
240            "json" => {
241                let mut example =
242                    serde_json::from_str::<Example>(&content).unwrap_or_else(|error| {
243                        panic!("Failed to parse example file: {}\n{error}", path.display())
244                    });
245                if example.spec.name.is_empty() {
246                    example.spec.name = filename;
247                }
248                examples.push(example);
249            }
250            "jsonl" => examples.extend(
251                content
252                    .lines()
253                    .enumerate()
254                    .map(|(line_ix, line)| {
255                        let mut example =
256                            serde_json::from_str::<Example>(line).unwrap_or_else(|error| {
257                                panic!(
258                                    "Failed to parse example on {}:{}\n{error}",
259                                    path.display(),
260                                    line_ix + 1
261                                )
262                            });
263                        if example.spec.name.is_empty() {
264                            example.spec.name = format!("{filename}-{line_ix}")
265                        }
266                        example
267                    })
268                    .collect::<Vec<Example>>(),
269            ),
270            "md" => {
271                let mut example = parse_markdown_example(&content).unwrap();
272                if example.spec.name.is_empty() {
273                    example.spec.name = filename;
274                }
275                examples.push(example);
276            }
277            ext => {
278                panic!("{} has invalid example extension `{ext}`", path.display())
279            }
280        }
281    }
282
283    examples
284}
285
286pub fn sort_examples_by_repo_and_rev(examples: &mut [Example]) {
287    examples.sort_by(|a, b| {
288        a.spec
289            .repository_url
290            .cmp(&b.spec.repository_url)
291            .then(b.spec.revision.cmp(&a.spec.revision))
292    });
293}
294
295pub fn group_examples_by_repo(examples: Vec<Example>) -> VecDeque<Vec<Example>> {
296    let mut examples_by_repo: HashMap<String, Vec<Example>> = HashMap::default();
297    let mut ungrouped = Vec::new();
298    for example in examples {
299        if example.spec.repository_url.is_empty() {
300            ungrouped.push(example);
301        } else {
302            examples_by_repo
303                .entry(example.spec.repository_url.clone())
304                .or_insert_with(Vec::new)
305                .push(example);
306        }
307    }
308    let mut result: VecDeque<Vec<Example>> = examples_by_repo.into_values().collect();
309    for example in ungrouped {
310        result.push_back(vec![example]);
311    }
312    result
313}
314
315fn parse_markdown_example(input: &str) -> Result<Example> {
316    let spec = ExampleSpec::from_markdown(input)?;
317    Ok(Example {
318        spec,
319        prompt_inputs: None,
320        prompt: None,
321        predictions: Vec::new(),
322        score: Vec::new(),
323        qa: Vec::new(),
324        state: None,
325        zed_version: None,
326    })
327}
328