1use std::{path::PathBuf, sync::Arc};
  2
  3use anyhow::{anyhow, Result};
  4use assistant_tooling::{AttachmentOutput, LanguageModelAttachment, ProjectContext};
  5use editor::Editor;
  6use gpui::{Render, Task, View, WeakModel, WeakView};
  7use language::Buffer;
  8use project::ProjectPath;
  9use serde::{Deserialize, Serialize};
 10use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
 11use util::maybe;
 12use workspace::Workspace;
 13
 14#[derive(Serialize, Deserialize)]
 15pub struct ActiveEditorAttachment {
 16    #[serde(skip)]
 17    buffer: Option<WeakModel<Buffer>>,
 18    path: Option<PathBuf>,
 19}
 20
 21pub struct FileAttachmentView {
 22    project_path: Option<ProjectPath>,
 23    buffer: Option<WeakModel<Buffer>>,
 24    error: Option<anyhow::Error>,
 25}
 26
 27impl Render for FileAttachmentView {
 28    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 29        if let Some(error) = &self.error {
 30            return div().child(error.to_string()).into_any_element();
 31        }
 32
 33        let filename: SharedString = self
 34            .project_path
 35            .as_ref()
 36            .and_then(|p| p.path.file_name()?.to_str())
 37            .unwrap_or("Untitled")
 38            .to_string()
 39            .into();
 40
 41        ButtonLike::new("file-attachment")
 42            .child(
 43                h_flex()
 44                    .gap_1()
 45                    .bg(cx.theme().colors().editor_background)
 46                    .rounded_md()
 47                    .child(ui::Icon::new(IconName::File))
 48                    .child(filename.clone()),
 49            )
 50            .tooltip(move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx))
 51            .into_any_element()
 52    }
 53}
 54
 55impl AttachmentOutput for FileAttachmentView {
 56    fn generate(&self, project: &mut ProjectContext, cx: &mut WindowContext) -> String {
 57        if let Some(path) = &self.project_path {
 58            project.add_file(path.clone());
 59            return format!("current file: {}", path.path.display());
 60        }
 61
 62        if let Some(buffer) = self.buffer.as_ref().and_then(|buffer| buffer.upgrade()) {
 63            return format!("current untitled buffer text:\n{}", buffer.read(cx).text());
 64        }
 65
 66        String::new()
 67    }
 68}
 69
 70pub struct ActiveEditorAttachmentTool {
 71    workspace: WeakView<Workspace>,
 72}
 73
 74impl ActiveEditorAttachmentTool {
 75    pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
 76        Self { workspace }
 77    }
 78}
 79
 80impl LanguageModelAttachment for ActiveEditorAttachmentTool {
 81    type Output = ActiveEditorAttachment;
 82    type View = FileAttachmentView;
 83
 84    fn name(&self) -> Arc<str> {
 85        "active-editor-attachment".into()
 86    }
 87
 88    fn run(&self, cx: &mut WindowContext) -> Task<Result<ActiveEditorAttachment>> {
 89        Task::ready(maybe!({
 90            let active_buffer = self
 91                .workspace
 92                .update(cx, |workspace, cx| {
 93                    workspace
 94                        .active_item(cx)
 95                        .and_then(|item| Some(item.act_as::<Editor>(cx)?.read(cx).buffer().clone()))
 96                })?
 97                .ok_or_else(|| anyhow!("no active buffer"))?;
 98
 99            let buffer = active_buffer.read(cx);
100
101            if let Some(buffer) = buffer.as_singleton() {
102                let path = project::File::from_dyn(buffer.read(cx).file())
103                    .and_then(|file| file.worktree.read(cx).absolutize(&file.path).ok());
104                return Ok(ActiveEditorAttachment {
105                    buffer: Some(buffer.downgrade()),
106                    path,
107                });
108            } else {
109                Err(anyhow!("no active buffer"))
110            }
111        }))
112    }
113
114    fn view(
115        &self,
116        output: Result<ActiveEditorAttachment>,
117        cx: &mut WindowContext,
118    ) -> View<Self::View> {
119        let error;
120        let project_path;
121        let buffer;
122        match output {
123            Ok(output) => {
124                error = None;
125                let workspace = self.workspace.upgrade().unwrap();
126                let project = workspace.read(cx).project();
127                project_path = output
128                    .path
129                    .and_then(|path| project.read(cx).project_path_for_absolute_path(&path, cx));
130                buffer = output.buffer;
131            }
132            Err(err) => {
133                error = Some(err);
134                buffer = None;
135                project_path = None;
136            }
137        }
138        cx.new_view(|_cx| FileAttachmentView {
139            project_path,
140            buffer,
141            error,
142        })
143    }
144}