1use crate::{
2 slash_command::SlashCommandCompletionProvider, AssistantPanel, CompletionProvider,
3 InlineAssist, InlineAssistant, LanguageModelRequest, LanguageModelRequestMessage, Role,
4};
5use anyhow::{anyhow, Result};
6use assistant_slash_command::SlashCommandRegistry;
7use chrono::{DateTime, Utc};
8use collections::HashMap;
9use editor::{actions::Tab, CurrentLineHighlight, Editor, EditorEvent};
10use futures::{
11 future::{self, BoxFuture, Shared},
12 FutureExt,
13};
14use fuzzy::StringMatchCandidate;
15use gpui::{
16 actions, point, size, AnyElement, AppContext, BackgroundExecutor, Bounds, DevicePixels,
17 EventEmitter, Global, PromptLevel, ReadGlobal, Subscription, Task, TitlebarOptions,
18 UpdateGlobal, View, WindowBounds, WindowHandle, WindowOptions,
19};
20use heed::{types::SerdeBincode, Database, RoTxn};
21use language::{language_settings::SoftWrap, Buffer, LanguageRegistry};
22use parking_lot::RwLock;
23use picker::{Picker, PickerDelegate};
24use rope::Rope;
25use serde::{Deserialize, Serialize};
26use std::{
27 future::Future,
28 path::PathBuf,
29 sync::{atomic::AtomicBool, Arc},
30 time::Duration,
31};
32use ui::{
33 div, prelude::*, IconButtonShape, ListHeader, ListItem, ListItemSpacing, ListSubHeader,
34 ParentElement, Render, SharedString, Styled, TitleBar, Tooltip, ViewContext, VisualContext,
35};
36use util::{paths::PROMPTS_DIR, ResultExt, TryFutureExt};
37use uuid::Uuid;
38use workspace::Workspace;
39
40actions!(
41 prompt_library,
42 [NewPrompt, DeletePrompt, ToggleDefaultPrompt]
43);
44
45/// Init starts loading the PromptStore in the background and assigns
46/// a shared future to a global.
47pub fn init(cx: &mut AppContext) {
48 let db_path = PROMPTS_DIR.join("prompts-library-db.0.mdb");
49 let prompt_store_future = PromptStore::new(db_path, cx.background_executor().clone())
50 .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
51 .boxed()
52 .shared();
53 cx.set_global(GlobalPromptStore(prompt_store_future))
54}
55
56/// This function opens a new prompt library window if one doesn't exist already.
57/// If one exists, it brings it to the foreground.
58///
59/// Note that, when opening a new window, this waits for the PromptStore to be
60/// initialized. If it was initialized successfully, it returns a window handle
61/// to a prompt library.
62pub fn open_prompt_library(
63 language_registry: Arc<LanguageRegistry>,
64 cx: &mut AppContext,
65) -> Task<Result<WindowHandle<PromptLibrary>>> {
66 let existing_window = cx
67 .windows()
68 .into_iter()
69 .find_map(|window| window.downcast::<PromptLibrary>());
70 if let Some(existing_window) = existing_window {
71 existing_window
72 .update(cx, |_, cx| cx.activate_window())
73 .ok();
74 Task::ready(Ok(existing_window))
75 } else {
76 let store = PromptStore::global(cx);
77 cx.spawn(|cx| async move {
78 let store = store.await?;
79 cx.update(|cx| {
80 let bounds = Bounds::centered(
81 None,
82 size(DevicePixels::from(1024), DevicePixels::from(768)),
83 cx,
84 );
85 cx.open_window(
86 WindowOptions {
87 titlebar: Some(TitlebarOptions {
88 title: Some("Prompt Library".into()),
89 appears_transparent: true,
90 traffic_light_position: Some(point(px(9.0), px(9.0))),
91 }),
92 window_bounds: Some(WindowBounds::Windowed(bounds)),
93 ..Default::default()
94 },
95 |cx| cx.new_view(|cx| PromptLibrary::new(store, language_registry, cx)),
96 )
97 })
98 })
99 }
100}
101
102pub struct PromptLibrary {
103 store: Arc<PromptStore>,
104 language_registry: Arc<LanguageRegistry>,
105 prompt_editors: HashMap<PromptId, PromptEditor>,
106 active_prompt_id: Option<PromptId>,
107 picker: View<Picker<PromptPickerDelegate>>,
108 pending_load: Task<()>,
109 _subscriptions: Vec<Subscription>,
110}
111
112struct PromptEditor {
113 editor: View<Editor>,
114 token_count: Option<usize>,
115 pending_token_count: Task<Option<()>>,
116 next_body_to_save: Option<Rope>,
117 pending_save: Option<Task<Option<()>>>,
118 _subscription: Subscription,
119}
120
121struct PromptPickerDelegate {
122 store: Arc<PromptStore>,
123 selected_index: usize,
124 entries: Vec<PromptPickerEntry>,
125}
126
127enum PromptPickerEvent {
128 Selected { prompt_id: Option<PromptId> },
129 Confirmed { prompt_id: PromptId },
130 Deleted { prompt_id: PromptId },
131 ToggledDefault { prompt_id: PromptId },
132}
133
134#[derive(Debug)]
135enum PromptPickerEntry {
136 DefaultPromptsHeader,
137 DefaultPromptsEmpty,
138 AllPromptsHeader,
139 AllPromptsEmpty,
140 Prompt(PromptMetadata),
141}
142
143impl PromptPickerEntry {
144 fn prompt_id(&self) -> Option<PromptId> {
145 match self {
146 PromptPickerEntry::Prompt(metadata) => Some(metadata.id),
147 _ => None,
148 }
149 }
150}
151
152impl EventEmitter<PromptPickerEvent> for Picker<PromptPickerDelegate> {}
153
154impl PickerDelegate for PromptPickerDelegate {
155 type ListItem = AnyElement;
156
157 fn match_count(&self) -> usize {
158 self.entries.len()
159 }
160
161 fn selected_index(&self) -> usize {
162 self.selected_index
163 }
164
165 fn set_selected_index(&mut self, ix: usize, cx: &mut ViewContext<Picker<Self>>) {
166 self.selected_index = ix;
167 let prompt_id = if let Some(PromptPickerEntry::Prompt(prompt)) =
168 self.entries.get(self.selected_index)
169 {
170 Some(prompt.id)
171 } else {
172 None
173 };
174 cx.emit(PromptPickerEvent::Selected { prompt_id });
175 }
176
177 fn placeholder_text(&self, _cx: &mut WindowContext) -> Arc<str> {
178 "Search...".into()
179 }
180
181 fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
182 let search = self.store.search(query);
183 let prev_prompt_id = self
184 .entries
185 .get(self.selected_index)
186 .and_then(|mat| mat.prompt_id());
187 cx.spawn(|this, mut cx| async move {
188 let (entries, selected_index) = cx
189 .background_executor()
190 .spawn(async move {
191 let prompts = search.await;
192 let (default_prompts, prompts) = prompts
193 .into_iter()
194 .partition::<Vec<_>, _>(|prompt| prompt.default);
195
196 let mut entries = Vec::new();
197 entries.push(PromptPickerEntry::DefaultPromptsHeader);
198 if default_prompts.is_empty() {
199 entries.push(PromptPickerEntry::DefaultPromptsEmpty);
200 } else {
201 entries.extend(default_prompts.into_iter().map(PromptPickerEntry::Prompt));
202 }
203
204 entries.push(PromptPickerEntry::AllPromptsHeader);
205 if prompts.is_empty() {
206 entries.push(PromptPickerEntry::AllPromptsEmpty);
207 } else {
208 entries.extend(prompts.into_iter().map(PromptPickerEntry::Prompt));
209 }
210
211 let selected_index = prev_prompt_id
212 .and_then(|prev_prompt_id| {
213 entries
214 .iter()
215 .position(|entry| entry.prompt_id() == Some(prev_prompt_id))
216 })
217 .or_else(|| entries.iter().position(|entry| entry.prompt_id().is_some()))
218 .unwrap_or(0);
219 (entries, selected_index)
220 })
221 .await;
222
223 this.update(&mut cx, |this, cx| {
224 this.delegate.entries = entries;
225 this.delegate.set_selected_index(selected_index, cx);
226 cx.notify();
227 })
228 .ok();
229 })
230 }
231
232 fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
233 if let Some(PromptPickerEntry::Prompt(prompt)) = self.entries.get(self.selected_index) {
234 cx.emit(PromptPickerEvent::Confirmed {
235 prompt_id: prompt.id,
236 });
237 }
238 }
239
240 fn dismissed(&mut self, _cx: &mut ViewContext<Picker<Self>>) {}
241
242 fn render_match(
243 &self,
244 ix: usize,
245 selected: bool,
246 cx: &mut ViewContext<Picker<Self>>,
247 ) -> Option<Self::ListItem> {
248 let prompt = self.entries.get(ix)?;
249 let element = match prompt {
250 PromptPickerEntry::DefaultPromptsHeader => ListHeader::new("Default Prompts")
251 .inset(true)
252 .start_slot(Icon::new(IconName::ZedAssistant))
253 .selected(selected)
254 .into_any_element(),
255 PromptPickerEntry::DefaultPromptsEmpty => {
256 ListSubHeader::new("Star a prompt to add it to the default context")
257 .inset(true)
258 .selected(selected)
259 .into_any_element()
260 }
261 PromptPickerEntry::AllPromptsHeader => ListHeader::new("All Prompts")
262 .inset(true)
263 .start_slot(Icon::new(IconName::Library))
264 .selected(selected)
265 .into_any_element(),
266 PromptPickerEntry::AllPromptsEmpty => ListSubHeader::new("No prompts")
267 .inset(true)
268 .selected(selected)
269 .into_any_element(),
270 PromptPickerEntry::Prompt(prompt) => {
271 let default = prompt.default;
272 let prompt_id = prompt.id;
273 ListItem::new(ix)
274 .inset(true)
275 .spacing(ListItemSpacing::Sparse)
276 .selected(selected)
277 .child(Label::new(
278 prompt.title.clone().unwrap_or("Untitled".into()),
279 ))
280 .end_hover_slot(
281 h_flex()
282 .gap_2()
283 .child(
284 IconButton::new("delete-prompt", IconName::Trash)
285 .shape(IconButtonShape::Square)
286 .tooltip(move |cx| Tooltip::text("Delete Prompt", cx))
287 .on_click(cx.listener(move |_, _, cx| {
288 cx.emit(PromptPickerEvent::Deleted { prompt_id })
289 })),
290 )
291 .child(
292 IconButton::new(
293 "toggle-default-prompt",
294 if default {
295 IconName::ZedAssistantFilled
296 } else {
297 IconName::ZedAssistant
298 },
299 )
300 .shape(IconButtonShape::Square)
301 .tooltip(move |cx| {
302 Tooltip::text(
303 if default {
304 "Remove from Default Prompt"
305 } else {
306 "Add to Default Prompt"
307 },
308 cx,
309 )
310 })
311 .on_click(cx.listener(
312 move |_, _, cx| {
313 cx.emit(PromptPickerEvent::ToggledDefault { prompt_id })
314 },
315 )),
316 ),
317 )
318 .into_any_element()
319 }
320 };
321 Some(element)
322 }
323}
324
325impl PromptLibrary {
326 fn new(
327 store: Arc<PromptStore>,
328 language_registry: Arc<LanguageRegistry>,
329 cx: &mut ViewContext<Self>,
330 ) -> Self {
331 let delegate = PromptPickerDelegate {
332 store: store.clone(),
333 selected_index: 0,
334 entries: Vec::new(),
335 };
336
337 let picker = cx.new_view(|cx| {
338 let picker = Picker::list(delegate, cx).modal(false).max_height(None);
339 picker.focus(cx);
340 picker
341 });
342 Self {
343 store: store.clone(),
344 language_registry,
345 prompt_editors: HashMap::default(),
346 active_prompt_id: None,
347 pending_load: Task::ready(()),
348 _subscriptions: vec![cx.subscribe(&picker, Self::handle_picker_event)],
349 picker,
350 }
351 }
352
353 fn handle_picker_event(
354 &mut self,
355 _: View<Picker<PromptPickerDelegate>>,
356 event: &PromptPickerEvent,
357 cx: &mut ViewContext<Self>,
358 ) {
359 match event {
360 PromptPickerEvent::Selected { prompt_id } => {
361 if let Some(prompt_id) = *prompt_id {
362 self.load_prompt(prompt_id, false, cx);
363 } else {
364 self.focus_picker(&Default::default(), cx);
365 }
366 }
367 PromptPickerEvent::Confirmed { prompt_id } => {
368 self.load_prompt(*prompt_id, true, cx);
369 }
370 PromptPickerEvent::ToggledDefault { prompt_id } => {
371 self.toggle_default_for_prompt(*prompt_id, cx);
372 }
373 PromptPickerEvent::Deleted { prompt_id } => {
374 self.delete_prompt(*prompt_id, cx);
375 }
376 }
377 }
378
379 pub fn new_prompt(&mut self, cx: &mut ViewContext<Self>) {
380 // If we already have an untitled prompt, use that instead
381 // of creating a new one.
382 if let Some(metadata) = self.store.first() {
383 if metadata.title.is_none() {
384 self.load_prompt(metadata.id, true, cx);
385 return;
386 }
387 }
388
389 let prompt_id = PromptId::new();
390 let save = self.store.save(prompt_id, None, false, "".into());
391 self.picker.update(cx, |picker, cx| picker.refresh(cx));
392 cx.spawn(|this, mut cx| async move {
393 save.await?;
394 this.update(&mut cx, |this, cx| this.load_prompt(prompt_id, true, cx))
395 })
396 .detach_and_log_err(cx);
397 }
398
399 pub fn save_prompt(&mut self, prompt_id: PromptId, cx: &mut ViewContext<Self>) {
400 const SAVE_THROTTLE: Duration = Duration::from_millis(500);
401
402 let prompt_metadata = self.store.metadata(prompt_id).unwrap();
403 let prompt_editor = self.prompt_editors.get_mut(&prompt_id).unwrap();
404 let body = prompt_editor.editor.update(cx, |editor, cx| {
405 editor
406 .buffer()
407 .read(cx)
408 .as_singleton()
409 .unwrap()
410 .read(cx)
411 .as_rope()
412 .clone()
413 });
414
415 let store = self.store.clone();
416 let executor = cx.background_executor().clone();
417
418 prompt_editor.next_body_to_save = Some(body);
419 if prompt_editor.pending_save.is_none() {
420 prompt_editor.pending_save = Some(cx.spawn(|this, mut cx| {
421 async move {
422 loop {
423 let next_body_to_save = this.update(&mut cx, |this, _| {
424 this.prompt_editors
425 .get_mut(&prompt_id)?
426 .next_body_to_save
427 .take()
428 })?;
429
430 if let Some(body) = next_body_to_save {
431 let title = title_from_body(body.chars_at(0));
432 store
433 .save(prompt_id, title, prompt_metadata.default, body)
434 .await
435 .log_err();
436 this.update(&mut cx, |this, cx| {
437 this.picker.update(cx, |picker, cx| picker.refresh(cx));
438 cx.notify();
439 })?;
440
441 executor.timer(SAVE_THROTTLE).await;
442 } else {
443 break;
444 }
445 }
446
447 this.update(&mut cx, |this, _cx| {
448 if let Some(prompt_editor) = this.prompt_editors.get_mut(&prompt_id) {
449 prompt_editor.pending_save = None;
450 }
451 })
452 }
453 .log_err()
454 }));
455 }
456 }
457
458 pub fn delete_active_prompt(&mut self, cx: &mut ViewContext<Self>) {
459 if let Some(active_prompt_id) = self.active_prompt_id {
460 self.delete_prompt(active_prompt_id, cx);
461 }
462 }
463
464 pub fn toggle_default_for_active_prompt(&mut self, cx: &mut ViewContext<Self>) {
465 if let Some(active_prompt_id) = self.active_prompt_id {
466 self.toggle_default_for_prompt(active_prompt_id, cx);
467 }
468 }
469
470 pub fn toggle_default_for_prompt(&mut self, prompt_id: PromptId, cx: &mut ViewContext<Self>) {
471 if let Some(prompt_metadata) = self.store.metadata(prompt_id) {
472 self.store
473 .save_metadata(prompt_id, prompt_metadata.title, !prompt_metadata.default)
474 .detach_and_log_err(cx);
475 self.picker.update(cx, |picker, cx| picker.refresh(cx));
476 cx.notify();
477 }
478 }
479
480 pub fn load_prompt(&mut self, prompt_id: PromptId, focus: bool, cx: &mut ViewContext<Self>) {
481 if let Some(prompt_editor) = self.prompt_editors.get(&prompt_id) {
482 if focus {
483 prompt_editor
484 .editor
485 .update(cx, |editor, cx| editor.focus(cx));
486 }
487 self.set_active_prompt(Some(prompt_id), cx);
488 } else {
489 let language_registry = self.language_registry.clone();
490 let commands = SlashCommandRegistry::global(cx);
491 let prompt = self.store.load(prompt_id);
492 self.pending_load = cx.spawn(|this, mut cx| async move {
493 let prompt = prompt.await;
494 let markdown = language_registry.language_for_name("Markdown").await;
495 this.update(&mut cx, |this, cx| match prompt {
496 Ok(prompt) => {
497 let buffer = cx.new_model(|cx| {
498 let mut buffer = Buffer::local(prompt, cx);
499 buffer.set_language(markdown.log_err(), cx);
500 buffer.set_language_registry(language_registry);
501 buffer
502 });
503 let editor = cx.new_view(|cx| {
504 let mut editor = Editor::for_buffer(buffer, None, cx);
505 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
506 editor.set_show_gutter(false, cx);
507 editor.set_show_wrap_guides(false, cx);
508 editor.set_show_indent_guides(false, cx);
509 editor.set_current_line_highlight(Some(CurrentLineHighlight::None));
510 editor.set_completion_provider(Box::new(
511 SlashCommandCompletionProvider::new(commands, None, None),
512 ));
513 if focus {
514 editor.focus(cx);
515 }
516 editor
517 });
518 let _subscription =
519 cx.subscribe(&editor, move |this, _editor, event, cx| {
520 this.handle_prompt_editor_event(prompt_id, event, cx)
521 });
522 this.prompt_editors.insert(
523 prompt_id,
524 PromptEditor {
525 editor,
526 next_body_to_save: None,
527 pending_save: None,
528 token_count: None,
529 pending_token_count: Task::ready(None),
530 _subscription,
531 },
532 );
533 this.set_active_prompt(Some(prompt_id), cx);
534 this.count_tokens(prompt_id, cx);
535 }
536 Err(error) => {
537 // TODO: we should show the error in the UI.
538 log::error!("error while loading prompt: {:?}", error);
539 }
540 })
541 .ok();
542 });
543 }
544 }
545
546 fn set_active_prompt(&mut self, prompt_id: Option<PromptId>, cx: &mut ViewContext<Self>) {
547 self.active_prompt_id = prompt_id;
548 self.picker.update(cx, |picker, cx| {
549 if let Some(prompt_id) = prompt_id {
550 if picker
551 .delegate
552 .entries
553 .get(picker.delegate.selected_index())
554 .map_or(true, |old_selected_prompt| {
555 old_selected_prompt.prompt_id() != Some(prompt_id)
556 })
557 {
558 if let Some(ix) = picker
559 .delegate
560 .entries
561 .iter()
562 .position(|mat| mat.prompt_id() == Some(prompt_id))
563 {
564 picker.set_selected_index(ix, true, cx);
565 }
566 }
567 }
568 });
569 cx.notify();
570 }
571
572 pub fn delete_prompt(&mut self, prompt_id: PromptId, cx: &mut ViewContext<Self>) {
573 if let Some(metadata) = self.store.metadata(prompt_id) {
574 let confirmation = cx.prompt(
575 PromptLevel::Warning,
576 &format!(
577 "Are you sure you want to delete {}",
578 metadata.title.unwrap_or("Untitled".into())
579 ),
580 None,
581 &["Delete", "Cancel"],
582 );
583
584 cx.spawn(|this, mut cx| async move {
585 if confirmation.await.ok() == Some(0) {
586 this.update(&mut cx, |this, cx| {
587 if this.active_prompt_id == Some(prompt_id) {
588 this.set_active_prompt(None, cx);
589 }
590 this.prompt_editors.remove(&prompt_id);
591 this.store.delete(prompt_id).detach_and_log_err(cx);
592 this.picker.update(cx, |picker, cx| picker.refresh(cx));
593 cx.notify();
594 })?;
595 }
596 anyhow::Ok(())
597 })
598 .detach_and_log_err(cx);
599 }
600 }
601
602 fn focus_active_prompt(&mut self, _: &Tab, cx: &mut ViewContext<Self>) {
603 if let Some(active_prompt) = self.active_prompt_id {
604 self.prompt_editors[&active_prompt]
605 .editor
606 .update(cx, |editor, cx| editor.focus(cx));
607 cx.stop_propagation();
608 }
609 }
610
611 fn focus_picker(&mut self, _: &menu::Cancel, cx: &mut ViewContext<Self>) {
612 self.picker.update(cx, |picker, cx| picker.focus(cx));
613 }
614
615 pub fn inline_assist(&mut self, _: &InlineAssist, cx: &mut ViewContext<Self>) {
616 let Some(active_prompt_id) = self.active_prompt_id else {
617 cx.propagate();
618 return;
619 };
620
621 let prompt_editor = &self.prompt_editors[&active_prompt_id].editor;
622 let provider = CompletionProvider::global(cx);
623 if provider.is_authenticated() {
624 InlineAssistant::update_global(cx, |assistant, cx| {
625 assistant.assist(&prompt_editor, None, false, cx)
626 })
627 } else {
628 for window in cx.windows() {
629 if let Some(workspace) = window.downcast::<Workspace>() {
630 let panel = workspace
631 .update(cx, |workspace, cx| {
632 cx.activate_window();
633 workspace.focus_panel::<AssistantPanel>(cx)
634 })
635 .ok()
636 .flatten();
637 if panel.is_some() {
638 return;
639 }
640 }
641 }
642 }
643 }
644
645 fn cancel_last_inline_assist(
646 &mut self,
647 _: &editor::actions::Cancel,
648 cx: &mut ViewContext<Self>,
649 ) {
650 let canceled = InlineAssistant::update_global(cx, |assistant, cx| {
651 assistant.cancel_last_inline_assist(cx)
652 });
653 if !canceled {
654 cx.propagate();
655 }
656 }
657
658 fn handle_prompt_editor_event(
659 &mut self,
660 prompt_id: PromptId,
661 event: &EditorEvent,
662 cx: &mut ViewContext<Self>,
663 ) {
664 if let EditorEvent::BufferEdited = event {
665 let prompt_editor = self.prompt_editors.get(&prompt_id).unwrap();
666 let buffer = prompt_editor
667 .editor
668 .read(cx)
669 .buffer()
670 .read(cx)
671 .as_singleton()
672 .unwrap();
673
674 buffer.update(cx, |buffer, cx| {
675 let mut chars = buffer.chars_at(0);
676 match chars.next() {
677 Some('#') => {
678 if chars.next() != Some(' ') {
679 drop(chars);
680 buffer.edit([(1..1, " ")], None, cx);
681 }
682 }
683 Some(' ') => {
684 drop(chars);
685 buffer.edit([(0..0, "#")], None, cx);
686 }
687 _ => {
688 drop(chars);
689 buffer.edit([(0..0, "# ")], None, cx);
690 }
691 }
692 });
693
694 self.save_prompt(prompt_id, cx);
695 self.count_tokens(prompt_id, cx);
696 }
697 }
698
699 fn count_tokens(&mut self, prompt_id: PromptId, cx: &mut ViewContext<Self>) {
700 if let Some(prompt) = self.prompt_editors.get_mut(&prompt_id) {
701 let editor = &prompt.editor.read(cx);
702 let buffer = &editor.buffer().read(cx).as_singleton().unwrap().read(cx);
703 let body = buffer.as_rope().clone();
704 prompt.pending_token_count = cx.spawn(|this, mut cx| {
705 async move {
706 const DEBOUNCE_TIMEOUT: Duration = Duration::from_secs(1);
707
708 cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
709 let token_count = cx
710 .update(|cx| {
711 let provider = CompletionProvider::global(cx);
712 let model = provider.model();
713 provider.count_tokens(
714 LanguageModelRequest {
715 model,
716 messages: vec![LanguageModelRequestMessage {
717 role: Role::System,
718 content: body.to_string(),
719 }],
720 stop: Vec::new(),
721 temperature: 1.,
722 },
723 cx,
724 )
725 })?
726 .await?;
727 this.update(&mut cx, |this, cx| {
728 let prompt_editor = this.prompt_editors.get_mut(&prompt_id).unwrap();
729 prompt_editor.token_count = Some(token_count);
730 cx.notify();
731 })
732 }
733 .log_err()
734 });
735 }
736 }
737
738 fn render_prompt_list(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
739 v_flex()
740 .id("prompt-list")
741 .capture_action(cx.listener(Self::focus_active_prompt))
742 .bg(cx.theme().colors().panel_background)
743 .h_full()
744 .w_1_3()
745 .overflow_x_hidden()
746 .child(
747 h_flex()
748 .p(Spacing::Small.rems(cx))
749 .border_b_1()
750 .border_color(cx.theme().colors().border)
751 .h(TitleBar::height(cx))
752 .w_full()
753 .flex_none()
754 .justify_end()
755 .child(
756 IconButton::new("new-prompt", IconName::Plus)
757 .shape(IconButtonShape::Square)
758 .tooltip(move |cx| Tooltip::for_action("New Prompt", &NewPrompt, cx))
759 .on_click(|_, cx| {
760 cx.dispatch_action(Box::new(NewPrompt));
761 }),
762 ),
763 )
764 .child(div().flex_grow().child(self.picker.clone()))
765 }
766
767 fn render_active_prompt(&mut self, cx: &mut ViewContext<PromptLibrary>) -> gpui::Stateful<Div> {
768 div()
769 .w_2_3()
770 .h_full()
771 .id("prompt-editor")
772 .border_l_1()
773 .border_color(cx.theme().colors().border)
774 .bg(cx.theme().colors().editor_background)
775 .flex_none()
776 .min_w_64()
777 .children(self.active_prompt_id.and_then(|prompt_id| {
778 let prompt_metadata = self.store.metadata(prompt_id)?;
779 let prompt_editor = &self.prompt_editors[&prompt_id];
780 Some(
781 h_flex()
782 .size_full()
783 .items_start()
784 .child(
785 div()
786 .on_action(cx.listener(Self::focus_picker))
787 .on_action(cx.listener(Self::inline_assist))
788 .on_action(cx.listener(Self::cancel_last_inline_assist))
789 .flex_grow()
790 .h_full()
791 .pt(Spacing::Large.rems(cx))
792 .pl(Spacing::Large.rems(cx))
793 .child(prompt_editor.editor.clone()),
794 )
795 .child(
796 v_flex()
797 .w_12()
798 .py(Spacing::Large.rems(cx))
799 .justify_start()
800 .items_center()
801 .gap_4()
802 .child(
803 IconButton::new(
804 "toggle-default-prompt",
805 if prompt_metadata.default {
806 IconName::ZedAssistantFilled
807 } else {
808 IconName::ZedAssistant
809 },
810 )
811 .size(ButtonSize::Large)
812 .shape(IconButtonShape::Square)
813 .tooltip(move |cx| {
814 Tooltip::for_action(
815 if prompt_metadata.default {
816 "Remove from Default Prompt"
817 } else {
818 "Add to Default Prompt"
819 },
820 &ToggleDefaultPrompt,
821 cx,
822 )
823 })
824 .on_click(|_, cx| {
825 cx.dispatch_action(Box::new(ToggleDefaultPrompt));
826 }),
827 )
828 .child(
829 IconButton::new("delete-prompt", IconName::Trash)
830 .shape(IconButtonShape::Square)
831 .tooltip(move |cx| {
832 Tooltip::for_action("Delete Prompt", &DeletePrompt, cx)
833 })
834 .on_click(|_, cx| {
835 cx.dispatch_action(Box::new(DeletePrompt));
836 }),
837 )
838 .children(prompt_editor.token_count.map(|token_count| {
839 h_flex()
840 .justify_center()
841 .child(Label::new(token_count.to_string()))
842 })),
843 ),
844 )
845 }))
846 }
847}
848
849impl Render for PromptLibrary {
850 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
851 h_flex()
852 .id("prompt-manager")
853 .key_context("PromptLibrary")
854 .on_action(cx.listener(|this, &NewPrompt, cx| this.new_prompt(cx)))
855 .on_action(cx.listener(|this, &DeletePrompt, cx| this.delete_active_prompt(cx)))
856 .on_action(cx.listener(|this, &ToggleDefaultPrompt, cx| {
857 this.toggle_default_for_active_prompt(cx)
858 }))
859 .size_full()
860 .overflow_hidden()
861 .child(self.render_prompt_list(cx))
862 .child(self.render_active_prompt(cx))
863 }
864}
865
866#[derive(Clone, Debug, Serialize, Deserialize)]
867pub struct PromptMetadata {
868 pub id: PromptId,
869 pub title: Option<SharedString>,
870 pub default: bool,
871 pub saved_at: DateTime<Utc>,
872}
873
874#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
875pub struct PromptId(Uuid);
876
877impl PromptId {
878 pub fn new() -> PromptId {
879 PromptId(Uuid::new_v4())
880 }
881}
882
883pub struct PromptStore {
884 executor: BackgroundExecutor,
885 env: heed::Env,
886 bodies: Database<SerdeBincode<PromptId>, SerdeBincode<String>>,
887 metadata: Database<SerdeBincode<PromptId>, SerdeBincode<PromptMetadata>>,
888 metadata_cache: RwLock<MetadataCache>,
889}
890
891#[derive(Default)]
892struct MetadataCache {
893 metadata: Vec<PromptMetadata>,
894 metadata_by_id: HashMap<PromptId, PromptMetadata>,
895}
896
897impl MetadataCache {
898 fn from_db(
899 db: Database<SerdeBincode<PromptId>, SerdeBincode<PromptMetadata>>,
900 txn: &RoTxn,
901 ) -> Result<Self> {
902 let mut cache = MetadataCache::default();
903 for result in db.iter(txn)? {
904 let (prompt_id, metadata) = result?;
905 cache.metadata.push(metadata.clone());
906 cache.metadata_by_id.insert(prompt_id, metadata);
907 }
908 cache.sort();
909 Ok(cache)
910 }
911
912 fn insert(&mut self, metadata: PromptMetadata) {
913 self.metadata_by_id.insert(metadata.id, metadata.clone());
914 if let Some(old_metadata) = self.metadata.iter_mut().find(|m| m.id == metadata.id) {
915 *old_metadata = metadata;
916 } else {
917 self.metadata.push(metadata);
918 }
919 self.sort();
920 }
921
922 fn remove(&mut self, id: PromptId) {
923 self.metadata.retain(|metadata| metadata.id != id);
924 self.metadata_by_id.remove(&id);
925 }
926
927 fn sort(&mut self) {
928 self.metadata.sort_unstable_by(|a, b| {
929 a.title
930 .cmp(&b.title)
931 .then_with(|| b.saved_at.cmp(&a.saved_at))
932 });
933 }
934}
935
936impl PromptStore {
937 pub fn global(cx: &AppContext) -> impl Future<Output = Result<Arc<Self>>> {
938 let store = GlobalPromptStore::global(cx).0.clone();
939 async move { store.await.map_err(|err| anyhow!(err)) }
940 }
941
942 pub fn new(db_path: PathBuf, executor: BackgroundExecutor) -> Task<Result<Self>> {
943 executor.spawn({
944 let executor = executor.clone();
945 async move {
946 std::fs::create_dir_all(&db_path)?;
947
948 let db_env = unsafe {
949 heed::EnvOpenOptions::new()
950 .map_size(1024 * 1024 * 1024) // 1GB
951 .max_dbs(2) // bodies and metadata
952 .open(db_path)?
953 };
954
955 let mut txn = db_env.write_txn()?;
956 let bodies = db_env.create_database(&mut txn, Some("bodies"))?;
957 let metadata = db_env.create_database(&mut txn, Some("metadata"))?;
958 let metadata_cache = MetadataCache::from_db(metadata, &txn)?;
959 txn.commit()?;
960
961 Ok(PromptStore {
962 executor,
963 env: db_env,
964 bodies,
965 metadata,
966 metadata_cache: RwLock::new(metadata_cache),
967 })
968 }
969 })
970 }
971
972 pub fn load(&self, id: PromptId) -> Task<Result<String>> {
973 let env = self.env.clone();
974 let bodies = self.bodies;
975 self.executor.spawn(async move {
976 let txn = env.read_txn()?;
977 bodies
978 .get(&txn, &id)?
979 .ok_or_else(|| anyhow!("prompt not found"))
980 })
981 }
982
983 pub fn default_prompt_metadata(&self) -> Vec<PromptMetadata> {
984 return self
985 .metadata_cache
986 .read()
987 .metadata
988 .iter()
989 .filter(|metadata| metadata.default)
990 .cloned()
991 .collect::<Vec<_>>();
992 }
993
994 pub fn delete(&self, id: PromptId) -> Task<Result<()>> {
995 self.metadata_cache.write().remove(id);
996
997 let db_connection = self.env.clone();
998 let bodies = self.bodies;
999 let metadata = self.metadata;
1000
1001 self.executor.spawn(async move {
1002 let mut txn = db_connection.write_txn()?;
1003
1004 metadata.delete(&mut txn, &id)?;
1005 bodies.delete(&mut txn, &id)?;
1006
1007 txn.commit()?;
1008 Ok(())
1009 })
1010 }
1011
1012 fn metadata(&self, id: PromptId) -> Option<PromptMetadata> {
1013 self.metadata_cache.read().metadata_by_id.get(&id).cloned()
1014 }
1015
1016 pub fn id_for_title(&self, title: &str) -> Option<PromptId> {
1017 let metadata_cache = self.metadata_cache.read();
1018 let metadata = metadata_cache
1019 .metadata
1020 .iter()
1021 .find(|metadata| metadata.title.as_ref().map(|title| &***title) == Some(title))?;
1022 Some(metadata.id)
1023 }
1024
1025 pub fn search(&self, query: String) -> Task<Vec<PromptMetadata>> {
1026 let cached_metadata = self.metadata_cache.read().metadata.clone();
1027 let executor = self.executor.clone();
1028 self.executor.spawn(async move {
1029 if query.is_empty() {
1030 cached_metadata
1031 } else {
1032 let candidates = cached_metadata
1033 .iter()
1034 .enumerate()
1035 .filter_map(|(ix, metadata)| {
1036 Some(StringMatchCandidate::new(
1037 ix,
1038 metadata.title.as_ref()?.to_string(),
1039 ))
1040 })
1041 .collect::<Vec<_>>();
1042 let matches = fuzzy::match_strings(
1043 &candidates,
1044 &query,
1045 false,
1046 100,
1047 &AtomicBool::default(),
1048 executor,
1049 )
1050 .await;
1051 matches
1052 .into_iter()
1053 .map(|mat| cached_metadata[mat.candidate_id].clone())
1054 .collect()
1055 }
1056 })
1057 }
1058
1059 fn save(
1060 &self,
1061 id: PromptId,
1062 title: Option<SharedString>,
1063 default: bool,
1064 body: Rope,
1065 ) -> Task<Result<()>> {
1066 let prompt_metadata = PromptMetadata {
1067 id,
1068 title,
1069 default,
1070 saved_at: Utc::now(),
1071 };
1072 self.metadata_cache.write().insert(prompt_metadata.clone());
1073
1074 let db_connection = self.env.clone();
1075 let bodies = self.bodies;
1076 let metadata = self.metadata;
1077
1078 self.executor.spawn(async move {
1079 let mut txn = db_connection.write_txn()?;
1080
1081 metadata.put(&mut txn, &id, &prompt_metadata)?;
1082 bodies.put(&mut txn, &id, &body.to_string())?;
1083
1084 txn.commit()?;
1085
1086 Ok(())
1087 })
1088 }
1089
1090 fn save_metadata(
1091 &self,
1092 id: PromptId,
1093 title: Option<SharedString>,
1094 default: bool,
1095 ) -> Task<Result<()>> {
1096 let prompt_metadata = PromptMetadata {
1097 id,
1098 title,
1099 default,
1100 saved_at: Utc::now(),
1101 };
1102 self.metadata_cache.write().insert(prompt_metadata.clone());
1103
1104 let db_connection = self.env.clone();
1105 let metadata = self.metadata;
1106
1107 self.executor.spawn(async move {
1108 let mut txn = db_connection.write_txn()?;
1109 metadata.put(&mut txn, &id, &prompt_metadata)?;
1110 txn.commit()?;
1111
1112 Ok(())
1113 })
1114 }
1115
1116 fn first(&self) -> Option<PromptMetadata> {
1117 self.metadata_cache.read().metadata.first().cloned()
1118 }
1119}
1120
1121/// Wraps a shared future to a prompt store so it can be assigned as a context global.
1122pub struct GlobalPromptStore(
1123 Shared<BoxFuture<'static, Result<Arc<PromptStore>, Arc<anyhow::Error>>>>,
1124);
1125
1126impl Global for GlobalPromptStore {}
1127
1128fn title_from_body(body: impl IntoIterator<Item = char>) -> Option<SharedString> {
1129 let mut chars = body.into_iter().take_while(|c| *c != '\n').peekable();
1130
1131 let mut level = 0;
1132 while let Some('#') = chars.peek() {
1133 level += 1;
1134 chars.next();
1135 }
1136
1137 if level > 0 {
1138 let title = chars.collect::<String>().trim().to_string();
1139 if title.is_empty() {
1140 None
1141 } else {
1142 Some(title.into())
1143 }
1144 } else {
1145 None
1146 }
1147}