Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Centralize project context provided to the assistant (#11471)

Browse source 
This PR restructures the way that tools and attachments add information
about the current project to a conversation with the assistant. Rather
than each tool call or attachment generating a new tool or system
message containing information about the project, they can all
collectively mutate a new type called a `ProjectContext`, which stores
all of the project data that should be sent to the assistant. That data
is then formatted in a single place, and passed to the assistant in one
system message.

This prevents multiple tools/attachments from including redundant
context.

Release Notes:

- N/A

---------

Co-authored-by: Kyle <kylek@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-05-06 17:01:50 -07:00 • committed by GitHub
parent f2a415135b
commit a64e20ed96
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 841 additions and 518 deletions
Download patch file Download diff file Expand all files Collapse all files

101
crates/assistant2/src/assistant2.rs
View file

@ -4,10 +4,16 @@ mod completion_provider;
mod tools;
pub mod ui;
use crate::{
attachments::ActiveEditorAttachmentTool,
tools::{CreateBufferTool, ProjectIndexTool},
ui::UserOrAssistant,
};
use ::ui::{div, prelude::*, Color, ViewContext};
use anyhow::{Context, Result};
use assistant_tooling::{ToolFunctionCall, ToolRegistry};
use attachments::{ActiveEditorAttachmentTool, UserAttachment, UserAttachmentStore};
use assistant_tooling::{
AttachmentRegistry, ProjectContext, ToolFunctionCall, ToolRegistry, UserAttachment,
};
use client::{proto, Client, UserStore};
use collections::HashMap;
use completion_provider::*;
@ -34,9 +40,6 @@ use workspace::{
pub use assistant_settings::AssistantSettings;
use crate::tools::{CreateBufferTool, ProjectIndexTool};
use crate::ui::UserOrAssistant;
const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
@ -85,10 +88,9 @@ pub fn init(client: Arc<Client>, cx: &mut AppContext) {
});
workspace.register_action(|workspace, _: &DebugProjectIndex, cx| {
if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
if let Some(index) = panel.read(cx).chat.read(cx).project_index.clone() {
let view = cx.new_view(|cx| ProjectIndexDebugView::new(index, cx));
workspace.add_item_to_center(Box::new(view), cx);
}
let index = panel.read(cx).chat.read(cx).project_index.clone();
let view = cx.new_view(|cx| ProjectIndexDebugView::new(index, cx));
workspace.add_item_to_center(Box::new(view), cx);
}
});
},
@ -122,10 +124,7 @@ impl AssistantPanel {
let mut tool_registry = ToolRegistry::new();
tool_registry
.register(
ProjectIndexTool::new(project_index.clone(), project.read(cx).fs().clone()),
cx,
)
.register(ProjectIndexTool::new(project_index.clone()), cx)
.context("failed to register ProjectIndexTool")
.log_err();
tool_registry
@ -136,7 +135,7 @@ impl AssistantPanel {
.context("failed to register CreateBufferTool")
.log_err();
let mut attachment_store = UserAttachmentStore::new();
let mut attachment_store = AttachmentRegistry::new();
attachment_store.register(ActiveEditorAttachmentTool::new(workspace.clone(), cx));
Self::new(
@ -144,7 +143,7 @@ impl AssistantPanel {
Arc::new(tool_registry),
Arc::new(attachment_store),
app_state.user_store.clone(),
Some(project_index),
project_index,
workspace,
cx,
)
@ -155,9 +154,9 @@ impl AssistantPanel {
pub fn new(
language_registry: Arc<LanguageRegistry>,
tool_registry: Arc<ToolRegistry>,
attachment_store: Arc<UserAttachmentStore>,
attachment_store: Arc<AttachmentRegistry>,
user_store: Model<UserStore>,
project_index: Option<Model<ProjectIndex>>,
project_index: Model<ProjectIndex>,
workspace: WeakView<Workspace>,
cx: &mut ViewContext<Self>,
) -> Self {
@ -241,16 +240,16 @@ pub struct AssistantChat {
list_state: ListState,
language_registry: Arc<LanguageRegistry>,
composer_editor: View<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
project_index_button: View<ProjectIndexButton>,
active_file_button: Option<View<ActiveFileButton>>,
user_store: Model<UserStore>,
next_message_id: MessageId,
collapsed_messages: HashMap<MessageId, bool>,
editing_message: Option<EditingMessage>,
pending_completion: Option<Task<()>>,
attachment_store: Arc<UserAttachmentStore>,
tool_registry: Arc<ToolRegistry>,
project_index: Option<Model<ProjectIndex>>,
attachment_registry: Arc<AttachmentRegistry>,
project_index: Model<ProjectIndex>,
}
struct EditingMessage {
@ -263,9 +262,9 @@ impl AssistantChat {
fn new(
language_registry: Arc<LanguageRegistry>,
tool_registry: Arc<ToolRegistry>,
attachment_store: Arc<UserAttachmentStore>,
attachment_registry: Arc<AttachmentRegistry>,
user_store: Model<UserStore>,
project_index: Option<Model<ProjectIndex>>,
project_index: Model<ProjectIndex>,
workspace: WeakView<Workspace>,
cx: &mut ViewContext<Self>,
) -> Self {
@ -281,14 +280,14 @@ impl AssistantChat {
},
);
let project_index_button = project_index.clone().map(|project_index| {
cx.new_view(|cx| ProjectIndexButton::new(project_index, tool_registry.clone(), cx))
let project_index_button = cx.new_view(|cx| {
ProjectIndexButton::new(project_index.clone(), tool_registry.clone(), cx)
});
let active_file_button = match workspace.upgrade() {
Some(workspace) => {
Some(cx.new_view(
|cx| ActiveFileButton::new(attachment_store.clone(), workspace, cx), //
|cx| ActiveFileButton::new(attachment_registry.clone(), workspace, cx), //
))
}
_ => None,
@ -313,7 +312,7 @@ impl AssistantChat {
editing_message: None,
collapsed_messages: HashMap::default(),
pending_completion: None,
attachment_store,
attachment_registry,
tool_registry,
}
}
@ -395,7 +394,7 @@ impl AssistantChat {
let mode = *mode;
self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {
let attachments_task = this.update(&mut cx, |this, cx| {
let attachment_store = this.attachment_store.clone();
let attachment_store = this.attachment_registry.clone();
attachment_store.call_all_attachment_tools(cx)
});
@ -443,7 +442,7 @@ impl AssistantChat {
let mut call_count = 0;
loop {
let complete = async {
let completion = this.update(cx, |this, cx| {
let (tool_definitions, model_name, messages) = this.update(cx, |this, cx| {
this.push_new_assistant_message(cx);
let definitions = if call_count < limit
@ -455,14 +454,22 @@ impl AssistantChat {
};
call_count += 1;
let messages = this.completion_messages(cx);
CompletionProvider::get(cx).complete(
(
definitions,
this.model.clone(),
this.completion_messages(cx),
)
})?;
let messages = messages.await?;
let completion = cx.update(|cx| {
CompletionProvider::get(cx).complete(
model_name,
messages,
Vec::new(),
1.0,
definitions,
tool_definitions,
)
});
@ -765,7 +772,12 @@ impl AssistantChat {
}
}
fn completion_messages(&self, cx: &mut WindowContext) -> Vec<CompletionMessage> {
fn completion_messages(&self, cx: &mut WindowContext) -> Task<Result<Vec<CompletionMessage>>> {
let project_index = self.project_index.read(cx);
let project = project_index.project();
let fs = project_index.fs();
let mut project_context = ProjectContext::new(project, fs);
let mut completion_messages = Vec::new();
for message in &self.messages {
@ -773,12 +785,11 @@ impl AssistantChat {
ChatMessage::User(UserMessage {
body, attachments, ..
}) => {
completion_messages.extend(
attachments
.into_iter()
.filter_map(|attachment| attachment.message.clone())
.map(|content| CompletionMessage::System { content }),
);
for attachment in attachments {
if let Some(content) = attachment.generate(&mut project_context, cx) {
completion_messages.push(CompletionMessage::System { content });
}
}
// Show user's message last so that the assistant is grounded in the user's request
completion_messages.push(CompletionMessage::User {
@ -815,7 +826,9 @@ impl AssistantChat {
for tool_call in tool_calls {
// Every tool call _must_ have a result by ID, otherwise OpenAI will error.
let content = match &tool_call.result {
Some(result) => result.format(&tool_call.name),
Some(result) => {
result.generate(&tool_call.name, &mut project_context, cx)
}
None => "".to_string(),
};
@ -828,7 +841,13 @@ impl AssistantChat {
}
}
completion_messages
let system_message = project_context.generate_system_message(cx);
cx.background_executor().spawn(async move {
let content = system_message.await?;
completion_messages.insert(0, CompletionMessage::System { content });
Ok(completion_messages)
})
}
}

206
crates/assistant2/src/attachments.rs
View file

@ -1,137 +1,18 @@
use std::{
any::TypeId,
sync::{
atomic::{AtomicBool, Ordering::SeqCst},
Arc,
},
};
pub mod active_file;
use anyhow::{anyhow, Result};
use collections::HashMap;
use assistant_tooling::{LanguageModelAttachment, ProjectContext, ToolOutput};
use editor::Editor;
use futures::future::join_all;
use gpui::{AnyView, Render, Task, View, WeakView};
use gpui::{Render, Task, View, WeakModel, WeakView};
use language::Buffer;
use project::ProjectPath;
use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
use util::{maybe, ResultExt};
use util::maybe;
use workspace::Workspace;
/// A collected attachment from running an attachment tool
pub struct UserAttachment {
pub message: Option<String>,
pub view: AnyView,
}
pub struct UserAttachmentStore {
attachment_tools: HashMap<TypeId, DynamicAttachment>,
}
/// Internal representation of an attachment tool to allow us to treat them dynamically
struct DynamicAttachment {
enabled: AtomicBool,
call: Box<dyn Fn(&mut WindowContext) -> Task<Result<UserAttachment>>>,
}
impl UserAttachmentStore {
pub fn new() -> Self {
Self {
attachment_tools: HashMap::default(),
}
}
pub fn register<A: AttachmentTool + 'static>(&mut self, attachment: A) {
let call = Box::new(move |cx: &mut WindowContext| {
let result = attachment.run(cx);
cx.spawn(move |mut cx| async move {
let result: Result<A::Output> = result.await;
let message = A::format(&result);
let view = cx.update(|cx| A::view(result, cx))?;
Ok(UserAttachment {
message,
view: view.into(),
})
})
});
self.attachment_tools.insert(
TypeId::of::<A>(),
DynamicAttachment {
call,
enabled: AtomicBool::new(true),
},
);
}
pub fn set_attachment_tool_enabled<A: AttachmentTool + 'static>(&self, is_enabled: bool) {
if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
attachment.enabled.store(is_enabled, SeqCst);
}
}
pub fn is_attachment_tool_enabled<A: AttachmentTool + 'static>(&self) -> bool {
if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
attachment.enabled.load(SeqCst)
} else {
false
}
}
pub fn call<A: AttachmentTool + 'static>(
&self,
cx: &mut WindowContext,
) -> Task<Result<UserAttachment>> {
let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) else {
return Task::ready(Err(anyhow!("no attachment tool")));
};
(attachment.call)(cx)
}
pub fn call_all_attachment_tools(
self: Arc<Self>,
cx: &mut WindowContext<'_>,
) -> Task<Result<Vec<UserAttachment>>> {
let this = self.clone();
cx.spawn(|mut cx| async move {
let attachment_tasks = cx.update(|cx| {
let mut tasks = Vec::new();
for attachment in this
.attachment_tools
.values()
.filter(|attachment| attachment.enabled.load(SeqCst))
{
tasks.push((attachment.call)(cx))
}
tasks
})?;
let attachments = join_all(attachment_tasks.into_iter()).await;
Ok(attachments
.into_iter()
.filter_map(|attachment| attachment.log_err())
.collect())
})
}
}
pub trait AttachmentTool {
type Output: 'static;
type View: Render;
fn run(&self, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
fn format(output: &Result<Self::Output>) -> Option<String>;
fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View>;
}
pub struct ActiveEditorAttachment {
filename: Arc<str>,
language: Arc<str>,
text: Arc<str>,
buffer: WeakModel<Buffer>,
path: Option<ProjectPath>,
}
pub struct FileAttachmentView {
@ -142,7 +23,13 @@ impl Render for FileAttachmentView {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
match &self.output {
Ok(attachment) => {
let filename = attachment.filename.clone();
let filename: SharedString = attachment
.path
.as_ref()
.and_then(|p| p.path.file_name()?.to_str())
.unwrap_or("Untitled")
.to_string()
.into();
// todo!(): make the button link to the actual file to open
ButtonLike::new("file-attachment")
@ -152,7 +39,7 @@ impl Render for FileAttachmentView {
.bg(cx.theme().colors().editor_background)
.rounded_md()
.child(ui::Icon::new(IconName::File))
.child(filename.to_string()),
.child(filename.clone()),
)
.tooltip({
move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
@ -164,6 +51,20 @@ impl Render for FileAttachmentView {
}
}
impl ToolOutput for FileAttachmentView {
fn generate(&self, project: &mut ProjectContext, cx: &mut WindowContext) -> String {
if let Ok(result) = &self.output {
if let Some(path) = &result.path {
project.add_file(path.clone());
return format!("current file: {}", path.path.display());
} else if let Some(buffer) = result.buffer.upgrade() {
return format!("current untitled buffer text:\n{}", buffer.read(cx).text());
}
}
String::new()
}
}
pub struct ActiveEditorAttachmentTool {
workspace: WeakView<Workspace>,
}
@ -174,7 +75,7 @@ impl ActiveEditorAttachmentTool {
}
}
impl AttachmentTool for ActiveEditorAttachmentTool {
impl LanguageModelAttachment for ActiveEditorAttachmentTool {
type Output = ActiveEditorAttachment;
type View = FileAttachmentView;
@ -191,47 +92,22 @@ impl AttachmentTool for ActiveEditorAttachmentTool {
let buffer = active_buffer.read(cx);
if let Some(singleton) = buffer.as_singleton() {
let singleton = singleton.read(cx);
let filename = singleton
.file()
.map(|file| file.path().to_string_lossy())
.unwrap_or("Untitled".into());
let text = singleton.text();
let language = singleton
.language()
.map(|l| {
let name = l.code_fence_block_name();
name.to_string()
})
.unwrap_or_default();
if let Some(buffer) = buffer.as_singleton() {
let path =
project::File::from_dyn(buffer.read(cx).file()).map(|file| ProjectPath {
worktree_id: file.worktree_id(cx),
path: file.path.clone(),
});
return Ok(ActiveEditorAttachment {
filename: filename.into(),
language: language.into(),
text: text.into(),
buffer: buffer.downgrade(),
path,
});
} else {
Err(anyhow!("no active buffer"))
}
Err(anyhow!("no active buffer"))
}))
}
fn format(output: &Result<Self::Output>) -> Option<String> {
let output = output.as_ref().ok()?;
let filename = &output.filename;
let language = &output.language;
let text = &output.text;
Some(format!(
"User's active file `{filename}`:\n\n```{language}\n{text}```\n\n"
))
}
fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
cx.new_view(|_cx| FileAttachmentView { output })
}

1
crates/assistant2/src/attachments/active_file.rs Normal file
View file

@ -0,0 +1 @@

39
crates/assistant2/src/tools/create_buffer.rs
View file

@ -1,5 +1,5 @@
use anyhow::Result;
use assistant_tooling::LanguageModelTool;
use assistant_tooling::{LanguageModelTool, ProjectContext, ToolOutput};
use editor::Editor;
use gpui::{prelude::*, Model, Task, View, WeakView};
use project::Project;
@ -31,11 +31,9 @@ pub struct CreateBufferInput {
language: String,
}
pub struct CreateBufferOutput {}
impl LanguageModelTool for CreateBufferTool {
type Input = CreateBufferInput;
type Output = CreateBufferOutput;
type Output = ();
type View = CreateBufferView;
fn name(&self) -> String {
@ -83,32 +81,39 @@ impl LanguageModelTool for CreateBufferTool {
})
.log_err();
Ok(CreateBufferOutput {})
Ok(())
}
})
}
fn format(input: &Self::Input, output: &Result<Self::Output>) -> String {
match output {
Ok(_) => format!("Created a new {} buffer", input.language),
Err(err) => format!("Failed to create buffer: {err:?}"),
}
}
fn output_view(
_tool_call_id: String,
_input: Self::Input,
_output: Result<Self::Output>,
input: Self::Input,
output: Result<Self::Output>,
cx: &mut WindowContext,
) -> View<Self::View> {
cx.new_view(|_cx| CreateBufferView {})
cx.new_view(|_cx| CreateBufferView {
language: input.language,
output,
})
}
}
pub struct CreateBufferView {}
pub struct CreateBufferView {
language: String,
output: Result<()>,
}
impl Render for CreateBufferView {
fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
div().child("Opening a buffer")
}
}
impl ToolOutput for CreateBufferView {
fn generate(&self, _: &mut ProjectContext, _: &mut WindowContext) -> String {
match &self.output {
Ok(_) => format!("Created a new {} buffer", self.language),
Err(err) => format!("Failed to create buffer: {err:?}"),
}
}
}

186
crates/assistant2/src/tools/project_index.rs
View file

@ -1,25 +1,18 @@
use anyhow::Result;
use assistant_tooling::LanguageModelTool;
use assistant_tooling::{LanguageModelTool, ToolOutput};
use collections::BTreeMap;
use gpui::{prelude::*, Model, Task};
use project::Fs;
use project::ProjectPath;
use schemars::JsonSchema;
use semantic_index::{ProjectIndex, Status};
use serde::Deserialize;
use std::{collections::HashSet, sync::Arc};
use ui::{
div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString,
WindowContext,
};
use util::ResultExt as _;
use std::{fmt::Write as _, ops::Range};
use ui::{div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
const DEFAULT_SEARCH_LIMIT: usize = 20;
#[derive(Clone)]
pub struct CodebaseExcerpt {
path: SharedString,
text: SharedString,
score: f32,
pub struct ProjectIndexTool {
project_index: Model<ProjectIndex>,
}
// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model.
@ -40,6 +33,11 @@ pub struct ProjectIndexView {
expanded_header: bool,
}
pub struct ProjectIndexOutput {
status: Status,
excerpts: BTreeMap<ProjectPath, Vec<Range<usize>>>,
}
impl ProjectIndexView {
fn new(input: CodebaseQuery, output: Result<ProjectIndexOutput>) -> Self {
let element_id = ElementId::Name(nanoid::nanoid!().into());
@ -71,19 +69,15 @@ impl Render for ProjectIndexView {
Ok(output) => output,
};
let num_files_searched = output.files_searched.len();
let file_count = output.excerpts.len();
let header = h_flex()
.gap_2()
.child(Icon::new(IconName::File))
.child(format!(
"Read {} {}",
num_files_searched,
if num_files_searched == 1 {
"file"
} else {
"files"
}
file_count,
if file_count == 1 { "file" } else { "files" }
));
v_flex().gap_3().child(
@ -102,36 +96,50 @@ impl Render for ProjectIndexView {
.child(Icon::new(IconName::MagnifyingGlass))
.child(Label::new(format!("`{}`", query)).color(Color::Muted)),
)
.child(v_flex().gap_2().children(output.files_searched.iter().map(
|path| {
h_flex()
.gap_2()
.child(Icon::new(IconName::File))
.child(Label::new(path.clone()).color(Color::Muted))
},
))),
.child(
v_flex()
.gap_2()
.children(output.excerpts.keys().map(|path| {
h_flex().gap_2().child(Icon::new(IconName::File)).child(
Label::new(path.path.to_string_lossy().to_string())
.color(Color::Muted),
)
})),
),
),
)
}
}
pub struct ProjectIndexTool {
project_index: Model<ProjectIndex>,
fs: Arc<dyn Fs>,
}
impl ToolOutput for ProjectIndexView {
fn generate(
&self,
context: &mut assistant_tooling::ProjectContext,
_: &mut WindowContext,
) -> String {
match &self.output {
Ok(output) => {
let mut body = "found results in the following paths:\n".to_string();
pub struct ProjectIndexOutput {
excerpts: Vec<CodebaseExcerpt>,
status: Status,
files_searched: HashSet<SharedString>,
for (project_path, ranges) in &output.excerpts {
context.add_excerpts(project_path.clone(), ranges);
writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
}
if output.status != Status::Idle {
body.push_str("Still indexing. Results may be incomplete.\n");
}
body
}
Err(err) => format!("Error: {}", err),
}
}
}
impl ProjectIndexTool {
pub fn new(project_index: Model<ProjectIndex>, fs: Arc<dyn Fs>) -> Self {
// Listen for project index status and update the ProjectIndexTool directly
// TODO: setup a better description based on the user's current codebase.
Self { project_index, fs }
pub fn new(project_index: Model<ProjectIndex>) -> Self {
Self { project_index }
}
}
@ -151,64 +159,42 @@ impl LanguageModelTool for ProjectIndexTool {
fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
let project_index = self.project_index.read(cx);
let status = project_index.status();
let results = project_index.search(
let search = project_index.search(
query.query.clone(),
query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
cx,
);
let fs = self.fs.clone();
cx.spawn(|mut cx| async move {
let search_results = search.await?;
cx.spawn(|cx| async move {
let results = results.await?;
cx.update(|cx| {
let mut output = ProjectIndexOutput {
status,
excerpts: Default::default(),
};
let excerpts = results.into_iter().map(|result| {
let abs_path = result
.worktree
.read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
let fs = fs.clone();
for search_result in search_results {
let path = ProjectPath {
worktree_id: search_result.worktree.read(cx).id(),
path: search_result.path.clone(),
};
async move {
let path = result.path.clone();
let text = fs.load(&abs_path?).await?;
let mut start = result.range.start;
let mut end = result.range.end.min(text.len());
while !text.is_char_boundary(start) {
start += 1;
}
while !text.is_char_boundary(end) {
end -= 1;
}
anyhow::Ok(CodebaseExcerpt {
path: path.to_string_lossy().to_string().into(),
text: SharedString::from(text[start..end].to_string()),
score: result.score,
})
let excerpts_for_path = output.excerpts.entry(path).or_default();
let ix = match excerpts_for_path
.binary_search_by_key(&search_result.range.start, |r| r.start)
{
Ok(ix) | Err(ix) => ix,
};
excerpts_for_path.insert(ix, search_result.range);
}
});
let mut files_searched = HashSet::new();
let excerpts = futures::future::join_all(excerpts)
.await
.into_iter()
.filter_map(|result| result.log_err())
.inspect(|excerpt| {
files_searched.insert(excerpt.path.clone());
})
.collect::<Vec<_>>();
anyhow::Ok(ProjectIndexOutput {
excerpts,
status,
files_searched,
output
})
})
}
fn output_view(
_tool_call_id: String,
input: Self::Input,
output: Result<Self::Output>,
cx: &mut WindowContext,
@ -220,34 +206,4 @@ impl LanguageModelTool for ProjectIndexTool {
CollapsibleContainer::new(ElementId::Name(nanoid::nanoid!().into()), false)
.start_slot("Searching code base")
}
fn format(_input: &Self::Input, output: &Result<Self::Output>) -> String {
match &output {
Ok(output) => {
let mut body = "Semantic search results:\n".to_string();
if output.status != Status::Idle {
body.push_str("Still indexing. Results may be incomplete.\n");
}
if output.excerpts.is_empty() {
body.push_str("No results found");
return body;
}
for excerpt in &output.excerpts {
body.push_str("Excerpt from ");
body.push_str(excerpt.path.as_ref());
body.push_str(", score ");
body.push_str(&excerpt.score.to_string());
body.push_str(":\n");
body.push_str("~~~\n");
body.push_str(excerpt.text.as_ref());
body.push_str("~~~\n");
}
body
}
Err(err) => format!("Error: {}", err),
}
}
}

13
crates/assistant2/src/ui/active_file_button.rs
View file

@ -1,4 +1,5 @@
use crate::attachments::{ActiveEditorAttachmentTool, UserAttachmentStore};
use crate::attachments::ActiveEditorAttachmentTool;
use assistant_tooling::AttachmentRegistry;
use editor::Editor;
use gpui::{prelude::*, Subscription, View};
use std::sync::Arc;
@ -13,7 +14,7 @@ enum Status {
}
pub struct ActiveFileButton {
attachment_store: Arc<UserAttachmentStore>,
attachment_registry: Arc<AttachmentRegistry>,
status: Status,
#[allow(dead_code)]
workspace_subscription: Subscription,
@ -21,7 +22,7 @@ pub struct ActiveFileButton {
impl ActiveFileButton {
pub fn new(
attachment_store: Arc<UserAttachmentStore>,
attachment_store: Arc<AttachmentRegistry>,
workspace: View<Workspace>,
cx: &mut ViewContext<Self>,
) -> Self {
@ -30,14 +31,14 @@ impl ActiveFileButton {
cx.defer(move |this, cx| this.update_active_buffer(workspace.clone(), cx));
Self {
attachment_store,
attachment_registry: attachment_store,
status: Status::NoFile,
workspace_subscription,
}
}
pub fn set_enabled(&mut self, enabled: bool) {
self.attachment_store
self.attachment_registry
.set_attachment_tool_enabled::<ActiveEditorAttachmentTool>(enabled);
}
@ -79,7 +80,7 @@ impl ActiveFileButton {
impl Render for ActiveFileButton {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
let is_enabled = self
.attachment_store
.attachment_registry
.is_attachment_tool_enabled::<ActiveEditorAttachmentTool>();
let icon = if is_enabled {

10
crates/assistant2/src/ui/composer.rs
View file

@ -11,7 +11,7 @@ use ui::{popover_menu, prelude::*, ButtonLike, ContextMenu, Divider, TextSize, T
#[derive(IntoElement)]
pub struct Composer {
editor: View<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
project_index_button: View<ProjectIndexButton>,
active_file_button: Option<View<ActiveFileButton>>,
model_selector: AnyElement,
}
@ -19,7 +19,7 @@ pub struct Composer {
impl Composer {
pub fn new(
editor: View<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
project_index_button: View<ProjectIndexButton>,
active_file_button: Option<View<ActiveFileButton>>,
model_selector: AnyElement,
) -> Self {
@ -32,11 +32,7 @@ impl Composer {
}
fn render_tools(&mut self, _cx: &mut WindowContext) -> impl IntoElement {
h_flex().children(
self.project_index_button
.clone()
.map(|view| view.into_any_element()),
)
h_flex().child(self.project_index_button.clone())
}
fn render_attachment_tools(&mut self, _cx: &mut WindowContext) -> impl IntoElement {