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::{LanguageModelRequestMessage, 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::{WebSearchCitation, WebSearchResponse};
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 _messages: &[LanguageModelRequestMessage],
57 _project: Entity<Project>,
58 _action_log: Entity<ActionLog>,
59 _window: Option<AnyWindowHandle>,
60 cx: &mut App,
61 ) -> ToolResult {
62 let input = match serde_json::from_value::<WebSearchToolInput>(input) {
63 Ok(input) => input,
64 Err(err) => return Task::ready(Err(anyhow!(err))).into(),
65 };
66 let Some(provider) = WebSearchRegistry::read_global(cx).active_provider() else {
67 return Task::ready(Err(anyhow!("Web search is not available."))).into();
68 };
69
70 let search_task = provider.search(input.query, cx).map_err(Arc::new).shared();
71 let output = cx.background_spawn({
72 let search_task = search_task.clone();
73 async move {
74 let response = search_task.await.map_err(|err| anyhow!(err))?;
75 serde_json::to_string(&response).context("Failed to serialize search results")
76 }
77 });
78
79 ToolResult {
80 output,
81 card: Some(cx.new(|cx| WebSearchToolCard::new(search_task, cx)).into()),
82 }
83 }
84}
85
86#[derive(RegisterComponent)]
87struct WebSearchToolCard {
88 response: Option<Result<WebSearchResponse>>,
89 _task: Task<()>,
90}
91
92impl WebSearchToolCard {
93 fn new(
94 search_task: impl 'static + Future<Output = Result<WebSearchResponse, Arc<anyhow::Error>>>,
95 cx: &mut Context<Self>,
96 ) -> Self {
97 let _task = cx.spawn(async move |this, cx| {
98 let response = search_task.await.map_err(|err| anyhow!(err));
99 this.update(cx, |this, cx| {
100 this.response = Some(response);
101 cx.notify();
102 })
103 .ok();
104 });
105
106 Self {
107 response: None,
108 _task,
109 }
110 }
111}
112
113impl ToolCard for WebSearchToolCard {
114 fn render(
115 &mut self,
116 _status: &ToolUseStatus,
117 _window: &mut Window,
118 _workspace: WeakEntity<Workspace>,
119 cx: &mut Context<Self>,
120 ) -> impl IntoElement {
121 let header = match self.response.as_ref() {
122 Some(Ok(response)) => {
123 let text: SharedString = if response.citations.len() == 1 {
124 "1 result".into()
125 } else {
126 format!("{} results", response.citations.len()).into()
127 };
128 ToolCallCardHeader::new(IconName::Globe, "Searched the Web")
129 .with_secondary_text(text)
130 }
131 Some(Err(error)) => {
132 ToolCallCardHeader::new(IconName::Globe, "Web Search").with_error(error.to_string())
133 }
134 None => ToolCallCardHeader::new(IconName::Globe, "Searching the Web").loading(),
135 };
136
137 let content =
138 self.response.as_ref().and_then(|response| match response {
139 Ok(response) => {
140 Some(
141 v_flex()
142 .overflow_hidden()
143 .ml_1p5()
144 .pl(px(5.))
145 .border_l_1()
146 .border_color(cx.theme().colors().border_variant)
147 .gap_1()
148 .children(response.citations.iter().enumerate().map(
149 |(index, citation)| {
150 let title = citation.title.clone();
151 let url = citation.url.clone();
152
153 Button::new(("citation", 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 "Citation Link",
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 ))
178 .into_any(),
179 )
180 }
181 Err(_) => None,
182 });
183
184 v_flex().mb_3().gap_1().child(header).children(content)
185 }
186}
187
188impl Component for WebSearchToolCard {
189 fn scope() -> ComponentScope {
190 ComponentScope::Agent
191 }
192
193 fn preview(window: &mut Window, cx: &mut App) -> Option<AnyElement> {
194 let in_progress_search = cx.new(|cx| WebSearchToolCard {
195 response: None,
196 _task: cx.spawn(async move |_this, cx| {
197 loop {
198 cx.background_executor()
199 .timer(Duration::from_secs(60))
200 .await
201 }
202 }),
203 });
204
205 let successful_search = cx.new(|_cx| WebSearchToolCard {
206 response: Some(Ok(example_search_response())),
207 _task: Task::ready(()),
208 });
209
210 let error_search = cx.new(|_cx| WebSearchToolCard {
211 response: Some(Err(anyhow!("Failed to resolve https://google.com"))),
212 _task: Task::ready(()),
213 });
214
215 Some(
216 v_flex()
217 .gap_6()
218 .children(vec![example_group(vec![
219 single_example(
220 "In Progress",
221 div()
222 .size_full()
223 .child(in_progress_search.update(cx, |tool, cx| {
224 tool.render(
225 &ToolUseStatus::Pending,
226 window,
227 WeakEntity::new_invalid(),
228 cx,
229 )
230 .into_any_element()
231 }))
232 .into_any_element(),
233 ),
234 single_example(
235 "Successful",
236 div()
237 .size_full()
238 .child(successful_search.update(cx, |tool, cx| {
239 tool.render(
240 &ToolUseStatus::Finished("".into()),
241 window,
242 WeakEntity::new_invalid(),
243 cx,
244 )
245 .into_any_element()
246 }))
247 .into_any_element(),
248 ),
249 single_example(
250 "Error",
251 div()
252 .size_full()
253 .child(error_search.update(cx, |tool, cx| {
254 tool.render(
255 &ToolUseStatus::Error("".into()),
256 window,
257 WeakEntity::new_invalid(),
258 cx,
259 )
260 .into_any_element()
261 }))
262 .into_any_element(),
263 ),
264 ])])
265 .into_any_element(),
266 )
267 }
268}
269
270fn example_search_response() -> WebSearchResponse {
271 WebSearchResponse {
272 summary: r#"Toronto boasts a vibrant culinary scene with a diverse array of..."#
273 .to_string(),
274 citations: vec![
275 WebSearchCitation {
276 title: "Alo".to_string(),
277 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
278 range: Some(147..213),
279 },
280 WebSearchCitation {
281 title: "Edulis".to_string(),
282 url: "https://www.google.com/maps/search/Edulis%2C+Toronto%2C+Canada".to_string(),
283 range: Some(447..519),
284 },
285 WebSearchCitation {
286 title: "Sushi Masaki Saito".to_string(),
287 url: "https://www.google.com/maps/search/Sushi+Masaki+Saito%2C+Toronto%2C+Canada"
288 .to_string(),
289 range: Some(776..872),
290 },
291 WebSearchCitation {
292 title: "Shoushin".to_string(),
293 url: "https://www.google.com/maps/search/Shoushin%2C+Toronto%2C+Canada".to_string(),
294 range: Some(1072..1148),
295 },
296 WebSearchCitation {
297 title: "Restaurant 20 Victoria".to_string(),
298 url:
299 "https://www.google.com/maps/search/Restaurant+20+Victoria%2C+Toronto%2C+Canada"
300 .to_string(),
301 range: Some(1291..1395),
302 },
303 ],
304 }
305}