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