use chrono::Duration; use serde::{Deserialize, Serialize}; use std::{ fmt::Display, ops::{Add, Range, Sub}, path::{Path, PathBuf}, sync::Arc, }; use strum::EnumIter; use uuid::Uuid; use crate::PredictEditsGitInfo; // TODO: snippet ordering within file / relative to excerpt #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PredictEditsRequest { pub excerpt: String, pub excerpt_path: Arc, /// Within file pub excerpt_range: Range, pub excerpt_line_range: Range, pub cursor_point: Point, /// Within `signatures` pub excerpt_parent: Option, pub signatures: Vec, pub referenced_declarations: Vec, pub events: Vec, #[serde(default)] pub can_collect_data: bool, #[serde(skip_serializing_if = "Vec::is_empty", default)] pub diagnostic_groups: Vec, #[serde(skip_serializing_if = "is_default", default)] pub diagnostic_groups_truncated: bool, /// Info about the git repository state, only present when can_collect_data is true. #[serde(skip_serializing_if = "Option::is_none", default)] pub git_info: Option, // Only available to staff #[serde(default)] pub debug_info: bool, #[serde(skip_serializing_if = "Option::is_none", default)] pub prompt_max_bytes: Option, #[serde(default)] pub prompt_format: PromptFormat, } #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, EnumIter)] pub enum PromptFormat { MarkedExcerpt, LabeledSections, NumLinesUniDiff, /// Prompt format intended for use via zeta_cli OnlySnippets, } impl PromptFormat { pub const DEFAULT: PromptFormat = PromptFormat::NumLinesUniDiff; } impl Default for PromptFormat { fn default() -> Self { Self::DEFAULT } } impl PromptFormat { pub fn iter() -> impl Iterator { ::iter() } } impl std::fmt::Display for PromptFormat { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { PromptFormat::MarkedExcerpt => write!(f, "Marked Excerpt"), PromptFormat::LabeledSections => write!(f, "Labeled Sections"), PromptFormat::OnlySnippets => write!(f, "Only Snippets"), PromptFormat::NumLinesUniDiff => write!(f, "Numbered Lines / Unified Diff"), } } } #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(any(test, feature = "test-support"), derive(PartialEq))] #[serde(tag = "event")] pub enum Event { BufferChange { path: Option, old_path: Option, diff: String, predicted: bool, }, } impl Display for Event { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Event::BufferChange { path, old_path, diff, predicted, } => { let new_path = path.as_deref().unwrap_or(Path::new("untitled")); let old_path = old_path.as_deref().unwrap_or(new_path); if *predicted { write!( f, "// User accepted prediction:\n--- a/{}\n+++ b/{}\n{diff}", old_path.display(), new_path.display() ) } else { write!( f, "--- a/{}\n+++ b/{}\n{diff}", old_path.display(), new_path.display() ) } } } } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Signature { pub text: String, pub text_is_truncated: bool, #[serde(skip_serializing_if = "Option::is_none", default)] pub parent_index: Option, /// Range of `text` within the file, possibly truncated according to `text_is_truncated`. The /// file is implicitly the file that contains the descendant declaration or excerpt. pub range: Range, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ReferencedDeclaration { pub path: Arc, pub text: String, pub text_is_truncated: bool, /// Range of `text` within file, possibly truncated according to `text_is_truncated` pub range: Range, /// Range within `text` pub signature_range: Range, /// Index within `signatures`. #[serde(skip_serializing_if = "Option::is_none", default)] pub parent_index: Option, pub score_components: DeclarationScoreComponents, pub signature_score: f32, pub declaration_score: f32, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DeclarationScoreComponents { pub is_same_file: bool, pub is_referenced_nearby: bool, pub is_referenced_in_breadcrumb: bool, pub reference_count: usize, pub same_file_declaration_count: usize, pub declaration_count: usize, pub reference_line_distance: u32, pub declaration_line_distance: u32, pub excerpt_vs_item_jaccard: f32, pub excerpt_vs_signature_jaccard: f32, pub adjacent_vs_item_jaccard: f32, pub adjacent_vs_signature_jaccard: f32, pub excerpt_vs_item_weighted_overlap: f32, pub excerpt_vs_signature_weighted_overlap: f32, pub adjacent_vs_item_weighted_overlap: f32, pub adjacent_vs_signature_weighted_overlap: f32, pub path_import_match_count: usize, pub wildcard_path_import_match_count: usize, pub import_similarity: f32, pub max_import_similarity: f32, pub normalized_import_similarity: f32, pub wildcard_import_similarity: f32, pub normalized_wildcard_import_similarity: f32, pub included_by_others: usize, pub includes_others: usize, } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(transparent)] pub struct DiagnosticGroup(pub Box); #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PredictEditsResponse { pub request_id: Uuid, pub edits: Vec, pub debug_info: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DebugInfo { pub prompt: String, pub prompt_planning_time: Duration, pub model_response: String, pub inference_time: Duration, pub parsing_time: Duration, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Edit { pub path: Arc, pub range: Range, pub content: String, } fn is_default(value: &T) -> bool { *value == T::default() } #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, PartialOrd, Eq, Ord)] pub struct Point { pub line: Line, pub column: u32, } #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, PartialOrd, Eq, Ord)] #[serde(transparent)] pub struct Line(pub u32); impl Add for Line { type Output = Self; fn add(self, rhs: Self) -> Self::Output { Self(self.0 + rhs.0) } } impl Sub for Line { type Output = Self; fn sub(self, rhs: Self) -> Self::Output { Self(self.0 - rhs.0) } } #[cfg(test)] mod tests { use super::*; use indoc::indoc; use pretty_assertions::assert_eq; #[test] fn test_event_display() { let ev = Event::BufferChange { path: None, old_path: None, diff: "@@ -1,2 +1,2 @@\n-a\n-b\n".into(), predicted: false, }; assert_eq!( ev.to_string(), indoc! {" --- a/untitled +++ b/untitled @@ -1,2 +1,2 @@ -a -b "} ); let ev = Event::BufferChange { path: Some(PathBuf::from("foo/bar.txt")), old_path: Some(PathBuf::from("foo/bar.txt")), diff: "@@ -1,2 +1,2 @@\n-a\n-b\n".into(), predicted: false, }; assert_eq!( ev.to_string(), indoc! {" --- a/foo/bar.txt +++ b/foo/bar.txt @@ -1,2 +1,2 @@ -a -b "} ); let ev = Event::BufferChange { path: Some(PathBuf::from("abc.txt")), old_path: Some(PathBuf::from("123.txt")), diff: "@@ -1,2 +1,2 @@\n-a\n-b\n".into(), predicted: false, }; assert_eq!( ev.to_string(), indoc! {" --- a/123.txt +++ b/abc.txt @@ -1,2 +1,2 @@ -a -b "} ); let ev = Event::BufferChange { path: Some(PathBuf::from("abc.txt")), old_path: Some(PathBuf::from("123.txt")), diff: "@@ -1,2 +1,2 @@\n-a\n-b\n".into(), predicted: true, }; assert_eq!( ev.to_string(), indoc! {" // User accepted prediction: --- a/123.txt +++ b/abc.txt @@ -1,2 +1,2 @@ -a -b "} ); } }