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::{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 _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.results.len() == 1 {
124 "1 result".into()
125 } else {
126 format!("{} results", response.results.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 = self.response.as_ref().and_then(|response| match response {
138 Ok(response) => Some(
139 v_flex()
140 .overflow_hidden()
141 .ml_1p5()
142 .pl(px(5.))
143 .border_l_1()
144 .border_color(cx.theme().colors().border_variant)
145 .gap_1()
146 .children(response.results.iter().enumerate().map(|(index, result)| {
147 let title = result.title.clone();
148 let url = result.url.clone();
149
150 Button::new(("result", index), title)
151 .label_size(LabelSize::Small)
152 .color(Color::Muted)
153 .icon(IconName::ArrowUpRight)
154 .icon_size(IconSize::XSmall)
155 .icon_position(IconPosition::End)
156 .truncate(true)
157 .tooltip({
158 let url = url.clone();
159 move |window, cx| {
160 Tooltip::with_meta(
161 "Web Search Result",
162 None,
163 url.clone(),
164 window,
165 cx,
166 )
167 }
168 })
169 .on_click({
170 let url = url.clone();
171 move |_, _, cx| cx.open_url(&url)
172 })
173 }))
174 .into_any(),
175 ),
176 Err(_) => None,
177 });
178
179 v_flex().mb_3().gap_1().child(header).children(content)
180 }
181}
182
183impl Component for WebSearchToolCard {
184 fn scope() -> ComponentScope {
185 ComponentScope::Agent
186 }
187
188 fn preview(window: &mut Window, cx: &mut App) -> Option<AnyElement> {
189 let in_progress_search = cx.new(|cx| WebSearchToolCard {
190 response: None,
191 _task: cx.spawn(async move |_this, cx| {
192 loop {
193 cx.background_executor()
194 .timer(Duration::from_secs(60))
195 .await
196 }
197 }),
198 });
199
200 let successful_search = cx.new(|_cx| WebSearchToolCard {
201 response: Some(Ok(example_search_response())),
202 _task: Task::ready(()),
203 });
204
205 let error_search = cx.new(|_cx| WebSearchToolCard {
206 response: Some(Err(anyhow!("Failed to resolve https://google.com"))),
207 _task: Task::ready(()),
208 });
209
210 Some(
211 v_flex()
212 .gap_6()
213 .children(vec![example_group(vec![
214 single_example(
215 "In Progress",
216 div()
217 .size_full()
218 .child(in_progress_search.update(cx, |tool, cx| {
219 tool.render(
220 &ToolUseStatus::Pending,
221 window,
222 WeakEntity::new_invalid(),
223 cx,
224 )
225 .into_any_element()
226 }))
227 .into_any_element(),
228 ),
229 single_example(
230 "Successful",
231 div()
232 .size_full()
233 .child(successful_search.update(cx, |tool, cx| {
234 tool.render(
235 &ToolUseStatus::Finished("".into()),
236 window,
237 WeakEntity::new_invalid(),
238 cx,
239 )
240 .into_any_element()
241 }))
242 .into_any_element(),
243 ),
244 single_example(
245 "Error",
246 div()
247 .size_full()
248 .child(error_search.update(cx, |tool, cx| {
249 tool.render(
250 &ToolUseStatus::Error("".into()),
251 window,
252 WeakEntity::new_invalid(),
253 cx,
254 )
255 .into_any_element()
256 }))
257 .into_any_element(),
258 ),
259 ])])
260 .into_any_element(),
261 )
262 }
263}
264
265fn example_search_response() -> WebSearchResponse {
266 WebSearchResponse {
267 results: vec![
268 WebSearchResult {
269 title: "Alo".to_string(),
270 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
271 text: "Alo is a popular restaurant in Toronto.".to_string(),
272 },
273 WebSearchResult {
274 title: "Alo".to_string(),
275 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
276 text: "Information about Alo restaurant in Toronto.".to_string(),
277 },
278 WebSearchResult {
279 title: "Edulis".to_string(),
280 url: "https://www.google.com/maps/search/Edulis%2C+Toronto%2C+Canada".to_string(),
281 text: "Details about Edulis restaurant in Toronto.".to_string(),
282 },
283 WebSearchResult {
284 title: "Sushi Masaki Saito".to_string(),
285 url: "https://www.google.com/maps/search/Sushi+Masaki+Saito%2C+Toronto%2C+Canada"
286 .to_string(),
287 text: "Information about Sushi Masaki Saito in Toronto.".to_string(),
288 },
289 WebSearchResult {
290 title: "Shoushin".to_string(),
291 url: "https://www.google.com/maps/search/Shoushin%2C+Toronto%2C+Canada".to_string(),
292 text: "Details about Shoushin restaurant in Toronto.".to_string(),
293 },
294 WebSearchResult {
295 title: "Restaurant 20 Victoria".to_string(),
296 url:
297 "https://www.google.com/maps/search/Restaurant+20+Victoria%2C+Toronto%2C+Canada"
298 .to_string(),
299 text: "Information about Restaurant 20 Victoria in Toronto.".to_string(),
300 },
301 ],
302 }
303}