Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

rpc: Remove llm module in favor of zed_llm_client (#28900)

Browse source 
This PR removes the `llm` module of the `rpc` crate in favor of using
the types from the `zed_llm_client`.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-16 16:22:44 -04:00 • committed by GitHub
parent 54b46fdfaa
commit fcb1efdf21
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 45 additions and 85 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/collab/Cargo.toml
View file

@ -75,6 +75,7 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json", "re
util.workspace = true
uuid.workspace = true
workspace-hack.workspace = true
zed_llm_client.workspace = true
[dev-dependencies]
assistant = { workspace = true, features = ["test-support"] }

6
crates/collab/src/api/billing.rs
View file

@ -330,8 +330,10 @@ async fn create_billing_subscription(
.await?
}
None => {
let default_model =
llm_db.model(rpc::LanguageModelProvider::Anthropic, "claude-3-7-sonnet")?;
let default_model = llm_db.model(
zed_llm_client::LanguageModelProvider::Anthropic,
"claude-3-7-sonnet",
)?;
let stripe_model = stripe_billing.register_model(default_model).await?;
stripe_billing
.checkout(customer_id, &user.github_login, &stripe_model, &success_url)

2
crates/collab/src/llm/db.rs
View file

@ -8,9 +8,9 @@ mod tests;
use collections::HashMap;
pub use ids::*;
use rpc::LanguageModelProvider;
pub use seed::*;
pub use tables::*;
use zed_llm_client::LanguageModelProvider;
#[cfg(test)]
pub use tests::TestLlmDb;

2
crates/collab/src/llm/db/tests/provider_tests.rs
View file

@ -1,5 +1,5 @@
use pretty_assertions::assert_eq;
use rpc::LanguageModelProvider;
use zed_llm_client::LanguageModelProvider;
use crate::llm::db::LlmDatabase;
use crate::test_llm_db;

37
crates/language_models/src/provider/cloud.rs
View file

@ -1,9 +1,6 @@
use anthropic::{AnthropicError, AnthropicModelMode, parse_prompt_too_long};
use anyhow::{Result, anyhow};
use client::{
Client, EXPIRED_LLM_TOKEN_HEADER_NAME, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME,
PerformCompletionParams, UserStore, zed_urls,
};
use client::{Client, UserStore, zed_urls};
use collections::BTreeMap;
use feature_flags::{FeatureFlagAppExt, LlmClosedBeta, ZedPro};
use futures::{
@ -26,7 +23,6 @@ use language_model::{
use proto::Plan;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::value::RawValue;
use settings::{Settings, SettingsStore};
use smol::Timer;
use smol::io::{AsyncReadExt, BufReader};
@ -38,7 +34,10 @@ use std::{
use strum::IntoEnumIterator;
use thiserror::Error;
use ui::{TintColor, prelude::*};
use zed_llm_client::{CURRENT_PLAN_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME};
use zed_llm_client::{
CURRENT_PLAN_HEADER_NAME, CompletionBody, EXPIRED_LLM_TOKEN_HEADER_NAME,
MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
};
use crate::AllLanguageModelSettings;
use crate::provider::anthropic::{count_anthropic_tokens, into_anthropic};
@ -517,7 +516,7 @@ impl CloudLanguageModel {
async fn perform_llm_completion(
client: Arc<Client>,
llm_api_token: LlmApiToken,
body: PerformCompletionParams,
body: CompletionBody,
) -> Result<Response<AsyncBody>> {
let http_client = &client.http_client();
@ -724,12 +723,10 @@ impl LanguageModel for CloudLanguageModel {
let response = Self::perform_llm_completion(
client.clone(),
llm_api_token,
PerformCompletionParams {
provider: client::LanguageModelProvider::Anthropic,
CompletionBody {
provider: zed_llm_client::LanguageModelProvider::Anthropic,
model: request.model.clone(),
provider_request: RawValue::from_string(serde_json::to_string(
&request,
)?)?,
provider_request: serde_json::to_value(&request)?,
},
)
.await
@ -765,12 +762,10 @@ impl LanguageModel for CloudLanguageModel {
let response = Self::perform_llm_completion(
client.clone(),
llm_api_token,
PerformCompletionParams {
provider: client::LanguageModelProvider::OpenAi,
CompletionBody {
provider: zed_llm_client::LanguageModelProvider::OpenAi,
model: request.model.clone(),
provider_request: RawValue::from_string(serde_json::to_string(
&request,
)?)?,
provider_request: serde_json::to_value(&request)?,
},
)
.await?;
@ -790,12 +785,10 @@ impl LanguageModel for CloudLanguageModel {
let response = Self::perform_llm_completion(
client.clone(),
llm_api_token,
PerformCompletionParams {
provider: client::LanguageModelProvider::Google,
CompletionBody {
provider: zed_llm_client::LanguageModelProvider::Google,
model: request.model.clone(),
provider_request: RawValue::from_string(serde_json::to_string(
&request,
)?)?,
provider_request: serde_json::to_value(&request)?,
},
)
.await?;

35
crates/rpc/src/llm.rs
View file

@ -1,35 +0,0 @@
use serde::{Deserialize, Serialize};
use strum::{Display, EnumIter, EnumString};
pub const EXPIRED_LLM_TOKEN_HEADER_NAME: &str = "x-zed-expired-token";
pub const MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME: &str = "x-zed-llm-max-monthly-spend-reached";
#[derive(
Debug, PartialEq, Eq, Hash, Clone, Copy, Serialize, Deserialize, EnumString, EnumIter, Display,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum LanguageModelProvider {
Anthropic,
OpenAi,
Google,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct LanguageModel {
pub provider: LanguageModelProvider,
pub name: String,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ListModelsResponse {
pub models: Vec<LanguageModel>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct PerformCompletionParams {
pub provider: LanguageModelProvider,
pub model: String,
pub provider_request: Box<serde_json::value::RawValue>,
}

2
crates/rpc/src/rpc.rs
View file

@ -1,14 +1,12 @@
pub mod auth;
mod conn;
mod extension;
mod llm;
mod message_stream;
mod notification;
mod peer;
pub use conn::Connection;
pub use extension::*;
pub use llm::*;
pub use notification::*;
pub use peer::*;
pub use proto;