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

Streaming tools (#11629)

Browse source 
Stream characters in for tool calls to allow rendering partial input.


https://github.com/zed-industries/zed/assets/836375/0f023a4b-9c46-4449-ae69-8b6bcab41673

Release Notes:

- N/A

---------

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
Co-authored-by: Marshall <marshall@zed.dev>
Co-authored-by: Max <max@zed.dev>
This commit is contained in:
Kyle Kelley 2024-05-09 15:57:14 -07:00 • committed by GitHub
parent 27ed0f4273
commit 50c45c7897
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 786 additions and 653 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/assistant2/Cargo.toml
View file

@ -29,7 +29,6 @@ fuzzy.workspace = true
gpui.workspace = true
language.workspace = true
log.workspace = true
nanoid.workspace = true
open_ai.workspace = true
picker.workspace = true
project.workspace = true

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

@ -536,25 +536,27 @@ impl AssistantChat {
body.push_str(content);
}
for tool_call in delta.tool_calls {
let index = tool_call.index as usize;
for tool_call_delta in delta.tool_calls {
let index = tool_call_delta.index as usize;
if index >= message.tool_calls.len() {
message.tool_calls.resize_with(index + 1, Default::default);
}
let call = &mut message.tool_calls[index];
let tool_call = &mut message.tool_calls[index];
if let Some(id) = &tool_call.id {
call.id.push_str(id);
if let Some(id) = &tool_call_delta.id {
tool_call.id.push_str(id);
}
match tool_call.variant {
Some(proto::tool_call_delta::Variant::Function(tool_call)) => {
if let Some(name) = &tool_call.name {
call.name.push_str(name);
}
if let Some(arguments) = &tool_call.arguments {
call.arguments.push_str(arguments);
}
match tool_call_delta.variant {
Some(proto::tool_call_delta::Variant::Function(
tool_call_delta,
)) => {
this.tool_registry.update_tool_call(
tool_call,
tool_call_delta.name.as_deref(),
tool_call_delta.arguments.as_deref(),
cx,
);
}
None => {}
}
@ -587,34 +589,20 @@ impl AssistantChat {
} else {
if let Some(current_message) = messages.last_mut() {
for tool_call in current_message.tool_calls.iter() {
tool_tasks.push(this.tool_registry.call(tool_call, cx));
tool_tasks
.extend(this.tool_registry.execute_tool_call(&tool_call, cx));
}
}
}
}
})?;
// This ends recursion on calling for responses after tools
if tool_tasks.is_empty() {
return Ok(());
}
let tools = join_all(tool_tasks.into_iter()).await;
// If the WindowContext went away for any tool's view we don't include it
// especially since the below call would fail for the same reason.
let tools = tools.into_iter().filter_map(|tool| tool.ok()).collect();
this.update(cx, |this, cx| {
if let Some(ChatMessage::Assistant(AssistantMessage { messages, .. })) =
this.messages.last_mut()
{
if let Some(current_message) = messages.last_mut() {
current_message.tool_calls = tools;
cx.notify();
} else {
unreachable!()
}
}
})?;
join_all(tool_tasks.into_iter()).await;
}
}
@ -948,13 +936,11 @@ impl AssistantChat {
for tool_call in &message.tool_calls {
// Every tool call _must_ have a result by ID, otherwise OpenAI will error.
let content = match &tool_call.result {
Some(result) => {
result.generate(&tool_call.name, &mut project_context, cx)
}
None => "".to_string(),
};
let content = self.tool_registry.content_for_tool_call(
tool_call,
&mut project_context,
cx,
);
completion_messages.push(CompletionMessage::Tool {
content,
tool_call_id: tool_call.id.clone(),
@ -1003,7 +989,11 @@ impl AssistantChat {
tool_calls: message
.tool_calls
.iter()
.map(|tool_call| self.tool_registry.serialize_tool_call(tool_call))
.filter_map(|tool_call| {
self.tool_registry
.serialize_tool_call(tool_call, cx)
.log_err()
})
.collect(),
})
.collect(),

4
crates/assistant2/src/attachments/active_file.rs
View file

@ -1,7 +1,7 @@
use std::{path::PathBuf, sync::Arc};
use anyhow::{anyhow, Result};
use assistant_tooling::{LanguageModelAttachment, ProjectContext, ToolOutput};
use assistant_tooling::{AttachmentOutput, LanguageModelAttachment, ProjectContext};
use editor::Editor;
use gpui::{Render, Task, View, WeakModel, WeakView};
use language::Buffer;
@ -52,7 +52,7 @@ impl Render for FileAttachmentView {
}
}
impl ToolOutput for FileAttachmentView {
impl AttachmentOutput for FileAttachmentView {
fn generate(&self, project: &mut ProjectContext, cx: &mut WindowContext) -> String {
if let Some(path) = &self.project_path {
project.add_file(path.clone());

259
crates/assistant2/src/tools/annotate_code.rs
View file

@ -4,7 +4,8 @@ use editor::{
display_map::{BlockContext, BlockDisposition, BlockProperties, BlockStyle},
Editor, MultiBuffer,
};
use gpui::{prelude::*, AnyElement, Model, Task, View, WeakView};
use futures::{channel::mpsc::UnboundedSender, StreamExt as _};
use gpui::{prelude::*, AnyElement, AsyncWindowContext, Model, Task, View, WeakView};
use language::ToPoint;
use project::{search::SearchQuery, Project, ProjectPath};
use schemars::JsonSchema;
@ -25,14 +26,19 @@ impl AnnotationTool {
}
}
#[derive(Debug, Deserialize, JsonSchema, Clone)]
#[derive(Default, Debug, Deserialize, JsonSchema, Clone)]
pub struct AnnotationInput {
/// Name for this set of annotations
#[serde(default = "default_title")]
title: String,
/// Excerpts from the file to show to the user.
excerpts: Vec<Excerpt>,
}
fn default_title() -> String {
"Untitled".to_string()
}
#[derive(Debug, Deserialize, JsonSchema, Clone)]
struct Excerpt {
/// Path to the file
@ -44,8 +50,6 @@ struct Excerpt {
}
impl LanguageModelTool for AnnotationTool {
type Input = AnnotationInput;
type Output = String;
type View = AnnotationResultView;
fn name(&self) -> String {
@ -56,67 +60,100 @@ impl LanguageModelTool for AnnotationTool {
"Dynamically annotate symbols in the current codebase. Opens a buffer in a panel in their editor, to the side of the conversation. The annotations are shown in the editor as a block decoration.".to_string()
}
fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
let workspace = self.workspace.clone();
let project = self.project.clone();
let excerpts = input.excerpts.clone();
let title = input.title.clone();
fn view(&self, cx: &mut WindowContext) -> View<Self::View> {
cx.new_view(|cx| {
let (tx, mut rx) = futures::channel::mpsc::unbounded();
cx.spawn(|view, mut cx| async move {
while let Some(excerpt) = rx.next().await {
AnnotationResultView::add_excerpt(view.clone(), excerpt, &mut cx).await?;
}
anyhow::Ok(())
})
.detach();
AnnotationResultView {
project: self.project.clone(),
workspace: self.workspace.clone(),
tx,
pending_excerpt: None,
added_editor_to_workspace: false,
editor: None,
error: None,
rendered_excerpt_count: 0,
}
})
}
}
pub struct AnnotationResultView {
workspace: WeakView<Workspace>,
project: Model<Project>,
pending_excerpt: Option<Excerpt>,
added_editor_to_workspace: bool,
editor: Option<View<Editor>>,
tx: UnboundedSender<Excerpt>,
error: Option<anyhow::Error>,
rendered_excerpt_count: usize,
}
impl AnnotationResultView {
async fn add_excerpt(
this: WeakView<Self>,
excerpt: Excerpt,
cx: &mut AsyncWindowContext,
) -> Result<()> {
let project = this.update(cx, |this, _cx| this.project.clone())?;
let worktree_id = project.update(cx, |project, cx| {
let worktree = project.worktrees().next()?;
let worktree_id = worktree.read(cx).id();
Some(worktree_id)
});
})?;
let worktree_id = if let Some(worktree_id) = worktree_id {
worktree_id
} else {
return Task::ready(Err(anyhow::anyhow!("No worktree found")));
return Err(anyhow::anyhow!("No worktree found"));
};
let buffer_tasks = project.update(cx, |project, cx| {
excerpts
.iter()
.map(|excerpt| {
project.open_buffer(
ProjectPath {
worktree_id,
path: Path::new(&excerpt.path).into(),
},
cx,
)
let buffer_task = project.update(cx, |project, cx| {
project.open_buffer(
ProjectPath {
worktree_id,
path: Path::new(&excerpt.path).into(),
},
cx,
)
})?;
let buffer = match buffer_task.await {
Ok(buffer) => buffer,
Err(error) => {
return this.update(cx, |this, cx| {
this.error = Some(error);
cx.notify();
})
.collect::<Vec<_>>()
});
}
};
cx.spawn(move |mut cx| async move {
let buffers = futures::future::try_join_all(buffer_tasks).await?;
let snapshot = buffer.update(cx, |buffer, _cx| buffer.snapshot())?;
let query = SearchQuery::text(&excerpt.text_passage, false, false, false, vec![], vec![])?;
let matches = query.search(&snapshot, None).await;
let Some(first_match) = matches.first() else {
log::warn!(
"text {:?} does not appear in '{}'",
excerpt.text_passage,
excerpt.path
);
return Ok(());
};
let multibuffer = cx.new_model(|_cx| {
MultiBuffer::new(0, language::Capability::ReadWrite).with_title(title)
})?;
let editor =
cx.new_view(|cx| Editor::for_multibuffer(multibuffer, Some(project), cx))?;
this.update(cx, |this, cx| {
let mut start = first_match.start.to_point(&snapshot);
start.column = 0;
for (excerpt, buffer) in excerpts.iter().zip(buffers.iter()) {
let snapshot = buffer.update(&mut cx, |buffer, _cx| buffer.snapshot())?;
let query =
SearchQuery::text(&excerpt.text_passage, false, false, false, vec![], vec![])?;
let matches = query.search(&snapshot, None).await;
let Some(first_match) = matches.first() else {
log::warn!(
"text {:?} does not appear in '{}'",
excerpt.text_passage,
excerpt.path
);
continue;
};
let mut start = first_match.start.to_point(&snapshot);
start.column = 0;
editor.update(&mut cx, |editor, cx| {
if let Some(editor) = &this.editor {
editor.update(cx, |editor, cx| {
let ranges = editor.buffer().update(cx, |multibuffer, cx| {
multibuffer.push_excerpts_with_context_lines(
buffer.clone(),
@ -125,7 +162,8 @@ impl LanguageModelTool for AnnotationTool {
cx,
)
});
let annotation = SharedString::from(excerpt.annotation.clone());
let annotation = SharedString::from(excerpt.annotation);
editor.insert_blocks(
[BlockProperties {
position: ranges[0].start,
@ -137,30 +175,22 @@ impl LanguageModelTool for AnnotationTool {
None,
cx,
);
})?;
});
if !this.added_editor_to_workspace {
this.added_editor_to_workspace = true;
this.workspace
.update(cx, |workspace, cx| {
workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
})
.log_err();
}
}
})?;
workspace
.update(&mut cx, |workspace, cx| {
workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
})
.log_err();
anyhow::Ok("showed comments to users in a new view".into())
})
Ok(())
}
fn view(
&self,
_: Self::Input,
output: Result<Self::Output>,
cx: &mut WindowContext,
) -> View<Self::View> {
cx.new_view(|_cx| AnnotationResultView { output })
}
}
impl AnnotationTool {
fn render_note_block(explanation: &SharedString, cx: &mut BlockContext) -> AnyElement {
let anchor_x = cx.anchor_x;
let gutter_width = cx.gutter_dimensions.width;
@ -186,24 +216,89 @@ impl AnnotationTool {
}
}
pub struct AnnotationResultView {
output: Result<String>,
}
impl Render for AnnotationResultView {
fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
match &self.output {
Ok(output) => div().child(output.clone().into_any_element()),
Err(error) => div().child(format!("failed to open path: {:?}", error)),
if let Some(error) = &self.error {
ui::Label::new(error.to_string()).into_any_element()
} else {
ui::Label::new(SharedString::from(format!(
"Opened a buffer with {} excerpts",
self.rendered_excerpt_count
)))
.into_any_element()
}
}
}
impl ToolOutput for AnnotationResultView {
fn generate(&self, _: &mut ProjectContext, _: &mut WindowContext) -> String {
match &self.output {
Ok(output) => output.clone(),
Err(err) => format!("Failed to create buffer: {err:?}"),
type Input = AnnotationInput;
type SerializedState = Option<String>;
fn generate(&self, _: &mut ProjectContext, _: &mut ViewContext<Self>) -> String {
if let Some(error) = &self.error {
format!("Failed to create buffer: {error:?}")
} else {
format!(
"opened {} excerpts in a buffer",
self.rendered_excerpt_count
)
}
}
fn set_input(&mut self, mut input: Self::Input, cx: &mut ViewContext<Self>) {
let editor = if let Some(editor) = &self.editor {
editor.clone()
} else {
let multibuffer = cx.new_model(|_cx| {
MultiBuffer::new(0, language::Capability::ReadWrite).with_title(String::new())
});
let editor = cx.new_view(|cx| {
Editor::for_multibuffer(multibuffer.clone(), Some(self.project.clone()), cx)
});
self.editor = Some(editor.clone());
editor
};
editor.update(cx, |editor, cx| {
editor.buffer().update(cx, |multibuffer, cx| {
if multibuffer.title(cx) != input.title {
multibuffer.set_title(input.title.clone(), cx);
}
});
self.pending_excerpt = input.excerpts.pop();
for excerpt in input.excerpts.iter().skip(self.rendered_excerpt_count) {
self.tx.unbounded_send(excerpt.clone()).ok();
}
self.rendered_excerpt_count = input.excerpts.len();
});
cx.notify();
}
fn execute(&mut self, _cx: &mut ViewContext<Self>) -> Task<Result<()>> {
if let Some(excerpt) = self.pending_excerpt.take() {
self.rendered_excerpt_count += 1;
self.tx.unbounded_send(excerpt.clone()).ok();
}
self.tx.close_channel();
Task::ready(Ok(()))
}
fn serialize(&self, _cx: &mut ViewContext<Self>) -> Self::SerializedState {
self.error.as_ref().map(|error| error.to_string())
}
fn deserialize(
&mut self,
output: Self::SerializedState,
_cx: &mut ViewContext<Self>,
) -> Result<()> {
if let Some(error_message) = output {
self.error = Some(anyhow::anyhow!("{}", error_message));
}
Ok(())
}
}

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

@ -1,4 +1,4 @@
use anyhow::Result;
use anyhow::{anyhow, Result};
use assistant_tooling::{LanguageModelTool, ProjectContext, ToolOutput};
use editor::Editor;
use gpui::{prelude::*, Model, Task, View, WeakView};
@ -20,7 +20,7 @@ impl CreateBufferTool {
}
}
#[derive(Debug, Deserialize, JsonSchema)]
#[derive(Debug, Clone, Deserialize, JsonSchema)]
pub struct CreateBufferInput {
/// The contents of the buffer.
text: String,
@ -32,8 +32,6 @@ pub struct CreateBufferInput {
}
impl LanguageModelTool for CreateBufferTool {
type Input = CreateBufferInput;
type Output = ();
type View = CreateBufferView;
fn name(&self) -> String {
@ -44,13 +42,59 @@ impl LanguageModelTool for CreateBufferTool {
"Create a new buffer in the current codebase".to_string()
}
fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
fn view(&self, cx: &mut WindowContext) -> View<Self::View> {
cx.new_view(|_cx| CreateBufferView {
workspace: self.workspace.clone(),
project: self.project.clone(),
input: None,
error: None,
})
}
}
pub struct CreateBufferView {
workspace: WeakView<Workspace>,
project: Model<Project>,
input: Option<CreateBufferInput>,
error: Option<anyhow::Error>,
}
impl Render for CreateBufferView {
fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
div().child("Opening a buffer")
}
}
impl ToolOutput for CreateBufferView {
type Input = CreateBufferInput;
type SerializedState = ();
fn generate(&self, _project: &mut ProjectContext, _cx: &mut ViewContext<Self>) -> String {
let Some(input) = self.input.as_ref() else {
return "No input".to_string();
};
match &self.error {
None => format!("Created a new {} buffer", input.language),
Some(err) => format!("Failed to create buffer: {err:?}"),
}
}
fn set_input(&mut self, input: Self::Input, _cx: &mut ViewContext<Self>) {
self.input = Some(input);
}
fn execute(&mut self, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
cx.spawn({
let workspace = self.workspace.clone();
let project = self.project.clone();
let text = input.text.clone();
let language_name = input.language.clone();
|mut cx| async move {
let input = self.input.clone();
|_this, mut cx| async move {
let input = input.ok_or_else(|| anyhow!("no input"))?;
let text = input.text.clone();
let language_name = input.language.clone();
let language = cx
.update(|cx| {
project
@ -86,35 +130,15 @@ impl LanguageModelTool for CreateBufferTool {
})
}
fn view(
&self,
input: Self::Input,
output: Result<Self::Output>,
cx: &mut WindowContext,
) -> View<Self::View> {
cx.new_view(|_cx| CreateBufferView {
language: input.language,
output,
})
}
}
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:?}"),
}
fn serialize(&self, _cx: &mut ViewContext<Self>) -> Self::SerializedState {
()
}
fn deserialize(
&mut self,
_output: Self::SerializedState,
_cx: &mut ViewContext<Self>,
) -> Result<()> {
Ok(())
}
}

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

@ -1,4 +1,4 @@
use anyhow::{anyhow, Result};
use anyhow::Result;
use assistant_tooling::{LanguageModelTool, ToolOutput};
use collections::BTreeMap;
use gpui::{prelude::*, Model, Task};
@ -6,9 +6,8 @@ use project::ProjectPath;
use schemars::JsonSchema;
use semantic_index::{ProjectIndex, Status};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
use ui::{div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
use ui::{prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
const DEFAULT_SEARCH_LIMIT: usize = 20;
@ -16,10 +15,26 @@ 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.
// Any changes or deletions to the `CodebaseQuery` comments will change model behavior.
#[derive(Default)]
enum ProjectIndexToolState {
#[default]
CollectingQuery,
Searching,
Error(anyhow::Error),
Finished {
excerpts: BTreeMap<ProjectPath, Vec<Range<usize>>>,
index_status: Status,
},
}
#[derive(Deserialize, JsonSchema)]
pub struct ProjectIndexView {
project_index: Model<ProjectIndex>,
input: CodebaseQuery,
expanded_header: bool,
state: ProjectIndexToolState,
}
#[derive(Default, Deserialize, JsonSchema)]
pub struct CodebaseQuery {
/// Semantic search query
query: String,
@ -27,21 +42,14 @@ pub struct CodebaseQuery {
limit: Option<usize>,
}
pub struct ProjectIndexView {
input: CodebaseQuery,
status: Status,
excerpts: Result<BTreeMap<ProjectPath, Vec<Range<usize>>>>,
element_id: ElementId,
expanded_header: bool,
}
#[derive(Serialize, Deserialize)]
pub struct ProjectIndexOutput {
status: Status,
pub struct SerializedState {
index_status: Status,
error_message: Option<String>,
worktrees: BTreeMap<Arc<Path>, WorktreeIndexOutput>,
}
#[derive(Serialize, Deserialize)]
#[derive(Default, Serialize, Deserialize)]
struct WorktreeIndexOutput {
excerpts: BTreeMap<Arc<Path>, Vec<Range<usize>>>,
}
@ -56,58 +64,80 @@ impl ProjectIndexView {
impl Render for ProjectIndexView {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
let query = self.input.query.clone();
let excerpts = match &self.excerpts {
Err(err) => {
return div().child(Label::new(format!("Error: {}", err)).color(Color::Error));
let (header_text, content) = match &self.state {
ProjectIndexToolState::Error(error) => {
return format!("failed to search: {error:?}").into_any_element()
}
ProjectIndexToolState::CollectingQuery | ProjectIndexToolState::Searching => {
("Searching...".to_string(), div())
}
ProjectIndexToolState::Finished { excerpts, .. } => {
let file_count = excerpts.len();
let header_text = format!(
"Read {} {}",
file_count,
if file_count == 1 { "file" } else { "files" }
);
let el = v_flex().gap_2().children(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),
)
}));
(header_text, el)
}
Ok(excerpts) => excerpts,
};
let file_count = excerpts.len();
let header = h_flex()
.gap_2()
.child(Icon::new(IconName::File))
.child(format!(
"Read {} {}",
file_count,
if file_count == 1 { "file" } else { "files" }
));
.child(header_text);
v_flex().gap_3().child(
CollapsibleContainer::new(self.element_id.clone(), self.expanded_header)
.start_slot(header)
.on_click(cx.listener(move |this, _, cx| {
this.toggle_header(cx);
}))
.child(
v_flex()
.gap_3()
.p_3()
.child(
h_flex()
.gap_2()
.child(Icon::new(IconName::MagnifyingGlass))
.child(Label::new(format!("`{}`", query)).color(Color::Muted)),
)
.child(v_flex().gap_2().children(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),
v_flex()
.gap_3()
.child(
CollapsibleContainer::new("collapsible-container", self.expanded_header)
.start_slot(header)
.on_click(cx.listener(move |this, _, cx| {
this.toggle_header(cx);
}))
.child(
v_flex()
.gap_3()
.p_3()
.child(
h_flex()
.gap_2()
.child(Icon::new(IconName::MagnifyingGlass))
.child(Label::new(format!("`{}`", query)).color(Color::Muted)),
)
}))),
),
)
.child(content),
),
)
.into_any_element()
}
}
impl ToolOutput for ProjectIndexView {
type Input = CodebaseQuery;
type SerializedState = SerializedState;
fn generate(
&self,
context: &mut assistant_tooling::ProjectContext,
_: &mut WindowContext,
_: &mut ViewContext<Self>,
) -> String {
match &self.excerpts {
Ok(excerpts) => {
match &self.state {
ProjectIndexToolState::CollectingQuery => String::new(),
ProjectIndexToolState::Searching => String::new(),
ProjectIndexToolState::Error(error) => format!("failed to search: {error:?}"),
ProjectIndexToolState::Finished {
excerpts,
index_status,
} => {
let mut body = "found results in the following paths:\n".to_string();
for (project_path, ranges) in excerpts {
@ -115,15 +145,126 @@ impl ToolOutput for ProjectIndexView {
writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
}
if self.status != Status::Idle {
if *index_status != Status::Idle {
body.push_str("Still indexing. Results may be incomplete.\n");
}
body
}
Err(err) => format!("Error: {}", err),
}
}
fn set_input(&mut self, input: Self::Input, cx: &mut ViewContext<Self>) {
self.input = input;
cx.notify();
}
fn execute(&mut self, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
self.state = ProjectIndexToolState::Searching;
cx.notify();
let project_index = self.project_index.read(cx);
let index_status = project_index.status();
let search = project_index.search(
self.input.query.clone(),
self.input.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
cx,
);
cx.spawn(|this, mut cx| async move {
let search_result = search.await;
this.update(&mut cx, |this, cx| {
match search_result {
Ok(search_results) => {
let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
for search_result in search_results {
let project_path = ProjectPath {
worktree_id: search_result.worktree.read(cx).id(),
path: search_result.path,
};
excerpts
.entry(project_path)
.or_default()
.push(search_result.range);
}
this.state = ProjectIndexToolState::Finished {
excerpts,
index_status,
};
}
Err(error) => {
this.state = ProjectIndexToolState::Error(error);
}
}
cx.notify();
})
})
}
fn serialize(&self, cx: &mut ViewContext<Self>) -> Self::SerializedState {
let mut serialized = SerializedState {
error_message: None,
index_status: Status::Idle,
worktrees: Default::default(),
};
match &self.state {
ProjectIndexToolState::Error(err) => serialized.error_message = Some(err.to_string()),
ProjectIndexToolState::Finished {
excerpts,
index_status,
} => {
serialized.index_status = *index_status;
if let Some(project) = self.project_index.read(cx).project().upgrade() {
let project = project.read(cx);
for (project_path, excerpts) in excerpts {
if let Some(worktree) =
project.worktree_for_id(project_path.worktree_id, cx)
{
let worktree_path = worktree.read(cx).abs_path();
serialized
.worktrees
.entry(worktree_path)
.or_default()
.excerpts
.insert(project_path.path.clone(), excerpts.clone());
}
}
}
}
_ => {}
}
serialized
}
fn deserialize(
&mut self,
serialized: Self::SerializedState,
cx: &mut ViewContext<Self>,
) -> Result<()> {
if !serialized.worktrees.is_empty() {
let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
if let Some(project) = self.project_index.read(cx).project().upgrade() {
let project = project.read(cx);
for (worktree_path, worktree_state) in serialized.worktrees {
if let Some(worktree) = project
.worktrees()
.find(|worktree| worktree.read(cx).abs_path() == worktree_path)
{
let worktree_id = worktree.read(cx).id();
for (path, serialized_excerpts) in worktree_state.excerpts {
excerpts.insert(ProjectPath { worktree_id, path }, serialized_excerpts);
}
}
}
}
self.state = ProjectIndexToolState::Finished {
excerpts,
index_status: serialized.index_status,
};
}
cx.notify();
Ok(())
}
}
impl ProjectIndexTool {
@ -133,8 +274,6 @@ impl ProjectIndexTool {
}
impl LanguageModelTool for ProjectIndexTool {
type Input = CodebaseQuery;
type Output = ProjectIndexOutput;
type View = ProjectIndexView;
fn name(&self) -> String {
@ -145,109 +284,12 @@ impl LanguageModelTool for ProjectIndexTool {
"Semantic search against the user's current codebase, returning excerpts related to the query by computing a dot product against embeddings of code chunks in the code base and an embedding of the query.".to_string()
}
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 search = project_index.search(
query.query.clone(),
query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
cx,
);
cx.spawn(|mut cx| async move {
let search_results = search.await?;
cx.update(|cx| {
let mut output = ProjectIndexOutput {
status,
worktrees: Default::default(),
};
for search_result in search_results {
let worktree_path = search_result.worktree.read(cx).abs_path();
let excerpts = &mut output
.worktrees
.entry(worktree_path)
.or_insert(WorktreeIndexOutput {
excerpts: Default::default(),
})
.excerpts;
let excerpts_for_path = excerpts.entry(search_result.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);
}
output
})
fn view(&self, cx: &mut WindowContext) -> gpui::View<Self::View> {
cx.new_view(|_| ProjectIndexView {
state: ProjectIndexToolState::CollectingQuery,
input: Default::default(),
expanded_header: false,
project_index: self.project_index.clone(),
})
}
fn view(
&self,
input: Self::Input,
output: Result<Self::Output>,
cx: &mut WindowContext,
) -> gpui::View<Self::View> {
cx.new_view(|cx| {
let status;
let excerpts;
match output {
Ok(output) => {
status = output.status;
let project_index = self.project_index.read(cx);
if let Some(project) = project_index.project().upgrade() {
let project = project.read(cx);
excerpts = Ok(output
.worktrees
.into_iter()
.filter_map(|(abs_path, output)| {
for worktree in project.worktrees() {
let worktree = worktree.read(cx);
if worktree.abs_path() == abs_path {
return Some((worktree.id(), output.excerpts));
}
}
None
})
.flat_map(|(worktree_id, excerpts)| {
excerpts.into_iter().map(move |(path, ranges)| {
(ProjectPath { worktree_id, path }, ranges)
})
})
.collect::<BTreeMap<_, _>>());
} else {
excerpts = Err(anyhow!("project was dropped"));
}
}
Err(err) => {
status = Status::Idle;
excerpts = Err(err);
}
};
ProjectIndexView {
input,
status,
excerpts,
element_id: ElementId::Name(nanoid::nanoid!().into()),
expanded_header: false,
}
})
}
fn render_running(arguments: &Option<Value>, _: &mut WindowContext) -> impl IntoElement {
let text: String = arguments
.as_ref()
.and_then(|arguments| arguments.get("query"))
.and_then(|query| query.as_str())
.map(|query| format!("Searching for: {}", query))
.unwrap_or_else(|| "Preparing search...".to_string());
CollapsibleContainer::new(ElementId::Name(nanoid::nanoid!().into()), false).start_slot(text)
}
}