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)
76 .context("Failed to serialize search results")
77 .map(Into::into)
78 }
79 });
80
81 ToolResult {
82 output,
83 card: Some(cx.new(|cx| WebSearchToolCard::new(search_task, cx)).into()),
84 }
85 }
86}
87
88#[derive(RegisterComponent)]
89struct WebSearchToolCard {
90 response: Option<Result<WebSearchResponse>>,
91 _task: Task<()>,
92}
93
94impl WebSearchToolCard {
95 fn new(
96 search_task: impl 'static + Future<Output = Result<WebSearchResponse, Arc<anyhow::Error>>>,
97 cx: &mut Context<Self>,
98 ) -> Self {
99 let _task = cx.spawn(async move |this, cx| {
100 let response = search_task.await.map_err(|err| anyhow!(err));
101 this.update(cx, |this, cx| {
102 this.response = Some(response);
103 cx.notify();
104 })
105 .ok();
106 });
107
108 Self {
109 response: None,
110 _task,
111 }
112 }
113}
114
115impl ToolCard for WebSearchToolCard {
116 fn render(
117 &mut self,
118 _status: &ToolUseStatus,
119 _window: &mut Window,
120 _workspace: WeakEntity<Workspace>,
121 cx: &mut Context<Self>,
122 ) -> impl IntoElement {
123 let header = match self.response.as_ref() {
124 Some(Ok(response)) => {
125 let text: SharedString = if response.results.len() == 1 {
126 "1 result".into()
127 } else {
128 format!("{} results", response.results.len()).into()
129 };
130 ToolCallCardHeader::new(IconName::Globe, "Searched the Web")
131 .with_secondary_text(text)
132 }
133 Some(Err(error)) => {
134 ToolCallCardHeader::new(IconName::Globe, "Web Search").with_error(error.to_string())
135 }
136 None => ToolCallCardHeader::new(IconName::Globe, "Searching the Web").loading(),
137 };
138
139 let content = self.response.as_ref().and_then(|response| match response {
140 Ok(response) => 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.results.iter().enumerate().map(|(index, result)| {
149 let title = result.title.clone();
150 let url = result.url.clone();
151
152 Button::new(("result", index), title)
153 .label_size(LabelSize::Small)
154 .color(Color::Muted)
155 .icon(IconName::ArrowUpRight)
156 .icon_size(IconSize::XSmall)
157 .icon_position(IconPosition::End)
158 .truncate(true)
159 .tooltip({
160 let url = url.clone();
161 move |window, cx| {
162 Tooltip::with_meta(
163 "Web Search Result",
164 None,
165 url.clone(),
166 window,
167 cx,
168 )
169 }
170 })
171 .on_click({
172 let url = url.clone();
173 move |_, _, cx| cx.open_url(&url)
174 })
175 }))
176 .into_any(),
177 ),
178 Err(_) => None,
179 });
180
181 v_flex().mb_3().gap_1().child(header).children(content)
182 }
183}
184
185impl Component for WebSearchToolCard {
186 fn scope() -> ComponentScope {
187 ComponentScope::Agent
188 }
189
190 fn preview(window: &mut Window, cx: &mut App) -> Option<AnyElement> {
191 let in_progress_search = cx.new(|cx| WebSearchToolCard {
192 response: None,
193 _task: cx.spawn(async move |_this, cx| {
194 loop {
195 cx.background_executor()
196 .timer(Duration::from_secs(60))
197 .await
198 }
199 }),
200 });
201
202 let successful_search = cx.new(|_cx| WebSearchToolCard {
203 response: Some(Ok(example_search_response())),
204 _task: Task::ready(()),
205 });
206
207 let error_search = cx.new(|_cx| WebSearchToolCard {
208 response: Some(Err(anyhow!("Failed to resolve https://google.com"))),
209 _task: Task::ready(()),
210 });
211
212 Some(
213 v_flex()
214 .gap_6()
215 .children(vec![example_group(vec![
216 single_example(
217 "In Progress",
218 div()
219 .size_full()
220 .child(in_progress_search.update(cx, |tool, cx| {
221 tool.render(
222 &ToolUseStatus::Pending,
223 window,
224 WeakEntity::new_invalid(),
225 cx,
226 )
227 .into_any_element()
228 }))
229 .into_any_element(),
230 ),
231 single_example(
232 "Successful",
233 div()
234 .size_full()
235 .child(successful_search.update(cx, |tool, cx| {
236 tool.render(
237 &ToolUseStatus::Finished("".into()),
238 window,
239 WeakEntity::new_invalid(),
240 cx,
241 )
242 .into_any_element()
243 }))
244 .into_any_element(),
245 ),
246 single_example(
247 "Error",
248 div()
249 .size_full()
250 .child(error_search.update(cx, |tool, cx| {
251 tool.render(
252 &ToolUseStatus::Error("".into()),
253 window,
254 WeakEntity::new_invalid(),
255 cx,
256 )
257 .into_any_element()
258 }))
259 .into_any_element(),
260 ),
261 ])])
262 .into_any_element(),
263 )
264 }
265}
266
267fn example_search_response() -> WebSearchResponse {
268 WebSearchResponse {
269 results: vec![
270 WebSearchResult {
271 title: "Alo".to_string(),
272 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
273 text: "Alo is a popular restaurant in Toronto.".to_string(),
274 },
275 WebSearchResult {
276 title: "Alo".to_string(),
277 url: "https://www.google.com/maps/search/Alo%2C+Toronto%2C+Canada".to_string(),
278 text: "Information about Alo restaurant in Toronto.".to_string(),
279 },
280 WebSearchResult {
281 title: "Edulis".to_string(),
282 url: "https://www.google.com/maps/search/Edulis%2C+Toronto%2C+Canada".to_string(),
283 text: "Details about Edulis restaurant in Toronto.".to_string(),
284 },
285 WebSearchResult {
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 text: "Information about Sushi Masaki Saito in Toronto.".to_string(),
290 },
291 WebSearchResult {
292 title: "Shoushin".to_string(),
293 url: "https://www.google.com/maps/search/Shoushin%2C+Toronto%2C+Canada".to_string(),
294 text: "Details about Shoushin restaurant in Toronto.".to_string(),
295 },
296 WebSearchResult {
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 text: "Information about Restaurant 20 Victoria in Toronto.".to_string(),
302 },
303 ],
304 }
305}