Marshall Bowers
f21780cef3
This PR removes the individual URL overrides for the LLM service. We initially had `ZED_PREDICT_EDITS_URL` to allow for directing traffic to the LLM Worker back when there was still the split of the Collab-based LLM Service and the Cloudflare-based LLM Worker. But now that all of the LLM functionality has been moved into the Worker, we can just direct all traffic there. Release Notes: - N/A
105 lines
3.4 KiB
Rust
105 lines
3.4 KiB
Rust
use std::sync::Arc;
|
|
|
|
use anyhow::{Context as _, Result, anyhow};
|
|
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?;
|
|
return Err(anyhow!(
|
|
"error performing web search.\nStatus: {:?}\nBody: {body}",
|
|
response.status(),
|
|
));
|
|
}
|
|
}
|