2680 lines · 96.6 KB · Rust
RawDownload
#[cfg(test)]mod context_tests; use crate::{    prompts::PromptBuilder, slash_command::SlashCommandLine, workflow::WorkflowStep, MessageId,    MessageStatus,};use anyhow::{anyhow, Context as _, Result};use assistant_slash_command::{    SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry,};use client::{self, proto, telemetry::Telemetry};use clock::ReplicaId;use collections::{HashMap, HashSet};use fs::{Fs, RemoveOptions};use futures::{future::Shared, stream::FuturesUnordered, FutureExt, StreamExt};use gpui::{    AppContext, Context as _, EventEmitter, Image, Model, ModelContext, RenderImage, SharedString,    Subscription, Task,}; use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};use language_model::{    LanguageModel, LanguageModelCacheConfiguration, LanguageModelImage, LanguageModelRegistry,    LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,};use open_ai::Model as OpenAiModel;use paths::{context_images_dir, contexts_dir};use project::Project;use serde::{Deserialize, Serialize};use smallvec::SmallVec;use std::{    cmp::{max, Ordering},    collections::hash_map,    fmt::Debug,    iter, mem,    ops::Range,    path::{Path, PathBuf},    sync::Arc,    time::{Duration, Instant},};use telemetry_events::AssistantKind;use text::BufferSnapshot;use util::{post_inc, ResultExt, TryFutureExt};use uuid::Uuid; #[derive(Clone, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]pub struct ContextId(String); impl ContextId {    pub fn new() -> Self {        Self(Uuid::new_v4().to_string())    }     pub fn from_proto(id: String) -> Self {        Self(id)    }     pub fn to_proto(&self) -> String {        self.0.clone()    }} #[derive(Clone, Debug)]pub enum ContextOperation {    InsertMessage {        anchor: MessageAnchor,        metadata: MessageMetadata,        version: clock::Global,    },    UpdateMessage {        message_id: MessageId,        metadata: MessageMetadata,        version: clock::Global,    },    UpdateSummary {        summary: ContextSummary,        version: clock::Global,    },    SlashCommandFinished {        id: SlashCommandId,        output_range: Range<language::Anchor>,        sections: Vec<SlashCommandOutputSection<language::Anchor>>,        version: clock::Global,    },    BufferOperation(language::Operation),} impl ContextOperation {    pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {        match op.variant.context("invalid variant")? {            proto::context_operation::Variant::InsertMessage(insert) => {                let message = insert.message.context("invalid message")?;                let id = MessageId(language::proto::deserialize_timestamp(                    message.id.context("invalid id")?,                ));                Ok(Self::InsertMessage {                    anchor: MessageAnchor {                        id,                        start: language::proto::deserialize_anchor(                            message.start.context("invalid anchor")?,                        )                        .context("invalid anchor")?,                    },                    metadata: MessageMetadata {                        role: Role::from_proto(message.role),                        status: MessageStatus::from_proto(                            message.status.context("invalid status")?,                        ),                        timestamp: id.0,                        cache: None,                    },                    version: language::proto::deserialize_version(&insert.version),                })            }            proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {                message_id: MessageId(language::proto::deserialize_timestamp(                    update.message_id.context("invalid message id")?,                )),                metadata: MessageMetadata {                    role: Role::from_proto(update.role),                    status: MessageStatus::from_proto(update.status.context("invalid status")?),                    timestamp: language::proto::deserialize_timestamp(                        update.timestamp.context("invalid timestamp")?,                    ),                    cache: None,                },                version: language::proto::deserialize_version(&update.version),            }),            proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {                summary: ContextSummary {                    text: update.summary,                    done: update.done,                    timestamp: language::proto::deserialize_timestamp(                        update.timestamp.context("invalid timestamp")?,                    ),                },                version: language::proto::deserialize_version(&update.version),            }),            proto::context_operation::Variant::SlashCommandFinished(finished) => {                Ok(Self::SlashCommandFinished {                    id: SlashCommandId(language::proto::deserialize_timestamp(                        finished.id.context("invalid id")?,                    )),                    output_range: language::proto::deserialize_anchor_range(                        finished.output_range.context("invalid range")?,                    )?,                    sections: finished                        .sections                        .into_iter()                        .map(|section| {                            Ok(SlashCommandOutputSection {                                range: language::proto::deserialize_anchor_range(                                    section.range.context("invalid range")?,                                )?,                                icon: section.icon_name.parse()?,                                label: section.label.into(),                            })                        })                        .collect::<Result<Vec<_>>>()?,                    version: language::proto::deserialize_version(&finished.version),                })            }            proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(                language::proto::deserialize_operation(                    op.operation.context("invalid buffer operation")?,                )?,            )),        }    }     pub fn to_proto(&self) -> proto::ContextOperation {        match self {            Self::InsertMessage {                anchor,                metadata,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::InsertMessage(                    proto::context_operation::InsertMessage {                        message: Some(proto::ContextMessage {                            id: Some(language::proto::serialize_timestamp(anchor.id.0)),                            start: Some(language::proto::serialize_anchor(&anchor.start)),                            role: metadata.role.to_proto() as i32,                            status: Some(metadata.status.to_proto()),                        }),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::UpdateMessage {                message_id,                metadata,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::UpdateMessage(                    proto::context_operation::UpdateMessage {                        message_id: Some(language::proto::serialize_timestamp(message_id.0)),                        role: metadata.role.to_proto() as i32,                        status: Some(metadata.status.to_proto()),                        timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::UpdateSummary { summary, version } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::UpdateSummary(                    proto::context_operation::UpdateSummary {                        summary: summary.text.clone(),                        done: summary.done,                        timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::SlashCommandFinished {                id,                output_range,                sections,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::SlashCommandFinished(                    proto::context_operation::SlashCommandFinished {                        id: Some(language::proto::serialize_timestamp(id.0)),                        output_range: Some(language::proto::serialize_anchor_range(                            output_range.clone(),                        )),                        sections: sections                            .iter()                            .map(|section| {                                let icon_name: &'static str = section.icon.into();                                proto::SlashCommandOutputSection {                                    range: Some(language::proto::serialize_anchor_range(                                        section.range.clone(),                                    )),                                    icon_name: icon_name.to_string(),                                    label: section.label.to_string(),                                }                            })                            .collect(),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::BufferOperation(operation) => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::BufferOperation(                    proto::context_operation::BufferOperation {                        operation: Some(language::proto::serialize_operation(operation)),                    },                )),            },        }    }     fn timestamp(&self) -> clock::Lamport {        match self {            Self::InsertMessage { anchor, .. } => anchor.id.0,            Self::UpdateMessage { metadata, .. } => metadata.timestamp,            Self::UpdateSummary { summary, .. } => summary.timestamp,            Self::SlashCommandFinished { id, .. } => id.0,            Self::BufferOperation(_) => {                panic!("reading the timestamp of a buffer operation is not supported")            }        }    }     /// Returns the current version of the context operation.    pub fn version(&self) -> &clock::Global {        match self {            Self::InsertMessage { version, .. }            | Self::UpdateMessage { version, .. }            | Self::UpdateSummary { version, .. }            | Self::SlashCommandFinished { version, .. } => version,            Self::BufferOperation(_) => {                panic!("reading the version of a buffer operation is not supported")            }        }    }} #[derive(Debug, Clone)]pub enum ContextEvent {    ShowAssistError(SharedString),    MessagesEdited,    SummaryChanged,    WorkflowStepsRemoved(Vec<Range<language::Anchor>>),    WorkflowStepUpdated(Range<language::Anchor>),    StreamedCompletion,    PendingSlashCommandsUpdated {        removed: Vec<Range<language::Anchor>>,        updated: Vec<PendingSlashCommand>,    },    SlashCommandFinished {        output_range: Range<language::Anchor>,        sections: Vec<SlashCommandOutputSection<language::Anchor>>,        run_commands_in_output: bool,        expand_result: bool,    },    Operation(ContextOperation),} #[derive(Clone, Default, Debug)]pub struct ContextSummary {    pub text: String,    done: bool,    timestamp: clock::Lamport,} #[derive(Clone, Debug, Eq, PartialEq)]pub struct MessageAnchor {    pub id: MessageId,    pub start: language::Anchor,} #[derive(Clone, Debug, Eq, PartialEq)]pub enum CacheStatus {    Pending,    Cached,} #[derive(Clone, Debug, Eq, PartialEq)]pub struct MessageCacheMetadata {    pub is_anchor: bool,    pub is_final_anchor: bool,    pub status: CacheStatus,    pub cached_at: clock::Global,} #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]pub struct MessageMetadata {    pub role: Role,    pub status: MessageStatus,    timestamp: clock::Lamport,    #[serde(skip)]    pub cache: Option<MessageCacheMetadata>,} impl MessageMetadata {    pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {        let result = match &self.cache {            Some(MessageCacheMetadata { cached_at, .. }) => !buffer.has_edits_since_in_range(                &cached_at,                Range {                    start: buffer.anchor_at(range.start, Bias::Right),                    end: buffer.anchor_at(range.end, Bias::Left),                },            ),            _ => false,        };        result    }} #[derive(Clone, Debug)]pub struct MessageImage {    image_id: u64,    image: Shared<Task<Option<LanguageModelImage>>>,} impl PartialEq for MessageImage {    fn eq(&self, other: &Self) -> bool {        self.image_id == other.image_id    }} impl Eq for MessageImage {} #[derive(Clone, Debug)]pub struct Message {    pub image_offsets: SmallVec<[(usize, MessageImage); 1]>,    pub offset_range: Range<usize>,    pub index_range: Range<usize>,    pub id: MessageId,    pub anchor: language::Anchor,    pub role: Role,    pub status: MessageStatus,    pub cache: Option<MessageCacheMetadata>,} impl Message {    fn to_request_message(&self, buffer: &Buffer) -> Option<LanguageModelRequestMessage> {        let mut content = Vec::new();         let mut range_start = self.offset_range.start;        for (image_offset, message_image) in self.image_offsets.iter() {            if *image_offset != range_start {                if let Some(text) = Self::collect_text_content(buffer, range_start..*image_offset) {                    content.push(text);                }            }             if let Some(image) = message_image.image.clone().now_or_never().flatten() {                content.push(language_model::MessageContent::Image(image));            }             range_start = *image_offset;        }        if range_start != self.offset_range.end {            if let Some(text) =                Self::collect_text_content(buffer, range_start..self.offset_range.end)            {                content.push(text);            }        }         if content.is_empty() {            return None;        }         Some(LanguageModelRequestMessage {            role: self.role,            content,            cache: self.cache.as_ref().map_or(false, |cache| cache.is_anchor),        })    }     fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<MessageContent> {        let text: String = buffer.text_for_range(range.clone()).collect();        if text.trim().is_empty() {            None        } else {            Some(MessageContent::Text(text))        }    }} #[derive(Clone, Debug)]pub struct ImageAnchor {    pub anchor: language::Anchor,    pub image_id: u64,    pub render_image: Arc<RenderImage>,    pub image: Shared<Task<Option<LanguageModelImage>>>,} struct PendingCompletion {    id: usize,    assistant_message_id: MessageId,    _task: Task<()>,} #[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]pub struct SlashCommandId(clock::Lamport); struct WorkflowStepEntry {    range: Range<language::Anchor>,    step: Model<WorkflowStep>,} pub struct Context {    id: ContextId,    timestamp: clock::Lamport,    version: clock::Global,    pending_ops: Vec<ContextOperation>,    operations: Vec<ContextOperation>,    buffer: Model<Buffer>,    pending_slash_commands: Vec<PendingSlashCommand>,    edits_since_last_slash_command_parse: language::Subscription,    finished_slash_commands: HashSet<SlashCommandId>,    slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,    message_anchors: Vec<MessageAnchor>,    images: HashMap<u64, (Arc<RenderImage>, Shared<Task<Option<LanguageModelImage>>>)>,    image_anchors: Vec<ImageAnchor>,    messages_metadata: HashMap<MessageId, MessageMetadata>,    summary: Option<ContextSummary>,    pending_summary: Task<Option<()>>,    completion_count: usize,    pending_completions: Vec<PendingCompletion>,    token_count: Option<usize>,    pending_token_count: Task<Option<()>>,    pending_save: Task<Result<()>>,    pending_cache_warming_task: Task<Option<()>>,    path: Option<PathBuf>,    _subscriptions: Vec<Subscription>,    telemetry: Option<Arc<Telemetry>>,    language_registry: Arc<LanguageRegistry>,    workflow_steps: Vec<WorkflowStepEntry>,    edits_since_last_workflow_step_prune: language::Subscription,    project: Option<Model<Project>>,    prompt_builder: Arc<PromptBuilder>,} impl EventEmitter<ContextEvent> for Context {} impl Context {    pub fn local(        language_registry: Arc<LanguageRegistry>,        project: Option<Model<Project>>,        telemetry: Option<Arc<Telemetry>>,        prompt_builder: Arc<PromptBuilder>,        cx: &mut ModelContext<Self>,    ) -> Self {        Self::new(            ContextId::new(),            ReplicaId::default(),            language::Capability::ReadWrite,            language_registry,            prompt_builder,            project,            telemetry,            cx,

Lines 1–500 of 2680.