use std::sync::Arc; use anyhow::{Context as _, Result};use client::Client;use futures::AsyncReadExt as _;use gpui::{App, AppContext, Context, Entity, Subscription, Task};use http_client::{HttpClient, Method};use language_model::{LlmApiToken, RefreshLlmTokenListener};use web_search::{WebSearchProvider, WebSearchProviderId};use zed_llm_client::{ CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, WebSearchBody, WebSearchResponse,}; pub struct CloudWebSearchProvider { state: Entity<State>,} impl CloudWebSearchProvider { pub fn new(client: Arc<Client>, cx: &mut App) -> Self { let state = cx.new(|cx| State::new(client, cx)); Self { state } }} pub struct State { client: Arc<Client>, llm_api_token: LlmApiToken, _llm_token_subscription: Subscription,} impl State { pub fn new(client: Arc<Client>, cx: &mut Context<Self>) -> Self { let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx); Self { client, llm_api_token: LlmApiToken::default(), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, |this, _, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); cx.spawn(async move |_this, _cx| { llm_api_token.refresh(&client).await?; anyhow::Ok(()) }) .detach_and_log_err(cx); }, ), } }} pub const ZED_WEB_SEARCH_PROVIDER_ID: &'static str = "zed.dev"; impl WebSearchProvider for CloudWebSearchProvider { fn id(&self) -> WebSearchProviderId { WebSearchProviderId(ZED_WEB_SEARCH_PROVIDER_ID.into()) } fn search(&self, query: String, cx: &mut App) -> Task<Result<WebSearchResponse>> { let state = self.state.read(cx); let client = state.client.clone(); let llm_api_token = state.llm_api_token.clone(); let body = WebSearchBody { query }; cx.background_spawn(async move { perform_web_search(client, llm_api_token, body).await }) }} async fn perform_web_search( client: Arc<Client>, llm_api_token: LlmApiToken, body: WebSearchBody,) -> Result<WebSearchResponse> { let http_client = &client.http_client(); let token = llm_api_token.acquire(&client).await?; let request = http_client::Request::builder() .method(Method::POST) .uri(http_client.build_zed_llm_url("/web_search", &[])?.as_ref()) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {token}")) .header(CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, "true") .body(serde_json::to_string(&body)?.into())?; let mut response = http_client .send(request) .await .context("failed to send web search request")?; if response.status().is_success() { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; return Ok(serde_json::from_str(&body)?); } else { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; anyhow::bail!( "error performing web search.\nStatus: {:?}\nBody: {body}", response.status(), ); }}
use std::sync::Arc;
use anyhow::{Context as _, Result};
use client::Client;
use futures::AsyncReadExt as _;
use gpui::{App, AppContext, Context, Entity, Subscription, Task};
use http_client::{HttpClient, Method};
use language_model::{LlmApiToken, RefreshLlmTokenListener};
use web_search::{WebSearchProvider, WebSearchProviderId};
use zed_llm_client::{
CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, WebSearchBody, WebSearchResponse,
};
pub struct CloudWebSearchProvider {
state: Entity<State>,
}
impl CloudWebSearchProvider {
pub fn new(client: Arc<Client>, cx: &mut App) -> Self {
let state = cx.new(|cx| State::new(client, cx));
Self { state }
}
}
pub struct State {
client: Arc<Client>,
llm_api_token: LlmApiToken,
_llm_token_subscription: Subscription,
}
impl State {
pub fn new(client: Arc<Client>, cx: &mut Context<Self>) -> Self {
let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
Self {
client,
llm_api_token: LlmApiToken::default(),
_llm_token_subscription: cx.subscribe(
&refresh_llm_token_listener,
|this, _, _event, cx| {
let client = this.client.clone();
let llm_api_token = this.llm_api_token.clone();
cx.spawn(async move |_this, _cx| {
llm_api_token.refresh(&client).await?;
anyhow::Ok(())
})
.detach_and_log_err(cx);
},
),
}
}
}
pub const ZED_WEB_SEARCH_PROVIDER_ID: &'static str = "zed.dev";
impl WebSearchProvider for CloudWebSearchProvider {
fn id(&self) -> WebSearchProviderId {
WebSearchProviderId(ZED_WEB_SEARCH_PROVIDER_ID.into())
}
fn search(&self, query: String, cx: &mut App) -> Task<Result<WebSearchResponse>> {
let state = self.state.read(cx);
let client = state.client.clone();
let llm_api_token = state.llm_api_token.clone();
let body = WebSearchBody { query };
cx.background_spawn(async move { perform_web_search(client, llm_api_token, body).await })
}
}
async fn perform_web_search(
client: Arc<Client>,
llm_api_token: LlmApiToken,
body: WebSearchBody,
) -> Result<WebSearchResponse> {
let http_client = &client.http_client();
let token = llm_api_token.acquire(&client).await?;
let request = http_client::Request::builder()
.method(Method::POST)
.uri(http_client.build_zed_llm_url("/web_search", &[])?.as_ref())
.header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {token}"))
.header(CLIENT_SUPPORTS_EXA_WEB_SEARCH_PROVIDER_HEADER_NAME, "true")
.body(serde_json::to_string(&body)?.into())?;
let mut response = http_client
.send(request)
.await
.context("failed to send web search request")?;
if response.status().is_success() {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
return Ok(serde_json::from_str(&body)?);
} else {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"error performing web search.\nStatus: {:?}\nBody: {body}",
response.status(),
);
}
}