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 && matches!(mode, SubmitMode::Codebase)
318 {
319 this.tool_registry.definitions()
320 } else {
321 &[]
322 };
323 call_count += 1;
324
325 let messages = this.completion_messages(cx);
326
327 CompletionProvider::get(cx).complete(
328 this.model.clone(),
329 messages,
330 Vec::new(),
331 1.0,
332 definitions,
333 )
334 });
335
336 let mut stream = completion?.await?;
337 let mut body = String::new();
338 while let Some(delta) = stream.next().await {
339 let delta = delta?;
340 this.update(cx, |this, cx| {
341 if let Some(ChatMessage::Assistant(AssistantMessage {
342 body: message_body,
343 tool_calls: message_tool_calls,
344 ..
345 })) = this.messages.last_mut()
346 {
347 if let Some(content) = &delta.content {
348 body.push_str(content);
349 }
350
351 for tool_call in delta.tool_calls {
352 let index = tool_call.index as usize;
353 if index >= message_tool_calls.len() {
354 message_tool_calls.resize_with(index + 1, Default::default);
355 }
356 let call = &mut message_tool_calls[index];
357
358 if let Some(id) = &tool_call.id {
359 call.id.push_str(id);
360 }
361
362 match tool_call.variant {
363 Some(proto::tool_call_delta::Variant::Function(tool_call)) => {
364 if let Some(name) = &tool_call.name {
365 call.name.push_str(name);
366 }
367 if let Some(arguments) = &tool_call.arguments {
368 call.arguments.push_str(arguments);
369 }
370 }
371 None => {}
372 }
373 }
374
375 *message_body =
376 RichText::new(body.clone(), &[], &this.language_registry);
377 cx.notify();
378 } else {
379 unreachable!()
380 }
381 })?;
382 }
383
384 anyhow::Ok(())
385 }
386 .await;
387
388 let mut tool_tasks = Vec::new();
389 this.update(cx, |this, cx| {
390 if let Some(ChatMessage::Assistant(AssistantMessage {
391 error: message_error,
392 tool_calls,
393 ..
394 })) = this.messages.last_mut()
395 {
396 if let Err(error) = complete {
397 message_error.replace(SharedString::from(error.to_string()));
398 cx.notify();
399 } else {
400 for tool_call in tool_calls.iter() {
401 tool_tasks.push(this.tool_registry.call(tool_call, cx));
402 }
403 }
404 }
405 })?;
406
407 if tool_tasks.is_empty() {
408 return Ok(());
409 }
410
411 let tools = join_all(tool_tasks.into_iter()).await;
412 // If the WindowContext went away for any tool's view we don't include it
413 // especially since the below call would fail for the same reason.
414 let tools = tools.into_iter().filter_map(|tool| tool.ok()).collect();
415
416 this.update(cx, |this, cx| {
417 if let Some(ChatMessage::Assistant(AssistantMessage { tool_calls, .. })) =
418 this.messages.last_mut()
419 {
420 *tool_calls = tools;
421 cx.notify();
422 }
423 })?;
424 }
425 }
426
427 fn user_message(&mut self, message_id: MessageId) -> &mut UserMessage {
428 self.messages
429 .iter_mut()
430 .find_map(|message| match message {
431 ChatMessage::User(user_message) if user_message.id == message_id => {
432 Some(user_message)
433 }
434 _ => None,
435 })
436 .expect("User message not found")
437 }
438
439 fn push_new_user_message(&mut self, focus: bool, cx: &mut ViewContext<Self>) {
440 let id = self.next_message_id.post_inc();
441 let body = cx.new_view(|cx| {
442 let mut editor = Editor::auto_height(80, cx);
443 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
444 if focus {
445 cx.focus_self();
446 }
447 editor
448 });
449 let message = ChatMessage::User(UserMessage {
450 id,
451 body,
452 contexts: Vec::new(),
453 });
454 self.push_message(message, cx);
455 }
456
457 fn push_new_assistant_message(&mut self, cx: &mut ViewContext<Self>) {
458 let message = ChatMessage::Assistant(AssistantMessage {
459 id: self.next_message_id.post_inc(),
460 body: RichText::default(),
461 tool_calls: Vec::new(),
462 error: None,
463 });
464 self.push_message(message, cx);
465 }
466
467 fn push_message(&mut self, message: ChatMessage, cx: &mut ViewContext<Self>) {
468 let old_len = self.messages.len();
469 let focus_handle = Some(message.focus_handle(cx));
470 self.messages.push(message);
471 self.list_state
472 .splice_focusable(old_len..old_len, focus_handle);
473 cx.notify();
474 }
475
476 fn pop_message(&mut self, cx: &mut ViewContext<Self>) {
477 if self.messages.is_empty() {
478 return;
479 }
480
481 self.messages.pop();
482 self.list_state
483 .splice(self.messages.len()..self.messages.len() + 1, 0);
484 cx.notify();
485 }
486
487 fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
488 if let Some(index) = self.messages.iter().position(|message| match message {
489 ChatMessage::User(message) => message.id == last_message_id,
490 ChatMessage::Assistant(message) => message.id == last_message_id,
491 }) {
492 self.list_state.splice(index + 1..self.messages.len(), 0);
493 self.messages.truncate(index + 1);
494 cx.notify();
495 }
496 }
497
498 fn render_error(
499 &self,
500 error: Option<SharedString>,
501 _ix: usize,
502 cx: &mut ViewContext<Self>,
503 ) -> AnyElement {
504 let theme = cx.theme();
505
506 if let Some(error) = error {
507 div()
508 .py_1()
509 .px_2()
510 .neg_mx_1()
511 .rounded_md()
512 .border()
513 .border_color(theme.status().error_border)
514 // .bg(theme.status().error_background)
515 .text_color(theme.status().error)
516 .child(error.clone())
517 .into_any_element()
518 } else {
519 div().into_any_element()
520 }
521 }
522
523 fn render_message(&self, ix: usize, cx: &mut ViewContext<Self>) -> AnyElement {
524 let is_last = ix == self.messages.len() - 1;
525
526 match &self.messages[ix] {
527 ChatMessage::User(UserMessage {
528 body,
529 contexts: _contexts,
530 ..
531 }) => div()
532 .when(!is_last, |element| element.mb_2())
533 .child(div().p_2().child(Label::new("You").color(Color::Default)))
534 .child(
535 div()
536 .on_action(cx.listener(Self::submit))
537 .p_2()
538 .text_color(cx.theme().colors().editor_foreground)
539 .font(ThemeSettings::get_global(cx).buffer_font.clone())
540 .bg(cx.theme().colors().editor_background)
541 .child(body.clone()), // .children(contexts.iter().map(|context| context.render(cx))),
542 )
543 .into_any(),
544 ChatMessage::Assistant(AssistantMessage {
545 id,
546 body,
547 error,
548 tool_calls,
549 ..
550 }) => {
551 let assistant_body = if body.text.is_empty() && !tool_calls.is_empty() {
552 div()
553 } else {
554 div().p_2().child(body.element(ElementId::from(id.0), cx))
555 };
556
557 div()
558 .when(!is_last, |element| element.mb_2())
559 .child(
560 div()
561 .p_2()
562 .child(Label::new("Assistant").color(Color::Modified)),
563 )
564 .child(assistant_body)
565 .child(self.render_error(error.clone(), ix, cx))
566 .children(tool_calls.iter().map(|tool_call| {
567 let result = &tool_call.result;
568 let name = tool_call.name.clone();
569 match result {
570 Some(result) => {
571 div().p_2().child(result.into_any_element(&name)).into_any()
572 }
573 None => div()
574 .p_2()
575 .child(Label::new(name).color(Color::Modified))
576 .child("Running...")
577 .into_any(),
578 }
579 }))
580 .into_any()
581 }
582 }
583 }
584
585 fn completion_messages(&self, cx: &mut WindowContext) -> Vec<CompletionMessage> {
586 let mut completion_messages = Vec::new();
587
588 for message in &self.messages {
589 match message {
590 ChatMessage::User(UserMessage { body, contexts, .. }) => {
591 // setup context for model
592 contexts.iter().for_each(|context| {
593 completion_messages.extend(context.completion_messages(cx))
594 });
595
596 // Show user's message last so that the assistant is grounded in the user's request
597 completion_messages.push(CompletionMessage::User {
598 content: body.read(cx).text(cx),
599 });
600 }
601 ChatMessage::Assistant(AssistantMessage {
602 body, tool_calls, ..
603 }) => {
604 // In no case do we want to send an empty message. This shouldn't happen, but we might as well
605 // not break the Chat API if it does.
606 if body.text.is_empty() && tool_calls.is_empty() {
607 continue;
608 }
609
610 let tool_calls_from_assistant = tool_calls
611 .iter()
612 .map(|tool_call| ToolCall {
613 content: ToolCallContent::Function {
614 function: FunctionContent {
615 name: tool_call.name.clone(),
616 arguments: tool_call.arguments.clone(),
617 },
618 },
619 id: tool_call.id.clone(),
620 })
621 .collect();
622
623 completion_messages.push(CompletionMessage::Assistant {
624 content: Some(body.text.to_string()),
625 tool_calls: tool_calls_from_assistant,
626 });
627
628 for tool_call in tool_calls {
629 // todo!(): we should not be sending when the tool is still running / has no result
630 // For now I'm going to have to assume we send an empty string because otherwise
631 // the Chat API will break -- there is a required message for every tool call by ID
632 let content = match &tool_call.result {
633 Some(result) => result.format(&tool_call.name),
634 None => "".to_string(),
635 };
636
637 completion_messages.push(CompletionMessage::Tool {
638 content,
639 tool_call_id: tool_call.id.clone(),
640 });
641 }
642 }
643 }
644 }
645
646 completion_messages
647 }
648
649 fn render_model_dropdown(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
650 let this = cx.view().downgrade();
651 div().h_flex().justify_end().child(
652 div().w_32().child(
653 popover_menu("user-menu")
654 .menu(move |cx| {
655 ContextMenu::build(cx, |mut menu, cx| {
656 for model in CompletionProvider::get(cx).available_models() {
657 menu = menu.custom_entry(
658 {
659 let model = model.clone();
660 move |_| Label::new(model.clone()).into_any_element()
661 },
662 {
663 let this = this.clone();
664 move |cx| {
665 _ = this.update(cx, |this, cx| {
666 this.model = model.clone();
667 cx.notify();
668 });
669 }
670 },
671 );
672 }
673 menu
674 })
675 .into()
676 })
677 .trigger(
678 ButtonLike::new("active-model")
679 .child(
680 h_flex()
681 .w_full()
682 .gap_0p5()
683 .child(
684 div()
685 .overflow_x_hidden()
686 .flex_grow()
687 .whitespace_nowrap()
688 .child(Label::new(self.model.clone())),
689 )
690 .child(div().child(
691 Icon::new(IconName::ChevronDown).color(Color::Muted),
692 )),
693 )
694 .style(ButtonStyle::Subtle)
695 .tooltip(move |cx| Tooltip::text("Change Model", cx)),
696 )
697 .anchor(gpui::AnchorCorner::TopRight),
698 ),
699 )
700 }
701}
702
703impl Render for AssistantChat {
704 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
705 div()
706 .relative()
707 .flex_1()
708 .v_flex()
709 .key_context("AssistantChat")
710 .on_action(cx.listener(Self::cancel))
711 .text_color(Color::Default.color(cx))
712 .child(self.render_model_dropdown(cx))
713 .child(list(self.list_state.clone()).flex_1())
714 }
715}
716
717#[derive(Copy, Clone, Eq, PartialEq)]
718struct MessageId(usize);
719
720impl MessageId {
721 fn post_inc(&mut self) -> Self {
722 let id = *self;
723 self.0 += 1;
724 id
725 }
726}
727
728enum ChatMessage {
729 User(UserMessage),
730 Assistant(AssistantMessage),
731}
732
733impl ChatMessage {
734 fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
735 match self {
736 ChatMessage::User(UserMessage { body, .. }) => Some(body.focus_handle(cx)),
737 ChatMessage::Assistant(_) => None,
738 }
739 }
740}
741
742struct UserMessage {
743 id: MessageId,
744 body: View<Editor>,
745 contexts: Vec<AssistantContext>,
746}
747
748struct AssistantMessage {
749 id: MessageId,
750 body: RichText,
751 tool_calls: Vec<ToolFunctionCall>,
752 error: Option<SharedString>,
753}
754
755// Since we're swapping out for direct query usage, we might not need to use this injected context
756// It will be useful though for when the user _definitely_ wants the model to see a specific file,
757// query, error, etc.
758#[allow(dead_code)]
759enum AssistantContext {
760 Codebase(View<CodebaseContext>),
761}
762
763#[allow(dead_code)]
764struct CodebaseExcerpt {
765 element_id: ElementId,
766 path: SharedString,
767 text: SharedString,
768 score: f32,
769 expanded: bool,
770}
771
772impl AssistantContext {
773 #[allow(dead_code)]
774 fn render(&self, _cx: &mut ViewContext<AssistantChat>) -> AnyElement {
775 match self {
776 AssistantContext::Codebase(context) => context.clone().into_any_element(),
777 }
778 }
779
780 fn completion_messages(&self, cx: &WindowContext) -> Vec<CompletionMessage> {
781 match self {
782 AssistantContext::Codebase(context) => context.read(cx).completion_messages(),
783 }
784 }
785}
786
787enum CodebaseContext {
788 Pending { _task: Task<()> },
789 Done(Result<Vec<CodebaseExcerpt>>),
790}
791
792impl CodebaseContext {
793 fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext<Self>) {
794 if let CodebaseContext::Done(Ok(excerpts)) = self {
795 if let Some(excerpt) = excerpts
796 .iter_mut()
797 .find(|excerpt| excerpt.element_id == element_id)
798 {
799 excerpt.expanded = !excerpt.expanded;
800 cx.notify();
801 }
802 }
803 }
804}
805
806impl Render for CodebaseContext {
807 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
808 match self {
809 CodebaseContext::Pending { .. } => div()
810 .h_flex()
811 .items_center()
812 .gap_1()
813 .child(Icon::new(IconName::Ai).color(Color::Muted).into_element())
814 .child("Searching codebase..."),
815 CodebaseContext::Done(Ok(excerpts)) => {
816 div()
817 .v_flex()
818 .gap_2()
819 .children(excerpts.iter().map(|excerpt| {
820 let expanded = excerpt.expanded;
821 let element_id = excerpt.element_id.clone();
822
823 CollapsibleContainer::new(element_id.clone(), expanded)
824 .start_slot(
825 h_flex()
826 .gap_1()
827 .child(Icon::new(IconName::File).color(Color::Muted))
828 .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
829 )
830 .on_click(cx.listener(move |this, _, cx| {
831 this.toggle_expanded(element_id.clone(), cx);
832 }))
833 .child(
834 div()
835 .p_2()
836 .rounded_md()
837 .bg(cx.theme().colors().editor_background)
838 .child(
839 excerpt.text.clone(), // todo!(): Show as an editor block
840 ),
841 )
842 }))
843 }
844 CodebaseContext::Done(Err(error)) => div().child(error.to_string()),
845 }
846 }
847}
848
849impl CodebaseContext {
850 #[allow(dead_code)]
851 fn new(
852 query: impl 'static + Future<Output = Result<String>>,
853 populated: oneshot::Sender<bool>,
854 project_index: Model<ProjectIndex>,
855 fs: Arc<dyn Fs>,
856 cx: &mut ViewContext<Self>,
857 ) -> Self {
858 let query = query.boxed_local();
859 let _task = cx.spawn(|this, mut cx| async move {
860 let result = async {
861 let query = query.await?;
862 let results = this
863 .update(&mut cx, |_this, cx| {
864 project_index.read(cx).search(&query, 16, cx)
865 })?
866 .await;
867
868 let excerpts = results.into_iter().map(|result| {
869 let abs_path = result
870 .worktree
871 .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
872 let fs = fs.clone();
873
874 async move {
875 let path = result.path.clone();
876 let text = fs.load(&abs_path?).await?;
877 // todo!("what should we do with stale ranges?");
878 let range = cmp::min(result.range.start, text.len())
879 ..cmp::min(result.range.end, text.len());
880
881 let text = SharedString::from(text[range].to_string());
882
883 anyhow::Ok(CodebaseExcerpt {
884 element_id: ElementId::Name(nanoid::nanoid!().into()),
885 path: path.to_string_lossy().to_string().into(),
886 text,
887 score: result.score,
888 expanded: false,
889 })
890 }
891 });
892
893 anyhow::Ok(
894 futures::future::join_all(excerpts)
895 .await
896 .into_iter()
897 .filter_map(|result| result.log_err())
898 .collect(),
899 )
900 }
901 .await;
902
903 this.update(&mut cx, |this, cx| {
904 this.populate(result, populated, cx);
905 })
906 .ok();
907 });
908
909 Self::Pending { _task }
910 }
911
912 #[allow(dead_code)]
913 fn populate(
914 &mut self,
915 result: Result<Vec<CodebaseExcerpt>>,
916 populated: oneshot::Sender<bool>,
917 cx: &mut ViewContext<Self>,
918 ) {
919 let success = result.is_ok();
920 *self = Self::Done(result);
921 populated.send(success).ok();
922 cx.notify();
923 }
924
925 fn completion_messages(&self) -> Vec<CompletionMessage> {
926 // One system message for the whole batch of excerpts:
927
928 // Semantic search results for user query:
929 //
930 // Excerpt from $path:
931 // ~~~
932 // `text`
933 // ~~~
934 //
935 // Excerpt from $path:
936
937 match self {
938 CodebaseContext::Done(Ok(excerpts)) => {
939 if excerpts.is_empty() {
940 return Vec::new();
941 }
942
943 let mut body = "Semantic search results for user query:\n".to_string();
944
945 for excerpt in excerpts {
946 body.push_str("Excerpt from ");
947 body.push_str(excerpt.path.as_ref());
948 body.push_str(", score ");
949 body.push_str(&excerpt.score.to_string());
950 body.push_str(":\n");
951 body.push_str("~~~\n");
952 body.push_str(excerpt.text.as_ref());
953 body.push_str("~~~\n");
954 }
955
956 vec![CompletionMessage::System { content: body }]
957 }
958 _ => vec![],
959 }
960 }
961}