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