assistant2.rs

  1mod assistant_settings;
  2mod completion_provider;
  3pub mod tools;
  4
  5use anyhow::{Context, Result};
  6use assistant_tooling::{ToolFunctionCall, ToolRegistry};
  7use client::{proto, Client};
  8use completion_provider::*;
  9use editor::{Editor, EditorEvent};
 10use feature_flags::FeatureFlagAppExt as _;
 11use futures::{channel::oneshot, future::join_all, Future, FutureExt, StreamExt};
 12use gpui::{
 13    list, prelude::*, AnyElement, AppContext, AsyncWindowContext, EventEmitter, FocusHandle,
 14    FocusableView, Global, ListAlignment, ListState, Model, Render, Task, View, WeakView,
 15};
 16use language::{language_settings::SoftWrap, LanguageRegistry};
 17use open_ai::{FunctionContent, ToolCall, ToolCallContent};
 18use project::Fs;
 19use rich_text::RichText;
 20use semantic_index::{CloudEmbeddingProvider, ProjectIndex, SemanticIndex};
 21use serde::Deserialize;
 22use settings::Settings;
 23use std::{cmp, sync::Arc};
 24use theme::ThemeSettings;
 25use tools::ProjectIndexTool;
 26use ui::{popover_menu, prelude::*, ButtonLike, CollapsibleContainer, Color, ContextMenu, Tooltip};
 27use util::{paths::EMBEDDINGS_DIR, ResultExt};
 28use workspace::{
 29    dock::{DockPosition, Panel, PanelEvent},
 30    Workspace,
 31};
 32
 33pub use assistant_settings::AssistantSettings;
 34
 35const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
 36
 37// gpui::actions!(assistant, [Submit]);
 38
 39#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
 40pub struct Submit(SubmitMode);
 41
 42/// There are multiple different ways to submit a model request, represented by this enum.
 43#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
 44pub enum SubmitMode {
 45    /// Only include the conversation.
 46    Simple,
 47    /// Send the current file as context.
 48    CurrentFile,
 49    /// Search the codebase and send relevant excerpts.
 50    Codebase,
 51}
 52
 53gpui::actions!(assistant2, [ToggleFocus]);
 54gpui::impl_actions!(assistant2, [Submit]);
 55
 56pub fn init(client: Arc<Client>, cx: &mut AppContext) {
 57    AssistantSettings::register(cx);
 58
 59    cx.spawn(|mut cx| {
 60        let client = client.clone();
 61        async move {
 62            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
 63            let semantic_index = SemanticIndex::new(
 64                EMBEDDINGS_DIR.join("semantic-index-db.0.mdb"),
 65                Arc::new(embedding_provider),
 66                &mut cx,
 67            )
 68            .await?;
 69            cx.update(|cx| cx.set_global(semantic_index))
 70        }
 71    })
 72    .detach();
 73
 74    cx.set_global(CompletionProvider::new(CloudCompletionProvider::new(
 75        client,
 76    )));
 77
 78    cx.observe_new_views(
 79        |workspace: &mut Workspace, _cx: &mut ViewContext<Workspace>| {
 80            workspace.register_action(|workspace, _: &ToggleFocus, cx| {
 81                workspace.toggle_panel_focus::<AssistantPanel>(cx);
 82            });
 83        },
 84    )
 85    .detach();
 86}
 87
 88pub fn enabled(cx: &AppContext) -> bool {
 89    cx.is_staff()
 90}
 91
 92pub struct AssistantPanel {
 93    chat: View<AssistantChat>,
 94    width: Option<Pixels>,
 95}
 96
 97impl AssistantPanel {
 98    pub fn load(
 99        workspace: WeakView<Workspace>,
100        cx: AsyncWindowContext,
101    ) -> Task<Result<View<Self>>> {
102        cx.spawn(|mut cx| async move {
103            let (app_state, project) = workspace.update(&mut cx, |workspace, _| {
104                (workspace.app_state().clone(), workspace.project().clone())
105            })?;
106
107            cx.new_view(|cx| {
108                // todo!("this will panic if the semantic index failed to load or has not loaded yet")
109                let project_index = cx.update_global(|semantic_index: &mut SemanticIndex, cx| {
110                    semantic_index.project_index(project.clone(), cx)
111                });
112
113                let mut tool_registry = ToolRegistry::new();
114                tool_registry
115                    .register(ProjectIndexTool::new(
116                        project_index.clone(),
117                        app_state.fs.clone(),
118                    ))
119                    .context("failed to register ProjectIndexTool")
120                    .log_err();
121
122                let tool_registry = Arc::new(tool_registry);
123
124                Self::new(app_state.languages.clone(), tool_registry, cx)
125            })
126        })
127    }
128
129    pub fn new(
130        language_registry: Arc<LanguageRegistry>,
131        tool_registry: Arc<ToolRegistry>,
132        cx: &mut ViewContext<Self>,
133    ) -> Self {
134        let chat = cx.new_view(|cx| {
135            AssistantChat::new(language_registry.clone(), tool_registry.clone(), cx)
136        });
137
138        Self { width: None, chat }
139    }
140}
141
142impl Render for AssistantPanel {
143    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
144        div()
145            .size_full()
146            .v_flex()
147            .p_2()
148            .bg(cx.theme().colors().background)
149            .child(self.chat.clone())
150    }
151}
152
153impl Panel for AssistantPanel {
154    fn persistent_name() -> &'static str {
155        "AssistantPanelv2"
156    }
157
158    fn position(&self, _cx: &WindowContext) -> workspace::dock::DockPosition {
159        // todo!("Add a setting / use assistant settings")
160        DockPosition::Right
161    }
162
163    fn position_is_valid(&self, position: workspace::dock::DockPosition) -> bool {
164        matches!(position, DockPosition::Right)
165    }
166
167    fn set_position(&mut self, _: workspace::dock::DockPosition, _: &mut ViewContext<Self>) {
168        // Do nothing until we have a setting for this
169    }
170
171    fn size(&self, _cx: &WindowContext) -> Pixels {
172        self.width.unwrap_or(px(400.))
173    }
174
175    fn set_size(&mut self, size: Option<Pixels>, cx: &mut ViewContext<Self>) {
176        self.width = size;
177        cx.notify();
178    }
179
180    fn icon(&self, _cx: &WindowContext) -> Option<ui::IconName> {
181        Some(IconName::Ai)
182    }
183
184    fn icon_tooltip(&self, _: &WindowContext) -> Option<&'static str> {
185        Some("Assistant Panel ✨")
186    }
187
188    fn toggle_action(&self) -> Box<dyn gpui::Action> {
189        Box::new(ToggleFocus)
190    }
191}
192
193impl EventEmitter<PanelEvent> for AssistantPanel {}
194
195impl FocusableView for AssistantPanel {
196    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
197        self.chat
198            .read(cx)
199            .messages
200            .iter()
201            .rev()
202            .find_map(|msg| msg.focus_handle(cx))
203            .expect("no user message in chat")
204    }
205}
206
207struct AssistantChat {
208    model: String,
209    messages: Vec<ChatMessage>,
210    list_state: ListState,
211    language_registry: Arc<LanguageRegistry>,
212    next_message_id: MessageId,
213    pending_completion: Option<Task<()>>,
214    tool_registry: Arc<ToolRegistry>,
215}
216
217impl AssistantChat {
218    fn new(
219        language_registry: Arc<LanguageRegistry>,
220        tool_registry: Arc<ToolRegistry>,
221        cx: &mut ViewContext<Self>,
222    ) -> Self {
223        let model = CompletionProvider::get(cx).default_model();
224        let view = cx.view().downgrade();
225        let list_state = ListState::new(
226            0,
227            ListAlignment::Bottom,
228            px(1024.),
229            move |ix, cx: &mut WindowContext| {
230                view.update(cx, |this, cx| this.render_message(ix, cx))
231                    .unwrap()
232            },
233        );
234
235        let mut this = Self {
236            model,
237            messages: Vec::new(),
238            list_state,
239            language_registry,
240            next_message_id: MessageId(0),
241            pending_completion: None,
242            tool_registry,
243        };
244        this.push_new_user_message(true, cx);
245        this
246    }
247
248    fn focused_message_id(&self, cx: &WindowContext) -> Option<MessageId> {
249        self.messages.iter().find_map(|message| match message {
250            ChatMessage::User(message) => message
251                .body
252                .focus_handle(cx)
253                .contains_focused(cx)
254                .then_some(message.id),
255            ChatMessage::Assistant(_) => None,
256        })
257    }
258
259    fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {
260        let Some(focused_message_id) = self.focused_message_id(cx) else {
261            log::error!("unexpected state: no user message editor is focused.");
262            return;
263        };
264
265        self.truncate_messages(focused_message_id, cx);
266
267        let mode = *mode;
268        self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {
269            Self::request_completion(
270                this.clone(),
271                mode,
272                MAX_COMPLETION_CALLS_PER_SUBMISSION,
273                &mut cx,
274            )
275            .await
276            .log_err();
277
278            this.update(&mut cx, |this, cx| {
279                let focus = this
280                    .user_message(focused_message_id)
281                    .body
282                    .focus_handle(cx)
283                    .contains_focused(cx);
284                this.push_new_user_message(focus, cx);
285            })
286            .context("Failed to push new user message")
287            .log_err();
288        }));
289    }
290
291    async fn request_completion(
292        this: WeakView<Self>,
293        mode: SubmitMode,
294        limit: usize,
295        cx: &mut AsyncWindowContext,
296    ) -> Result<()> {
297        let mut call_count = 0;
298        loop {
299            let complete = async {
300                let completion = this.update(cx, |this, cx| {
301                    this.push_new_assistant_message(cx);
302
303                    let definitions = if call_count < limit && matches!(mode, SubmitMode::Codebase)
304                    {
305                        this.tool_registry.definitions()
306                    } else {
307                        &[]
308                    };
309                    call_count += 1;
310
311                    CompletionProvider::get(cx).complete(
312                        this.model.clone(),
313                        this.completion_messages(cx),
314                        Vec::new(),
315                        1.0,
316                        definitions,
317                    )
318                });
319
320                let mut stream = completion?.await?;
321                let mut body = String::new();
322                while let Some(delta) = stream.next().await {
323                    let delta = delta?;
324                    this.update(cx, |this, cx| {
325                        if let Some(ChatMessage::Assistant(AssistantMessage {
326                            body: message_body,
327                            tool_calls: message_tool_calls,
328                            ..
329                        })) = this.messages.last_mut()
330                        {
331                            if let Some(content) = &delta.content {
332                                body.push_str(content);
333                            }
334
335                            for tool_call in delta.tool_calls {
336                                let index = tool_call.index as usize;
337                                if index >= message_tool_calls.len() {
338                                    message_tool_calls.resize_with(index + 1, Default::default);
339                                }
340                                let call = &mut message_tool_calls[index];
341
342                                if let Some(id) = &tool_call.id {
343                                    call.id.push_str(id);
344                                }
345
346                                match tool_call.variant {
347                                    Some(proto::tool_call_delta::Variant::Function(tool_call)) => {
348                                        if let Some(name) = &tool_call.name {
349                                            call.name.push_str(name);
350                                        }
351                                        if let Some(arguments) = &tool_call.arguments {
352                                            call.arguments.push_str(arguments);
353                                        }
354                                    }
355                                    None => {}
356                                }
357                            }
358
359                            *message_body =
360                                RichText::new(body.clone(), &[], &this.language_registry);
361                            cx.notify();
362                        } else {
363                            unreachable!()
364                        }
365                    })?;
366                }
367
368                anyhow::Ok(())
369            }
370            .await;
371
372            let mut tool_tasks = Vec::new();
373            this.update(cx, |this, cx| {
374                if let Some(ChatMessage::Assistant(AssistantMessage {
375                    error: message_error,
376                    tool_calls,
377                    ..
378                })) = this.messages.last_mut()
379                {
380                    if let Err(error) = complete {
381                        message_error.replace(SharedString::from(error.to_string()));
382                        cx.notify();
383                    } else {
384                        for tool_call in tool_calls.iter() {
385                            tool_tasks.push(this.tool_registry.call(tool_call, cx));
386                        }
387                    }
388                }
389            })?;
390
391            if tool_tasks.is_empty() {
392                return Ok(());
393            }
394
395            let tools = join_all(tool_tasks.into_iter()).await;
396            this.update(cx, |this, cx| {
397                if let Some(ChatMessage::Assistant(AssistantMessage { tool_calls, .. })) =
398                    this.messages.last_mut()
399                {
400                    *tool_calls = tools;
401                    cx.notify();
402                }
403            })?;
404        }
405    }
406
407    fn user_message(&mut self, message_id: MessageId) -> &mut UserMessage {
408        self.messages
409            .iter_mut()
410            .find_map(|message| match message {
411                ChatMessage::User(user_message) if user_message.id == message_id => {
412                    Some(user_message)
413                }
414                _ => None,
415            })
416            .expect("User message not found")
417    }
418
419    fn push_new_user_message(&mut self, focus: bool, cx: &mut ViewContext<Self>) {
420        let id = self.next_message_id.post_inc();
421        let body = cx.new_view(|cx| {
422            let mut editor = Editor::auto_height(80, cx);
423            editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
424            if focus {
425                cx.focus_self();
426            }
427            editor
428        });
429        let _subscription = cx.subscribe(&body, move |this, editor, event, cx| match event {
430            EditorEvent::SelectionsChanged { .. } => {
431                if editor.read(cx).is_focused(cx) {
432                    let (message_ix, _message) = this
433                        .messages
434                        .iter()
435                        .enumerate()
436                        .find_map(|(ix, message)| match message {
437                            ChatMessage::User(user_message) if user_message.id == id => {
438                                Some((ix, user_message))
439                            }
440                            _ => None,
441                        })
442                        .expect("user message not found");
443
444                    this.list_state.scroll_to_reveal_item(message_ix);
445                }
446            }
447            _ => {}
448        });
449        let message = ChatMessage::User(UserMessage {
450            id,
451            body,
452            contexts: Vec::new(),
453            _subscription,
454        });
455        self.push_message(message, cx);
456    }
457
458    fn push_new_assistant_message(&mut self, cx: &mut ViewContext<Self>) {
459        let message = ChatMessage::Assistant(AssistantMessage {
460            id: self.next_message_id.post_inc(),
461            body: RichText::default(),
462            tool_calls: Vec::new(),
463            error: None,
464        });
465        self.push_message(message, cx);
466    }
467
468    fn push_message(&mut self, message: ChatMessage, cx: &mut ViewContext<Self>) {
469        let old_len = self.messages.len();
470        let focus_handle = Some(message.focus_handle(cx));
471        self.messages.push(message);
472        self.list_state
473            .splice_focusable(old_len..old_len, focus_handle);
474        cx.notify();
475    }
476
477    fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
478        if let Some(index) = self.messages.iter().position(|message| match message {
479            ChatMessage::User(message) => message.id == last_message_id,
480            ChatMessage::Assistant(message) => message.id == last_message_id,
481        }) {
482            self.list_state.splice(index + 1..self.messages.len(), 0);
483            self.messages.truncate(index + 1);
484            cx.notify();
485        }
486    }
487
488    fn render_error(
489        &self,
490        error: Option<SharedString>,
491        _ix: usize,
492        cx: &mut ViewContext<Self>,
493    ) -> AnyElement {
494        let theme = cx.theme();
495
496        if let Some(error) = error {
497            div()
498                .py_1()
499                .px_2()
500                .neg_mx_1()
501                .rounded_md()
502                .border()
503                .border_color(theme.status().error_border)
504                // .bg(theme.status().error_background)
505                .text_color(theme.status().error)
506                .child(error.clone())
507                .into_any_element()
508        } else {
509            div().into_any_element()
510        }
511    }
512
513    fn render_message(&self, ix: usize, cx: &mut ViewContext<Self>) -> AnyElement {
514        let is_last = ix == self.messages.len() - 1;
515
516        match &self.messages[ix] {
517            ChatMessage::User(UserMessage {
518                body,
519                contexts: _contexts,
520                ..
521            }) => div()
522                .when(!is_last, |element| element.mb_2())
523                .child(div().p_2().child(Label::new("You").color(Color::Default)))
524                .child(
525                    div()
526                        .on_action(cx.listener(Self::submit))
527                        .p_2()
528                        .text_color(cx.theme().colors().editor_foreground)
529                        .font(ThemeSettings::get_global(cx).buffer_font.clone())
530                        .bg(cx.theme().colors().editor_background)
531                        .child(body.clone()), // .children(contexts.iter().map(|context| context.render(cx))),
532                )
533                .into_any(),
534            ChatMessage::Assistant(AssistantMessage {
535                id,
536                body,
537                error,
538                tool_calls,
539                ..
540            }) => {
541                let assistant_body = if body.text.is_empty() && !tool_calls.is_empty() {
542                    div()
543                } else {
544                    div().p_2().child(body.element(ElementId::from(id.0), cx))
545                };
546
547                div()
548                    .when(!is_last, |element| element.mb_2())
549                    .child(
550                        div()
551                            .p_2()
552                            .child(Label::new("Assistant").color(Color::Modified)),
553                    )
554                    .child(assistant_body)
555                    .child(self.render_error(error.clone(), ix, cx))
556                    .children(tool_calls.iter().map(|tool_call| {
557                        let result = &tool_call.result;
558                        let name = tool_call.name.clone();
559                        match result {
560                            Some(result) => div()
561                                .p_2()
562                                .child(result.render(&name, &tool_call.id, cx))
563                                .into_any(),
564                            None => div()
565                                .p_2()
566                                .child(Label::new(name).color(Color::Modified))
567                                .child("Running...")
568                                .into_any(),
569                        }
570                    }))
571                    .into_any()
572            }
573        }
574    }
575
576    fn completion_messages(&self, cx: &WindowContext) -> Vec<CompletionMessage> {
577        let mut completion_messages = Vec::new();
578
579        for message in &self.messages {
580            match message {
581                ChatMessage::User(UserMessage { body, contexts, .. }) => {
582                    // setup context for model
583                    contexts.iter().for_each(|context| {
584                        completion_messages.extend(context.completion_messages(cx))
585                    });
586
587                    // Show user's message last so that the assistant is grounded in the user's request
588                    completion_messages.push(CompletionMessage::User {
589                        content: body.read(cx).text(cx),
590                    });
591                }
592                ChatMessage::Assistant(AssistantMessage {
593                    body, tool_calls, ..
594                }) => {
595                    // In no case do we want to send an empty message. This shouldn't happen, but we might as well
596                    // not break the Chat API if it does.
597                    if body.text.is_empty() && tool_calls.is_empty() {
598                        continue;
599                    }
600
601                    let tool_calls_from_assistant = tool_calls
602                        .iter()
603                        .map(|tool_call| ToolCall {
604                            content: ToolCallContent::Function {
605                                function: FunctionContent {
606                                    name: tool_call.name.clone(),
607                                    arguments: tool_call.arguments.clone(),
608                                },
609                            },
610                            id: tool_call.id.clone(),
611                        })
612                        .collect();
613
614                    completion_messages.push(CompletionMessage::Assistant {
615                        content: Some(body.text.to_string()),
616                        tool_calls: tool_calls_from_assistant,
617                    });
618
619                    for tool_call in tool_calls {
620                        // todo!(): we should not be sending when the tool is still running / has no result
621                        // For now I'm going to have to assume we send an empty string because otherwise
622                        // the Chat API will break -- there is a required message for every tool call by ID
623                        let content = match &tool_call.result {
624                            Some(result) => result.format(&tool_call.name),
625                            None => "".to_string(),
626                        };
627
628                        completion_messages.push(CompletionMessage::Tool {
629                            content,
630                            tool_call_id: tool_call.id.clone(),
631                        });
632                    }
633                }
634            }
635        }
636
637        completion_messages
638    }
639
640    fn render_model_dropdown(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
641        let this = cx.view().downgrade();
642        div().h_flex().justify_end().child(
643            div().w_32().child(
644                popover_menu("user-menu")
645                    .menu(move |cx| {
646                        ContextMenu::build(cx, |mut menu, cx| {
647                            for model in CompletionProvider::get(cx).available_models() {
648                                menu = menu.custom_entry(
649                                    {
650                                        let model = model.clone();
651                                        move |_| Label::new(model.clone()).into_any_element()
652                                    },
653                                    {
654                                        let this = this.clone();
655                                        move |cx| {
656                                            _ = this.update(cx, |this, cx| {
657                                                this.model = model.clone();
658                                                cx.notify();
659                                            });
660                                        }
661                                    },
662                                );
663                            }
664                            menu
665                        })
666                        .into()
667                    })
668                    .trigger(
669                        ButtonLike::new("active-model")
670                            .child(
671                                h_flex()
672                                    .w_full()
673                                    .gap_0p5()
674                                    .child(
675                                        div()
676                                            .overflow_x_hidden()
677                                            .flex_grow()
678                                            .whitespace_nowrap()
679                                            .child(Label::new(self.model.clone())),
680                                    )
681                                    .child(div().child(
682                                        Icon::new(IconName::ChevronDown).color(Color::Muted),
683                                    )),
684                            )
685                            .style(ButtonStyle::Subtle)
686                            .tooltip(move |cx| Tooltip::text("Change Model", cx)),
687                    )
688                    .anchor(gpui::AnchorCorner::TopRight),
689            ),
690        )
691    }
692}
693
694impl Render for AssistantChat {
695    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
696        div()
697            .relative()
698            .flex_1()
699            .v_flex()
700            .key_context("AssistantChat")
701            .text_color(Color::Default.color(cx))
702            .child(self.render_model_dropdown(cx))
703            .child(list(self.list_state.clone()).flex_1())
704    }
705}
706
707#[derive(Copy, Clone, Eq, PartialEq)]
708struct MessageId(usize);
709
710impl MessageId {
711    fn post_inc(&mut self) -> Self {
712        let id = *self;
713        self.0 += 1;
714        id
715    }
716}
717
718enum ChatMessage {
719    User(UserMessage),
720    Assistant(AssistantMessage),
721}
722
723impl ChatMessage {
724    fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
725        match self {
726            ChatMessage::User(UserMessage { body, .. }) => Some(body.focus_handle(cx)),
727            ChatMessage::Assistant(_) => None,
728        }
729    }
730}
731
732struct UserMessage {
733    id: MessageId,
734    body: View<Editor>,
735    contexts: Vec<AssistantContext>,
736    _subscription: gpui::Subscription,
737}
738
739struct AssistantMessage {
740    id: MessageId,
741    body: RichText,
742    tool_calls: Vec<ToolFunctionCall>,
743    error: Option<SharedString>,
744}
745
746// Since we're swapping out for direct query usage, we might not need to use this injected context
747// It will be useful though for when the user _definitely_ wants the model to see a specific file,
748// query, error, etc.
749#[allow(dead_code)]
750enum AssistantContext {
751    Codebase(View<CodebaseContext>),
752}
753
754#[allow(dead_code)]
755struct CodebaseExcerpt {
756    element_id: ElementId,
757    path: SharedString,
758    text: SharedString,
759    score: f32,
760    expanded: bool,
761}
762
763impl AssistantContext {
764    #[allow(dead_code)]
765    fn render(&self, _cx: &mut ViewContext<AssistantChat>) -> AnyElement {
766        match self {
767            AssistantContext::Codebase(context) => context.clone().into_any_element(),
768        }
769    }
770
771    fn completion_messages(&self, cx: &WindowContext) -> Vec<CompletionMessage> {
772        match self {
773            AssistantContext::Codebase(context) => context.read(cx).completion_messages(),
774        }
775    }
776}
777
778enum CodebaseContext {
779    Pending { _task: Task<()> },
780    Done(Result<Vec<CodebaseExcerpt>>),
781}
782
783impl CodebaseContext {
784    fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext<Self>) {
785        if let CodebaseContext::Done(Ok(excerpts)) = self {
786            if let Some(excerpt) = excerpts
787                .iter_mut()
788                .find(|excerpt| excerpt.element_id == element_id)
789            {
790                excerpt.expanded = !excerpt.expanded;
791                cx.notify();
792            }
793        }
794    }
795}
796
797impl Render for CodebaseContext {
798    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
799        match self {
800            CodebaseContext::Pending { .. } => div()
801                .h_flex()
802                .items_center()
803                .gap_1()
804                .child(Icon::new(IconName::Ai).color(Color::Muted).into_element())
805                .child("Searching codebase..."),
806            CodebaseContext::Done(Ok(excerpts)) => {
807                div()
808                    .v_flex()
809                    .gap_2()
810                    .children(excerpts.iter().map(|excerpt| {
811                        let expanded = excerpt.expanded;
812                        let element_id = excerpt.element_id.clone();
813
814                        CollapsibleContainer::new(element_id.clone(), expanded)
815                            .start_slot(
816                                h_flex()
817                                    .gap_1()
818                                    .child(Icon::new(IconName::File).color(Color::Muted))
819                                    .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
820                            )
821                            .on_click(cx.listener(move |this, _, cx| {
822                                this.toggle_expanded(element_id.clone(), cx);
823                            }))
824                            .child(
825                                div()
826                                    .p_2()
827                                    .rounded_md()
828                                    .bg(cx.theme().colors().editor_background)
829                                    .child(
830                                        excerpt.text.clone(), // todo!(): Show as an editor block
831                                    ),
832                            )
833                    }))
834            }
835            CodebaseContext::Done(Err(error)) => div().child(error.to_string()),
836        }
837    }
838}
839
840impl CodebaseContext {
841    #[allow(dead_code)]
842    fn new(
843        query: impl 'static + Future<Output = Result<String>>,
844        populated: oneshot::Sender<bool>,
845        project_index: Model<ProjectIndex>,
846        fs: Arc<dyn Fs>,
847        cx: &mut ViewContext<Self>,
848    ) -> Self {
849        let query = query.boxed_local();
850        let _task = cx.spawn(|this, mut cx| async move {
851            let result = async {
852                let query = query.await?;
853                let results = this
854                    .update(&mut cx, |_this, cx| {
855                        project_index.read(cx).search(&query, 16, cx)
856                    })?
857                    .await;
858
859                let excerpts = results.into_iter().map(|result| {
860                    let abs_path = result
861                        .worktree
862                        .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
863                    let fs = fs.clone();
864
865                    async move {
866                        let path = result.path.clone();
867                        let text = fs.load(&abs_path?).await?;
868                        // todo!("what should we do with stale ranges?");
869                        let range = cmp::min(result.range.start, text.len())
870                            ..cmp::min(result.range.end, text.len());
871
872                        let text = SharedString::from(text[range].to_string());
873
874                        anyhow::Ok(CodebaseExcerpt {
875                            element_id: ElementId::Name(nanoid::nanoid!().into()),
876                            path: path.to_string_lossy().to_string().into(),
877                            text,
878                            score: result.score,
879                            expanded: false,
880                        })
881                    }
882                });
883
884                anyhow::Ok(
885                    futures::future::join_all(excerpts)
886                        .await
887                        .into_iter()
888                        .filter_map(|result| result.log_err())
889                        .collect(),
890                )
891            }
892            .await;
893
894            this.update(&mut cx, |this, cx| {
895                this.populate(result, populated, cx);
896            })
897            .ok();
898        });
899
900        Self::Pending { _task }
901    }
902
903    #[allow(dead_code)]
904    fn populate(
905        &mut self,
906        result: Result<Vec<CodebaseExcerpt>>,
907        populated: oneshot::Sender<bool>,
908        cx: &mut ViewContext<Self>,
909    ) {
910        let success = result.is_ok();
911        *self = Self::Done(result);
912        populated.send(success).ok();
913        cx.notify();
914    }
915
916    fn completion_messages(&self) -> Vec<CompletionMessage> {
917        // One system message for the whole batch of excerpts:
918
919        // Semantic search results for user query:
920        //
921        // Excerpt from $path:
922        // ~~~
923        // `text`
924        // ~~~
925        //
926        // Excerpt from $path:
927
928        match self {
929            CodebaseContext::Done(Ok(excerpts)) => {
930                if excerpts.is_empty() {
931                    return Vec::new();
932                }
933
934                let mut body = "Semantic search results for user query:\n".to_string();
935
936                for excerpt in excerpts {
937                    body.push_str("Excerpt from ");
938                    body.push_str(excerpt.path.as_ref());
939                    body.push_str(", score ");
940                    body.push_str(&excerpt.score.to_string());
941                    body.push_str(":\n");
942                    body.push_str("~~~\n");
943                    body.push_str(excerpt.text.as_ref());
944                    body.push_str("~~~\n");
945                }
946
947                vec![CompletionMessage::System { content: body }]
948            }
949            _ => vec![],
950        }
951    }
952}