tenant.openagents/omega

1use std::sync::Arc;
2
3use anyhow::{Result, anyhow};
4use client::{Client, UserStore, global_llm_token};
5use cloud_api_client::LlmApiToken;
6use cloud_api_types::OrganizationId;
7use cloud_llm_client::{WebSearchBody, WebSearchResponse};
8use futures::AsyncReadExt as _;
9use gpui::{App, AppContext, Context, Entity, Task};
10use http_client::Method;
11use web_search::{WebSearchProvider, WebSearchProviderId};
12
13pub struct CloudWebSearchProvider {
14    state: Entity<State>,
15}
16
17impl CloudWebSearchProvider {
18    pub fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) -> Self {
19        let state = cx.new(|cx| State::new(client, user_store, cx));
20
21        Self { state }
22    }
23}
24
25pub struct State {
26    client: Arc<Client>,
27    user_store: Entity<UserStore>,
28    llm_api_token: LlmApiToken,
29}
30
31impl State {
32    pub fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
33        let llm_api_token = global_llm_token(cx);
34
35        Self {
36            client,
37            user_store,
38            llm_api_token,
39        }
40    }
41}
42
43pub const ZED_WEB_SEARCH_PROVIDER_ID: &str = "zed.dev";
44
45impl WebSearchProvider for CloudWebSearchProvider {
46    fn id(&self) -> WebSearchProviderId {
47        WebSearchProviderId(ZED_WEB_SEARCH_PROVIDER_ID.into())
48    }
49
50    fn search(&self, query: String, cx: &mut App) -> Task<Result<WebSearchResponse>> {
51        let state = self.state.read(cx);
52        let client = state.client.clone();
53        let llm_api_token = state.llm_api_token.clone();
54        let organization_id = state
55            .user_store
56            .read(cx)
57            .current_organization()
58            .map(|organization| organization.id.clone());
59        let body = WebSearchBody { query };
60        cx.background_spawn(async move {
61            perform_web_search(client, llm_api_token, organization_id, body).await
62        })
63    }
64}
65
66async fn perform_web_search(
67    client: Arc<Client>,
68    llm_api_token: LlmApiToken,
69    organization_id: Option<OrganizationId>,
70    body: WebSearchBody,
71) -> Result<WebSearchResponse> {
72    let organization_id = organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
73
74    let url = client.http_client().build_zed_llm_url("/web_search", &[])?;
75    let body = serde_json::to_string(&body)?;
76    let mut response = client
77        .authenticated_llm_request(&llm_api_token, organization_id, |token| {
78            Ok(http_client::Request::builder()
79                .method(Method::POST)
80                .uri(url.as_ref())
81                .header("Content-Type", "application/json")
82                .header("Authorization", format!("Bearer {token}"))
83                .body(body.clone().into())?)
84        })
85        .await?;
86
87    if response.status().is_success() {
88        let mut body = String::new();
89        response.body_mut().read_to_string(&mut body).await?;
90        Ok(serde_json::from_str(&body)?)
91    } else {
92        let status = response.status();
93        let mut body = String::new();
94        response.body_mut().read_to_string(&mut body).await?;
95        anyhow::bail!("error performing web search.\nStatus: {status:?}\nBody: {body}");
96    }
97}
98