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