1992 lines · 66.0 KB · Rust
RawDownload
mod completion_diff_element;mod init;mod input_excerpt;mod license_detection;mod onboarding_banner;mod onboarding_modal;mod onboarding_telemetry;mod rate_completion_modal; pub(crate) use completion_diff_element::*;use db::kvp::KEY_VALUE_STORE;pub use init::*;use inline_completion::DataCollectionState;pub use license_detection::is_license_eligible_for_data_collection;use license_detection::LICENSE_FILES_TO_CHECK;pub use onboarding_banner::*;pub use rate_completion_modal::*; use anyhow::{anyhow, Context as _, Result};use arrayvec::ArrayVec;use client::{Client, UserStore};use collections::{HashMap, HashSet, VecDeque};use futures::AsyncReadExt;use gpui::{    actions, App, AppContext as _, AsyncApp, Context, Entity, EntityId, Global, SemanticVersion,    Subscription, Task, WeakEntity,};use http_client::{HttpClient, Method};use input_excerpt::excerpt_for_cursor_position;use language::{    text_diff, Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, ToOffset, ToPoint,};use language_model::{LlmApiToken, RefreshLlmTokenListener};use postage::watch;use project::Project;use release_channel::AppVersion;use settings::WorktreeId;use std::str::FromStr;use std::{    borrow::Cow,    cmp,    fmt::Write,    future::Future,    mem,    ops::Range,    path::Path,    rc::Rc,    sync::Arc,    time::{Duration, Instant},};use telemetry_events::InlineCompletionRating;use thiserror::Error;use util::ResultExt;use uuid::Uuid;use workspace::notifications::{ErrorMessagePrompt, NotificationId};use workspace::Workspace;use worktree::Worktree;use zed_llm_client::{    PredictEditsBody, PredictEditsResponse, EXPIRED_LLM_TOKEN_HEADER_NAME,    MINIMUM_REQUIRED_VERSION_HEADER_NAME,}; const CURSOR_MARKER: &'static str = "<|user_cursor_is_here|>";const START_OF_FILE_MARKER: &'static str = "<|start_of_file|>";const EDITABLE_REGION_START_MARKER: &'static str = "<|editable_region_start|>";const EDITABLE_REGION_END_MARKER: &'static str = "<|editable_region_end|>";const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1);const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice"; const MAX_CONTEXT_TOKENS: usize = 150;const MAX_REWRITE_TOKENS: usize = 350;const MAX_EVENT_TOKENS: usize = 500; /// Maximum number of events to track.const MAX_EVENT_COUNT: usize = 16; actions!(edit_prediction, [ClearHistory]); #[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]pub struct InlineCompletionId(Uuid); impl From<InlineCompletionId> for gpui::ElementId {    fn from(value: InlineCompletionId) -> Self {        gpui::ElementId::Uuid(value.0)    }} impl std::fmt::Display for InlineCompletionId {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}", self.0)    }} #[derive(Clone)]struct ZetaGlobal(Entity<Zeta>); impl Global for ZetaGlobal {} #[derive(Clone)]pub struct InlineCompletion {    id: InlineCompletionId,    path: Arc<Path>,    excerpt_range: Range<usize>,    cursor_offset: usize,    edits: Arc<[(Range<Anchor>, String)]>,    snapshot: BufferSnapshot,    edit_preview: EditPreview,    input_outline: Arc<str>,    input_events: Arc<str>,    input_excerpt: Arc<str>,    output_excerpt: Arc<str>,    request_sent_at: Instant,    response_received_at: Instant,} impl InlineCompletion {    fn latency(&self) -> Duration {        self.response_received_at            .duration_since(self.request_sent_at)    }     fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, String)>> {        interpolate(&self.snapshot, new_snapshot, self.edits.clone())    }} fn interpolate(    old_snapshot: &BufferSnapshot,    new_snapshot: &BufferSnapshot,    current_edits: Arc<[(Range<Anchor>, String)]>,) -> Option<Vec<(Range<Anchor>, String)>> {    let mut edits = Vec::new();     let mut model_edits = current_edits.into_iter().peekable();    for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {        while let Some((model_old_range, _)) = model_edits.peek() {            let model_old_range = model_old_range.to_offset(old_snapshot);            if model_old_range.end < user_edit.old.start {                let (model_old_range, model_new_text) = model_edits.next().unwrap();                edits.push((model_old_range.clone(), model_new_text.clone()));            } else {                break;            }        }         if let Some((model_old_range, model_new_text)) = model_edits.peek() {            let model_old_offset_range = model_old_range.to_offset(old_snapshot);            if user_edit.old == model_old_offset_range {                let user_new_text = new_snapshot                    .text_for_range(user_edit.new.clone())                    .collect::<String>();                 if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {                    if !model_suffix.is_empty() {                        let anchor = old_snapshot.anchor_after(user_edit.old.end);                        edits.push((anchor..anchor, model_suffix.to_string()));                    }                     model_edits.next();                    continue;                }            }        }         return None;    }     edits.extend(model_edits.cloned());     if edits.is_empty() {        None    } else {        Some(edits)    }} impl std::fmt::Debug for InlineCompletion {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("InlineCompletion")            .field("id", &self.id)            .field("path", &self.path)            .field("edits", &self.edits)            .finish_non_exhaustive()    }} pub struct Zeta {    workspace: Option<WeakEntity<Workspace>>,    client: Arc<Client>,    events: VecDeque<Event>,    registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,    shown_completions: VecDeque<InlineCompletion>,    rated_completions: HashSet<InlineCompletionId>,    data_collection_choice: Entity<DataCollectionChoice>,    llm_token: LlmApiToken,    _llm_token_subscription: Subscription,    /// Whether the terms of service have been accepted.    tos_accepted: bool,    /// Whether an update to a newer version of Zed is required to continue using Zeta.    update_required: bool,    _user_store_subscription: Subscription,    license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,} impl Zeta {    pub fn global(cx: &mut App) -> Option<Entity<Self>> {        cx.try_global::<ZetaGlobal>().map(|global| global.0.clone())    }     pub fn register(        workspace: Option<WeakEntity<Workspace>>,        worktree: Option<Entity<Worktree>>,        client: Arc<Client>,        user_store: Entity<UserStore>,        cx: &mut App,    ) -> Entity<Self> {        let this = Self::global(cx).unwrap_or_else(|| {            let entity = cx.new(|cx| Self::new(workspace, client, user_store, cx));            cx.set_global(ZetaGlobal(entity.clone()));            entity        });         this.update(cx, move |this, cx| {            if let Some(worktree) = worktree {                worktree.update(cx, |worktree, cx| {                    this.license_detection_watchers                        .entry(worktree.id())                        .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(worktree, cx)));                });            }        });         this    }     pub fn clear_history(&mut self) {        self.events.clear();    }     fn new(        workspace: Option<WeakEntity<Workspace>>,        client: Arc<Client>,        user_store: Entity<UserStore>,        cx: &mut Context<Self>,    ) -> Self {        let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);         let data_collection_choice = Self::load_data_collection_choices();        let data_collection_choice = cx.new(|_| data_collection_choice);         Self {            workspace,            client,            events: VecDeque::new(),            shown_completions: VecDeque::new(),            rated_completions: HashSet::default(),            registered_buffers: HashMap::default(),            data_collection_choice,            llm_token: LlmApiToken::default(),            _llm_token_subscription: cx.subscribe(                &refresh_llm_token_listener,                |this, _listener, _event, cx| {                    let client = this.client.clone();                    let llm_token = this.llm_token.clone();                    cx.spawn(|_this, _cx| async move {                        llm_token.refresh(&client).await?;                        anyhow::Ok(())                    })                    .detach_and_log_err(cx);                },            ),            tos_accepted: user_store                .read(cx)                .current_user_has_accepted_terms()                .unwrap_or(false),            update_required: false,            _user_store_subscription: cx.subscribe(&user_store, |this, user_store, event, cx| {                match event {                    client::user::Event::PrivateUserInfoUpdated => {                        this.tos_accepted = user_store                            .read(cx)                            .current_user_has_accepted_terms()                            .unwrap_or(false);                    }                    _ => {}                }            }),            license_detection_watchers: HashMap::default(),        }    }     fn push_event(&mut self, event: Event) {        if let Some(Event::BufferChange {            new_snapshot: last_new_snapshot,            timestamp: last_timestamp,            ..        }) = self.events.back_mut()        {            // Coalesce edits for the same buffer when they happen one after the other.            let Event::BufferChange {                old_snapshot,                new_snapshot,                timestamp,            } = &event;             if timestamp.duration_since(*last_timestamp) <= BUFFER_CHANGE_GROUPING_INTERVAL                && old_snapshot.remote_id() == last_new_snapshot.remote_id()                && old_snapshot.version == last_new_snapshot.version            {                *last_new_snapshot = new_snapshot.clone();                *last_timestamp = *timestamp;                return;            }        }         self.events.push_back(event);        if self.events.len() >= MAX_EVENT_COUNT {            self.events.drain(..MAX_EVENT_COUNT / 2);        }    }     pub fn register_buffer(&mut self, buffer: &Entity<Buffer>, cx: &mut Context<Self>) {        let buffer_id = buffer.entity_id();        let weak_buffer = buffer.downgrade();         if let std::collections::hash_map::Entry::Vacant(entry) =            self.registered_buffers.entry(buffer_id)        {            let snapshot = buffer.read(cx).snapshot();             entry.insert(RegisteredBuffer {                snapshot,                _subscriptions: [                    cx.subscribe(buffer, move |this, buffer, event, cx| {                        this.handle_buffer_event(buffer, event, cx);                    }),                    cx.observe_release(buffer, move |this, _buffer, _cx| {                        this.registered_buffers.remove(&weak_buffer.entity_id());                    }),                ],            });        };    }     fn handle_buffer_event(        &mut self,        buffer: Entity<Buffer>,        event: &language::BufferEvent,        cx: &mut Context<Self>,    ) {        if let language::BufferEvent::Edited = event {            self.report_changes_for_buffer(&buffer, cx);        }    }     #[allow(clippy::too_many_arguments)]    fn request_completion_impl<F, R>(        &mut self,        workspace: Option<Entity<Workspace>>,        project: Option<&Entity<Project>>,        buffer: &Entity<Buffer>,        cursor: language::Anchor,        can_collect_data: bool,        cx: &mut Context<Self>,        perform_predict_edits: F,    ) -> Task<Result<Option<InlineCompletion>>>    where        F: FnOnce(PerformPredictEditsParams) -> R + 'static,        R: Future<Output = Result<PredictEditsResponse>> + Send + 'static,    {        let snapshot = self.report_changes_for_buffer(&buffer, cx);        let diagnostic_groups = snapshot.diagnostic_groups(None);        let cursor_point = cursor.to_point(&snapshot);        let cursor_offset = cursor_point.to_offset(&snapshot);        let events = self.events.clone();        let path: Arc<Path> = snapshot            .file()            .map(|f| Arc::from(f.full_path(cx).as_path()))            .unwrap_or_else(|| Arc::from(Path::new("untitled")));         let zeta = cx.entity();        let client = self.client.clone();        let llm_token = self.llm_token.clone();        let app_version = AppVersion::global(cx);         let buffer = buffer.clone();         let local_lsp_store =            project.and_then(|project| project.read(cx).lsp_store().read(cx).as_local());        let diagnostic_groups = if let Some(local_lsp_store) = local_lsp_store {            Some(                diagnostic_groups                    .into_iter()                    .filter_map(|(language_server_id, diagnostic_group)| {                        let language_server =                            local_lsp_store.running_language_server_for_id(language_server_id)?;                         Some((                            language_server.name(),                            diagnostic_group.resolve::<usize>(&snapshot),                        ))                    })                    .collect::<Vec<_>>(),            )        } else {            None        };         cx.spawn(|_, cx| async move {            let request_sent_at = Instant::now();             struct BackgroundValues {                input_events: String,                input_excerpt: String,                speculated_output: String,                editable_range: Range<usize>,                input_outline: String,            }             let values = cx                .background_spawn({                    let snapshot = snapshot.clone();                    let path = path.clone();                    async move {                        let path = path.to_string_lossy();                        let input_excerpt = excerpt_for_cursor_position(                            cursor_point,                            &path,                            &snapshot,                            MAX_REWRITE_TOKENS,                            MAX_CONTEXT_TOKENS,                        );                        let input_events = prompt_for_events(&events, MAX_EVENT_TOKENS);                        let input_outline = prompt_for_outline(&snapshot);                         anyhow::Ok(BackgroundValues {                            input_events,                            input_excerpt: input_excerpt.prompt,                            speculated_output: input_excerpt.speculated_output,                            editable_range: input_excerpt.editable_range.to_offset(&snapshot),                            input_outline,                        })                    }                })                .await?;             log::debug!(                "Events:\n{}\nExcerpt:\n{:?}",                values.input_events,                values.input_excerpt            );             let body = PredictEditsBody {                input_events: values.input_events.clone(),                input_excerpt: values.input_excerpt.clone(),                speculated_output: Some(values.speculated_output),                outline: Some(values.input_outline.clone()),                can_collect_data,                diagnostic_groups: diagnostic_groups.and_then(|diagnostic_groups| {                    diagnostic_groups                        .into_iter()                        .map(|(name, diagnostic_group)| {                            Ok((name.to_string(), serde_json::to_value(diagnostic_group)?))                        })                        .collect::<Result<Vec<_>>>()                        .log_err()                }),            };             let response = perform_predict_edits(PerformPredictEditsParams {                client,                llm_token,                app_version,                body,            })            .await;            let response = match response {                Ok(response) => response,                Err(err) => {                    if err.is::<ZedUpdateRequiredError>() {                        cx.update(|cx| {                            zeta.update(cx, |zeta, _cx| {                                zeta.update_required = true;                            });                             if let Some(workspace) = workspace {                                workspace.update(cx, |workspace, cx| {                                    workspace.show_notification(                                        NotificationId::unique::<ZedUpdateRequiredError>(),                                        cx,                                        |cx| {                                            cx.new(|_| {                                                ErrorMessagePrompt::new(err.to_string())                                                    .with_link_button(                                                        "Update Zed",                                                        "https://zed.dev/releases",                                                    )                                            })                                        },                                    );

Lines 1–500 of 1992.