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

context_servers: Support tools (#19548)

Browse source 
This PR depends on #19547 

This PR adds support for tools from context servers. Context servers are
free to expose tools that Zed can pass to models. When called by the
model, Zed forwards the request to context servers. This allows for some
interesting techniques. Context servers can easily expose tools such as
querying local databases, reading or writing local files, reading
resources over authenticated APIs (e.g. kubernetes, asana, etc).

This is currently experimental. 

Things to discuss
* I want to still add a confirm dialog asking people if a server is
allows to use the tool. Should do this or just use the tool and assume
trustworthyness of context servers?
* Can we add tool use behind a local setting flag?

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
This commit is contained in:
David Soria Parra 2024-10-28 14:37:58 +00:00 • committed by GitHub
parent cdddb4d360
commit 8a96ea25c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 235 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

85
crates/assistant/src/assistant.rs
View file

@ -298,25 +298,64 @@ fn register_context_server_handlers(cx: &mut AppContext) {
return;
};
if let Some(prompts) = protocol.list_prompts().await.log_err() {
for prompt in prompts
.into_iter()
.filter(context_server_command::acceptable_prompt)
{
log::info!(
"registering context server command: {:?}",
prompt.name
);
context_server_registry.register_command(
server.id.clone(),
prompt.name.as_str(),
);
slash_command_registry.register_command(
context_server_command::ContextServerSlashCommand::new(
&server, prompt,
),
true,
);
if protocol.capable(context_servers::protocol::ServerCapability::Prompts) {
if let Some(prompts) = protocol.list_prompts().await.log_err() {
for prompt in prompts
.into_iter()
.filter(context_server_command::acceptable_prompt)
{
log::info!(
"registering context server command: {:?}",
prompt.name
);
context_server_registry.register_command(
server.id.clone(),
prompt.name.as_str(),
);
slash_command_registry.register_command(
context_server_command::ContextServerSlashCommand::new(
&server, prompt,
),
true,
);
}
}
}
})
.detach();
}
},
);
cx.update_model(
&manager,
|manager: &mut context_servers::manager::ContextServerManager, cx| {
let tool_registry = ToolRegistry::global(cx);
let context_server_registry = ContextServerRegistry::global(cx);
if let Some(server) = manager.get_server(server_id) {
cx.spawn(|_, _| async move {
let Some(protocol) = server.client.read().clone() else {
return;
};
if protocol.capable(context_servers::protocol::ServerCapability::Tools) {
if let Some(tools) = protocol.list_tools().await.log_err() {
for tool in tools.tools {
log::info!(
"registering context server tool: {:?}",
tool.name
);
context_server_registry.register_tool(
server.id.clone(),
tool.name.as_str(),
);
tool_registry.register_tool(
tools::context_server_tool::ContextServerTool::new(
server.id.clone(),
tool
),
);
}
}
}
})
@ -334,6 +373,14 @@ fn register_context_server_handlers(cx: &mut AppContext) {
context_server_registry.unregister_command(&server_id, &command_name);
}
}
if let Some(tools) = context_server_registry.get_tools(server_id) {
let tool_registry = ToolRegistry::global(cx);
for tool_name in tools {
tool_registry.unregister_tool_by_name(&tool_name);
context_server_registry.unregister_tool(&server_id, &tool_name);
}
}
}
},
)

1
crates/assistant/src/tools.rs
View file

@ -1 +1,2 @@
pub mod context_server_tool;
pub mod now_tool;

82
crates/assistant/src/tools/context_server_tool.rs Normal file
View file

@ -0,0 +1,82 @@
use anyhow::{anyhow, bail};
use assistant_tool::Tool;
use context_servers::manager::ContextServerManager;
use context_servers::types;
use gpui::Task;
pub struct ContextServerTool {
server_id: String,
tool: types::Tool,
}
impl ContextServerTool {
pub fn new(server_id: impl Into<String>, tool: types::Tool) -> Self {
Self {
server_id: server_id.into(),
tool,
}
}
}
impl Tool for ContextServerTool {
fn name(&self) -> String {
self.tool.name.clone()
}
fn description(&self) -> String {
self.tool.description.clone().unwrap_or_default()
}
fn input_schema(&self) -> serde_json::Value {
match &self.tool.input_schema {
serde_json::Value::Null => {
serde_json::json!({ "type": "object", "properties": [] })
}
serde_json::Value::Object(map) if map.is_empty() => {
serde_json::json!({ "type": "object", "properties": [] })
}
_ => self.tool.input_schema.clone(),
}
}
fn run(
self: std::sync::Arc<Self>,
input: serde_json::Value,
_workspace: gpui::WeakView<workspace::Workspace>,
cx: &mut ui::WindowContext,
) -> gpui::Task<gpui::Result<String>> {
let manager = ContextServerManager::global(cx);
let manager = manager.read(cx);
if let Some(server) = manager.get_server(&self.server_id) {
cx.foreground_executor().spawn({
let tool_name = self.tool.name.clone();
async move {
let Some(protocol) = server.client.read().clone() else {
bail!("Context server not initialized");
};
let arguments = if let serde_json::Value::Object(map) = input {
Some(map.into_iter().collect())
} else {
None
};
log::trace!(
"Running tool: {} with arguments: {:?}",
tool_name,
arguments
);
let response = protocol.run_tool(tool_name, arguments).await?;
let tool_result = match response.tool_result {
serde_json::Value::String(s) => s,
_ => serde_json::to_string(&response.tool_result)?,
};
Ok(tool_result)
}
})
} else {
Task::ready(Err(anyhow!("Context server not found")))
}
}
}