1use std::{sync::Arc, time::Duration};
2
3use crate::schema::json_schema_for;
4use crate::ui::ToolCallCardHeader;
5use anyhow::{Context as _, Result, anyhow};
6use assistant_tool::{
7 ActionLog, Tool, ToolCard, ToolResult, ToolResultContent, ToolResultOutput, ToolUseStatus,
8};
9use futures::{Future, FutureExt, TryFutureExt};
10use gpui::{
11 AnyWindowHandle, App, AppContext, Context, Entity, IntoElement, Task, WeakEntity, Window,
12};
13use language_model::{LanguageModel, LanguageModelRequest, LanguageModelToolSchemaFormat};
14use project::Project;
15use schemars::JsonSchema;
16use serde::{Deserialize, Serialize};
17use ui::{IconName, Tooltip, prelude::*};
18use web_search::WebSearchRegistry;
19use workspace::Workspace;
20use zed_llm_client::{WebSearchResponse, WebSearchResult};
21
22#[derive(Debug, Serialize, Deserialize, JsonSchema)]
23pub struct WebSearchToolInput {
24 /// The search term or question to query on the web.
25 query: String,
26}
27
28pub struct WebSearchTool;
29
30impl Tool for WebSearchTool {
31 fn name(&self) -> String {
32 "web_search".into()
33 }
34
35 fn needs_confirmation(&self, _: &serde_json::Value, _: &App) -> bool {
36 false
37 }
38
39 fn may_perform_edits(&self) -> bool {
40 false
41 }
42
43 fn description(&self) -> String {
44 "Search the web for information using your query. Use this when you need real-time information, facts, or data that might not be in your training. Results will include snippets and links from relevant web pages.".into()
45 }
46
47 fn icon(&self) -> IconName {
48 IconName::Globe
49 }
50
51 fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
52 json_schema_for::<WebSearchToolInput>(format)
53 }
54
55 fn ui_text(&self, _input: &serde_json::Value) -> String {
56 "Searching the Web".to_string()
57 }
58
59 fn run(
60 self: Arc<Self>,
61 input: serde_json::Value,
62 _request: Arc<LanguageModelRequest>,
63 _project: Entity<Project>,
64 _action_log: Entity<ActionLog>,
65 _model: Arc<dyn LanguageModel>,
66 _window: Option<AnyWindowHandle>,
67 cx: &mut App,
68 ) -> ToolResult {
69 let input = match serde_json::from_value::<WebSearchToolInput>(input) {
70 Ok(input) => input,
71 Err(err) => return Task::ready(Err(anyhow!(err))).into(),
72 };
73 let Some(provider) = WebSearchRegistry::read_global(cx).active_provider() else {
74 return Task::ready(Err(anyhow!("Web search is not available."))).into();
75 };
76
77 let search_task = provider.search(input.query, cx).map_err(Arc::new).shared();
78 let output = cx.background_spawn({
79 let search_task = search_task.clone();
80 async move {
81 let response = search_task.await.map_err(|err| anyhow!(err))?;
82 Ok(ToolResultOutput {
83 content: ToolResultContent::Text(
84 serde_json::to_string(&response)
85 .context("Failed to serialize search results")?,
86 ),
87 output: Some(serde_json::to_value(response)?),
88 })
89 }
90 });
91
92 ToolResult {
93 output,
94 card: Some(cx.new(|cx| WebSearchToolCard::new(search_task, cx)).into()),
95 }
96 }
97
98 fn deserialize_card(
99 self: Arc<Self>,
100 output: serde_json::Value,
101 _project: Entity<Project>,
102 _window: &mut Window,
103 cx: &mut App,
104 ) -> Option<assistant_tool::AnyToolCard> {
105 let output = serde_json::from_value::<WebSearchResponse>(output).ok()?;
106 let card = cx.new(|cx| WebSearchToolCard::new(Task::ready(Ok(output)), cx));
107 Some(card.into())
108 }
109}
110
111#[derive(RegisterComponent)]
112struct WebSearchToolCard {
113 response: Option<Result<WebSearchResponse>>,
114 _task: Task<()>,
115}
116
117impl WebSearchToolCard {
118 fn new(
119 search_task: impl 'static + Future<Output = Result<WebSearchResponse, Arc<anyhow::Error>>>,
120 cx: &mut Context<Self>,
121 ) -> Self {
122 let _task = cx.spawn(async move |this, cx| {
123 let response = search_task.await.map_err(|err| anyhow!(err));
124 this.update(cx, |this, cx| {
125 this.response = Some(response);
126 cx.notify();
127 })
128 .ok();
129 });
130
131 Self {
132 response: None,
133 _task,
134 }
135 }
136}
137
138impl ToolCard for WebSearchToolCard {
139 fn render(
140 &mut self,
141 _status: &ToolUseStatus,
142 _window: &mut Window,
143 _workspace: WeakEntity<Workspace>,
144 cx: &mut Context<Self>,
145 ) -> impl IntoElement {
146 let header = match self.response.as_ref() {
147 Some(Ok(response)) => {
148 let text: SharedString = if response.results.len() == 1 {
149 "1 result".into()
150 } else {
151 format!("{} results", response.results.len()).into()
152 };
153 ToolCallCardHeader::new(IconName::Globe, "Searched the Web")
154 .with_secondary_text(text)
155 }
156 Some(Err(error)) => {
157 ToolCallCardHeader::new(IconName::Globe, "Web Search").with_error(error.to_string())
158 }
159 None => ToolCallCardHeader::new(IconName::Globe, "Searching the Web").loading(),
160 };
161
162 let content = self.response.as_ref().and_then(|response| match response {
163 Ok(response) => Some(
164 v_flex()
165 .overflow_hidden()
166 .ml_1p5()
167 .pl(px(5.))
168 .border_l_1()
169 .border_color(cx.theme().colors().border_variant)
170 .gap_1()
171 .children(response.results.iter().enumerate().map(|(index, result)| {
172 let title = result.title.clone();
173 let url = SharedString::from(result.url.clone());
174
175 Button::new(("result", index), title)
176 .label_size(LabelSize::Small)
177 .color(Color::Muted)
178 .icon(IconName::ArrowUpRight)
179 .icon_size(IconSize::XSmall)
180 .icon_position(IconPosition::End)
181 .truncate(true)
182 .tooltip({
183 let url = url.clone();
184 move |window, cx| {
185 Tooltip::with_meta(
186 "Web Search Result",
187 None,
188 url.clone(),
189 window,
190 cx,
191 )
192 }
193 })
194 .on_click({
195 let url = url.clone();
196 move |_, _, cx| cx.open_url(&url)
197 })
198 }))
199 .into_any(),
200 ),
201 Err(_) => None,
202 });
203
204 v_flex().mb_3().gap_1().child(header).children(content)
205 }
206}
207
208impl Component for WebSearchToolCard {
209 fn scope() -> ComponentScope {
210 ComponentScope::Agent
211 }
212
213 fn preview(window: &mut Window, cx: &mut App) -> Option<AnyElement> {
214 let in_progress_search = cx.new(|cx| WebSearchToolCard {
215 response: None,
216 _task: cx.spawn(async move |_this, cx| {
217 loop {
218 cx.background_executor()
219 .timer(Duration::from_secs(60))
220 .await
221 }
222 }),
223 });
224
225 let successful_search = cx.new(|_cx| WebSearchToolCard {
226 response: Some(Ok(example_search_response())),
227 _task: Task::ready(()),
228 });
229
230 let error_search = cx.new(|_cx| WebSearchToolCard {
231 response: Some(Err(anyhow!("Failed to resolve https://google.com"))),
232 _task: Task::ready(()),
233 });
234
235 Some(
236 v_flex()
237 .gap_6()
238 .children(vec![example_group(vec![
239 single_example(
240 "In Progress",
241 div()
242 .size_full()
243 .child(in_progress_search.update(cx, |tool, cx| {
244 tool.render(
245 &ToolUseStatus::Pending,
246 window,
247 WeakEntity::new_invalid(),
248 cx,
249 )
250 .into_any_element()
251 }))
252 .into_any_element(),
253 ),
254 single_example(
255 "Successful",
256 div()
257 .size_full()
258 .child(successful_search.update(cx, |tool, cx| {
259 tool.render(
260 &ToolUseStatus::Finished("".into()),
261 window,
262 WeakEntity::new_invalid(),
263 cx,
264 )
265 .into_any_element()
266 }))
267 .into_any_element(),
268 ),
269 single_example(
270 "Error",
271 div()
272 .size_full()
273 .child(error_search.update(cx, |tool, cx| {
274 tool.render(
275 &ToolUseStatus::Error("".into()),
276 window,
277 WeakEntity::new_invalid(),
278 cx,
279 )
280 .into_any_element()
281 }))
282 .into_any_element(),
283 ),
284 ])])
285 .into_any_element(),
286 )
287 }
288}
289
290fn example_search_response() -> WebSearchResponse {
291 WebSearchResponse {
292 results: vec![
293 WebSearchResult {
294 title: "Alo".to_string(),
295 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
296 text: "Alo is a popular restaurant in Toronto.".to_string(),
297 },
298 WebSearchResult {
299 title: "Alo".to_string(),
300 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
301 text: "Information about Alo restaurant in Toronto.".to_string(),
302 },
303 WebSearchResult {
304 title: "Edulis".to_string(),
305 url: "https://www.google.com/maps/search/Edulis%2C+Toronto%2C+Canada".to_string(),
306 text: "Details about Edulis restaurant in Toronto.".to_string(),
307 },
308 WebSearchResult {
309 title: "Sushi Masaki Saito".to_string(),
310 url: "https://www.google.com/maps/search/Sushi+Masaki+Saito%2C+Toronto%2C+Canada"
311 .to_string(),
312 text: "Information about Sushi Masaki Saito in Toronto.".to_string(),
313 },
314 WebSearchResult {
315 title: "Shoushin".to_string(),
316 url: "https://www.google.com/maps/search/Shoushin%2C+Toronto%2C+Canada".to_string(),
317 text: "Details about Shoushin restaurant in Toronto.".to_string(),
318 },
319 WebSearchResult {
320 title: "Restaurant 20 Victoria".to_string(),
321 url:
322 "https://www.google.com/maps/search/Restaurant+20+Victoria%2C+Toronto%2C+Canada"
323 .to_string(),
324 text: "Information about Restaurant 20 Victoria in Toronto.".to_string(),
325 },
326 ],
327 }
328}