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