tenant.openagents/omega

1use std::sync::Arc;
2
3use crate::{AgentTool, ToolCallEventStream, ToolInput};
4use agent_client_protocol::schema::v1 as acp;
5use anyhow::Result;
6use cloud_llm_client::WebSearchResponse;
7use futures::FutureExt as _;
8use gpui::{App, Task};
9use language_model::{
10    LanguageModelProviderId, LanguageModelToolResultContent, ZED_CLOUD_PROVIDER_ID,
11};
12use schemars::JsonSchema;
13use serde::{Deserialize, Serialize};
14use ui::prelude::*;
15use util::markdown::MarkdownInlineCode;
16use web_search::WebSearchRegistry;
17
18/// Search the web for information using your query.
19/// Use this when you need real-time information, facts, or data that might not be in your training.
20/// Results will include snippets and links from relevant web pages.
21#[derive(Debug, Serialize, Deserialize, JsonSchema)]
22pub struct WebSearchToolInput {
23    /// The search term or question to query on the web.
24    query: String,
25}
26
27#[derive(Debug, Serialize, Deserialize)]
28#[serde(untagged)]
29pub enum WebSearchToolOutput {
30    Success(WebSearchResponse),
31    Error { error: String },
32}
33
34impl From<WebSearchToolOutput> for LanguageModelToolResultContent {
35    fn from(value: WebSearchToolOutput) -> Self {
36        match value {
37            WebSearchToolOutput::Success(response) => serde_json::to_string(&response)
38                .unwrap_or_else(|e| format!("Failed to serialize web search response: {e}"))
39                .into(),
40            WebSearchToolOutput::Error { error } => error.into(),
41        }
42    }
43}
44
45pub struct WebSearchTool;
46
47impl AgentTool for WebSearchTool {
48    type Input = WebSearchToolInput;
49    type Output = WebSearchToolOutput;
50
51    const NAME: &'static str = "search_web";
52
53    fn kind() -> acp::ToolKind {
54        acp::ToolKind::Fetch
55    }
56
57    fn initial_title(
58        &self,
59        _input: Result<Self::Input, serde_json::Value>,
60        _cx: &mut App,
61    ) -> SharedString {
62        "Searching the Web".into()
63    }
64
65    /// We currently only support Zed Cloud as a provider.
66    fn supports_provider(provider: &LanguageModelProviderId) -> bool {
67        provider == &ZED_CLOUD_PROVIDER_ID
68    }
69
70    fn run(
71        self: Arc<Self>,
72        input: ToolInput<Self::Input>,
73        event_stream: ToolCallEventStream,
74        cx: &mut App,
75    ) -> Task<Result<Self::Output, Self::Output>> {
76        cx.spawn(async move |cx| {
77            let input = input
78                .recv()
79                .await
80                .map_err(|e| WebSearchToolOutput::Error {
81                    error: e.to_string(),
82                })?;
83
84            let authorize = cx.update(|cx| {
85                let context =
86                    crate::ToolPermissionContext::new(Self::NAME, vec![input.query.clone()]);
87                event_stream.authorize(
88                    format!("Search the web for {}", MarkdownInlineCode(&input.query)),
89                    context,
90                    cx,
91                )
92            });
93            authorize
94                .await
95                .map_err(|e| WebSearchToolOutput::Error { error: e.to_string() })?;
96
97            let search_task = cx.update(|cx| {
98                let Some(provider) = WebSearchRegistry::read_global(cx).active_provider() else {
99                    return Err(WebSearchToolOutput::Error {
100                        error: "Web search is not available.".to_string(),
101                    });
102                };
103                Ok(provider.search(input.query, cx))
104            })?;
105
106            let response = futures::select! {
107                result = search_task.fuse() => {
108                    match result {
109                        Ok(response) => response,
110                        Err(err) => {
111                            event_stream
112                                .update_fields(acp::ToolCallUpdateFields::new().title("Web Search Failed"));
113                            return Err(WebSearchToolOutput::Error { error: err.to_string() });
114                        }
115                    }
116                }
117                _ = event_stream.cancelled_by_user().fuse() => {
118                    return Err(WebSearchToolOutput::Error { error: "Web search cancelled by user".to_string() });
119                }
120            };
121
122            emit_update(&response, &event_stream);
123            Ok(WebSearchToolOutput::Success(response))
124        })
125    }
126
127    fn replay(
128        &self,
129        _input: Self::Input,
130        output: Self::Output,
131        event_stream: ToolCallEventStream,
132        _cx: &mut App,
133    ) -> Result<()> {
134        if let WebSearchToolOutput::Success(response) = &output {
135            emit_update(response, &event_stream);
136        }
137        Ok(())
138    }
139}
140
141fn emit_update(response: &WebSearchResponse, event_stream: &ToolCallEventStream) {
142    let result_text = if response.results.len() == 1 {
143        "1 result".to_string()
144    } else {
145        format!("{} results", response.results.len())
146    };
147    event_stream.update_fields(
148        acp::ToolCallUpdateFields::new()
149            .title(format!("Searched the web: {result_text}"))
150            .content(
151                response
152                    .results
153                    .iter()
154                    .map(|result| {
155                        acp::ToolCallContent::Content(acp::Content::new(
156                            acp::ContentBlock::ResourceLink(
157                                acp::ResourceLink::new(result.title.clone(), result.url.clone())
158                                    .title(result.title.clone())
159                                    .description(result.text.clone()),
160                            ),
161                        ))
162                    })
163                    .collect::<Vec<_>>(),
164            ),
165    );
166}
167