Enable Claude 3 models to be used via the Zed server if "language-mod…

…els" feature flag is enabled for user (zed-industries#10015)

Release Notes:

- N/A

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "crates/activity_indicator",
+    "crates/anthropic",
     "crates/assets",
     "crates/assistant",
     "crates/audio",
@@ -119,6 +120,7 @@ resolver = "2"
 [workspace.dependencies]
 activity_indicator = { path = "crates/activity_indicator" }
 ai = { path = "crates/ai" }
+anthropic = { path = "crates/anthropic" }
 assets = { path = "crates/assets" }
 assistant = { path = "crates/assistant" }
 audio = { path = "crates/audio" }

crates/anthropic/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "anthropic"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+path = "src/anthropic.rs"
+
+[dependencies]
+anyhow.workspace = true
+futures.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+util.workspace = true
+
+[dev-dependencies]
+tokio.workspace = true
+
+[lints]
+workspace = true

crates/anthropic/src/anthropic.rs

@@ -0,0 +1,234 @@
+use anyhow::{anyhow, Result};
+use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, StreamExt};
+use serde::{Deserialize, Serialize};
+use std::convert::TryFrom;
+use util::http::{AsyncBody, HttpClient, Method, Request as HttpRequest};
+
+#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
+pub enum Model {
+    #[default]
+    #[serde(rename = "claude-3-opus-20240229")]
+    Claude3Opus,
+    #[serde(rename = "claude-3-sonnet-20240229")]
+    Claude3Sonnet,
+    #[serde(rename = "claude-3-haiku-20240307")]
+    Claude3Haiku,
+}
+
+impl Model {
+    pub fn from_id(id: &str) -> Result<Self> {
+        if id.starts_with("claude-3-opus") {
+            Ok(Self::Claude3Opus)
+        } else if id.starts_with("claude-3-sonnet") {
+            Ok(Self::Claude3Sonnet)
+        } else if id.starts_with("claude-3-haiku") {
+            Ok(Self::Claude3Haiku)
+        } else {
+            Err(anyhow!("Invalid model id: {}", id))
+        }
+    }
+
+    pub fn display_name(&self) -> &'static str {
+        match self {
+            Self::Claude3Opus => "Claude 3 Opus",
+            Self::Claude3Sonnet => "Claude 3 Sonnet",
+            Self::Claude3Haiku => "Claude 3 Haiku",
+        }
+    }
+
+    pub fn max_token_count(&self) -> usize {
+        200_000
+    }
+}
+
+#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    User,
+    Assistant,
+}
+
+impl TryFrom<String> for Role {
+    type Error = anyhow::Error;
+
+    fn try_from(value: String) -> Result<Self> {
+        match value.as_str() {
+            "user" => Ok(Self::User),
+            "assistant" => Ok(Self::Assistant),
+            _ => Err(anyhow!("invalid role '{value}'")),
+        }
+    }
+}
+
+impl From<Role> for String {
+    fn from(val: Role) -> Self {
+        match val {
+            Role::User => "user".to_owned(),
+            Role::Assistant => "assistant".to_owned(),
+        }
+    }
+}
+
+#[derive(Debug, Serialize)]
+pub struct Request {
+    pub model: Model,
+    pub messages: Vec<RequestMessage>,
+    pub stream: bool,
+    pub system: String,
+    pub max_tokens: u32,
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+pub struct RequestMessage {
+    pub role: Role,
+    pub content: String,
+}
+
+#[derive(Deserialize, Debug)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ResponseEvent {
+    MessageStart {
+        message: ResponseMessage,
+    },
+    ContentBlockStart {
+        index: u32,
+        content_block: ContentBlock,
+    },
+    Ping {},
+    ContentBlockDelta {
+        index: u32,
+        delta: TextDelta,
+    },
+    ContentBlockStop {
+        index: u32,
+    },
+    MessageDelta {
+        delta: ResponseMessage,
+        usage: Usage,
+    },
+    MessageStop {},
+}
+
+#[derive(Deserialize, Debug)]
+pub struct ResponseMessage {
+    #[serde(rename = "type")]
+    pub message_type: Option<String>,
+    pub id: Option<String>,
+    pub role: Option<String>,
+    pub content: Option<Vec<String>>,
+    pub model: Option<String>,
+    pub stop_reason: Option<String>,
+    pub stop_sequence: Option<String>,
+    pub usage: Option<Usage>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Usage {
+    pub input_tokens: Option<u32>,
+    pub output_tokens: Option<u32>,
+}
+
+#[derive(Deserialize, Debug)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ContentBlock {
+    Text { text: String },
+}
+
+#[derive(Deserialize, Debug)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum TextDelta {
+    TextDelta { text: String },
+}
+
+pub async fn stream_completion(
+    client: &dyn HttpClient,
+    api_url: &str,
+    api_key: &str,
+    request: Request,
+) -> Result<BoxStream<'static, Result<ResponseEvent>>> {
+    let uri = format!("{api_url}/v1/messages");
+    let request = HttpRequest::builder()
+        .method(Method::POST)
+        .uri(uri)
+        .header("Anthropic-Version", "2023-06-01")
+        .header("Anthropic-Beta", "messages-2023-12-15")
+        .header("X-Api-Key", api_key)
+        .header("Content-Type", "application/json")
+        .body(AsyncBody::from(serde_json::to_string(&request)?))?;
+    let mut response = client.send(request).await?;
+    if response.status().is_success() {
+        let reader = BufReader::new(response.into_body());
+        Ok(reader
+            .lines()
+            .filter_map(|line| async move {
+                match line {
+                    Ok(line) => {
+                        let line = line.strip_prefix("data: ")?;
+                        match serde_json::from_str(line) {
+                            Ok(response) => Some(Ok(response)),
+                            Err(error) => Some(Err(anyhow!(error))),
+                        }
+                    }
+                    Err(error) => Some(Err(anyhow!(error))),
+                }
+            })
+            .boxed())
+    } else {
+        let mut body = Vec::new();
+        response.body_mut().read_to_end(&mut body).await?;
+
+        let body_str = std::str::from_utf8(&body)?;
+
+        match serde_json::from_str::<ResponseEvent>(body_str) {
+            Ok(_) => Err(anyhow!(
+                "Unexpected success response while expecting an error: {}",
+                body_str,
+            )),
+            Err(_) => Err(anyhow!(
+                "Failed to connect to API: {} {}",
+                response.status(),
+                body_str,
+            )),
+        }
+    }
+}
+
+// #[cfg(test)]
+// mod tests {
+//     use super::*;
+//     use util::http::IsahcHttpClient;
+
+//     #[tokio::test]
+//     async fn stream_completion_success() {
+//         let http_client = IsahcHttpClient::new().unwrap();
+
+//         let request = Request {
+//             model: Model::Claude3Opus,
+//             messages: vec![RequestMessage {
+//                 role: Role::User,
+//                 content: "Ping".to_string(),
+//             }],
+//             stream: true,
+//             system: "Respond to ping with pong".to_string(),
+//             max_tokens: 4096,
+//         };
+
+//         let stream = stream_completion(
+//             &http_client,
+//             "https://api.anthropic.com",
+//             &std::env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY not set"),
+//             request,
+//         )
+//         .await
+//         .unwrap();
+
+//         stream
+//             .for_each(|event| async {
+//                 match event {
+//                     Ok(event) => println!("{:?}", event),
+//                     Err(e) => eprintln!("Error: {:?}", e),
+//                 }
+//             })
+//             .await;
+//     }
+// }

crates/assistant/src/assistant_panel.rs

@@ -768,15 +768,18 @@ impl AssistantPanel {
                 open_ai::Model::FourTurbo => open_ai::Model::ThreePointFiveTurbo,
             }),
             LanguageModel::ZedDotDev(model) => LanguageModel::ZedDotDev(match &model {
-                ZedDotDevModel::GptThreePointFiveTurbo => ZedDotDevModel::GptFour,
-                ZedDotDevModel::GptFour => ZedDotDevModel::GptFourTurbo,
-                ZedDotDevModel::GptFourTurbo => {
+                ZedDotDevModel::Gpt3Point5Turbo => ZedDotDevModel::Gpt4,
+                ZedDotDevModel::Gpt4 => ZedDotDevModel::Gpt4Turbo,
+                ZedDotDevModel::Gpt4Turbo => ZedDotDevModel::Claude3Opus,
+                ZedDotDevModel::Claude3Opus => ZedDotDevModel::Claude3Sonnet,
+                ZedDotDevModel::Claude3Sonnet => ZedDotDevModel::Claude3Haiku,
+                ZedDotDevModel::Claude3Haiku => {
                     match CompletionProvider::global(cx).default_model() {
                         LanguageModel::ZedDotDev(custom) => custom,
-                        _ => ZedDotDevModel::GptThreePointFiveTurbo,
+                        _ => ZedDotDevModel::Gpt3Point5Turbo,
                     }
                 }
-                ZedDotDevModel::Custom(_) => ZedDotDevModel::GptThreePointFiveTurbo,
+                ZedDotDevModel::Custom(_) => ZedDotDevModel::Gpt3Point5Turbo,
             }),
         };
 

crates/assistant/src/assistant_settings.rs

@@ -14,10 +14,13 @@ use settings::Settings;
 
 #[derive(Clone, Debug, Default, PartialEq)]
 pub enum ZedDotDevModel {
-    GptThreePointFiveTurbo,
-    GptFour,
+    Gpt3Point5Turbo,
+    Gpt4,
     #[default]
-    GptFourTurbo,
+    Gpt4Turbo,
+    Claude3Opus,
+    Claude3Sonnet,
+    Claude3Haiku,
     Custom(String),
 }
 
@@ -49,9 +52,9 @@ impl<'de> Deserialize<'de> for ZedDotDevModel {
                 E: de::Error,
             {
                 match value {
-                    "gpt-3.5-turbo" => Ok(ZedDotDevModel::GptThreePointFiveTurbo),
-                    "gpt-4" => Ok(ZedDotDevModel::GptFour),
-                    "gpt-4-turbo-preview" => Ok(ZedDotDevModel::GptFourTurbo),
+                    "gpt-3.5-turbo" => Ok(ZedDotDevModel::Gpt3Point5Turbo),
+                    "gpt-4" => Ok(ZedDotDevModel::Gpt4),
+                    "gpt-4-turbo-preview" => Ok(ZedDotDevModel::Gpt4Turbo),
                     _ => Ok(ZedDotDevModel::Custom(value.to_owned())),
                 }
             }
@@ -94,27 +97,34 @@ impl JsonSchema for ZedDotDevModel {
 impl ZedDotDevModel {
     pub fn id(&self) -> &str {
         match self {
-            Self::GptThreePointFiveTurbo => "gpt-3.5-turbo",
-            Self::GptFour => "gpt-4",
-            Self::GptFourTurbo => "gpt-4-turbo-preview",
+            Self::Gpt3Point5Turbo => "gpt-3.5-turbo",
+            Self::Gpt4 => "gpt-4",
+            Self::Gpt4Turbo => "gpt-4-turbo-preview",
+            Self::Claude3Opus => "claude-3-opus",
+            Self::Claude3Sonnet => "claude-3-sonnet",
+            Self::Claude3Haiku => "claude-3-haiku",
             Self::Custom(id) => id,
         }
     }
 
     pub fn display_name(&self) -> &str {
         match self {
-            Self::GptThreePointFiveTurbo => "gpt-3.5-turbo",
-            Self::GptFour => "gpt-4",
-            Self::GptFourTurbo => "gpt-4-turbo",
+            Self::Gpt3Point5Turbo => "GPT 3.5 Turbo",
+            Self::Gpt4 => "GPT 4",
+            Self::Gpt4Turbo => "GPT 4 Turbo",
+            Self::Claude3Opus => "Claude 3 Opus",
+            Self::Claude3Sonnet => "Claude 3 Sonnet",
+            Self::Claude3Haiku => "Claude 3 Haiku",
             Self::Custom(id) => id.as_str(),
         }
     }
 
     pub fn max_token_count(&self) -> usize {
         match self {
-            Self::GptThreePointFiveTurbo => 2048,
-            Self::GptFour => 4096,
-            Self::GptFourTurbo => 128000,
+            Self::Gpt3Point5Turbo => 2048,
+            Self::Gpt4 => 4096,
+            Self::Gpt4Turbo => 128000,
+            Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3Haiku => 200000,
             Self::Custom(_) => 4096, // TODO: Make this configurable
         }
     }