oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix issues with Claude in Assistant2 (#12619)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Nathan <nathan@zed.dev>
This commit is contained in:
Mikayla Maki
2024-06-03 16:30:09 -07:00
committed by GitHub
co-authored by Nathan
parent afc0650a49
commit 3cd6719b30
11 changed files with 276 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/assistant/src/completion_provider/anthropic.rs
+66 -39
View File
@@ -1,9 +1,9 @@
use crate::count_open_ai_tokens;
use crate::{
assistant_settings::AnthropicModel, CompletionProvider, LanguageModel, LanguageModelRequest,
Role,
};
use anthropic::{stream_completion, Request, RequestMessage, Role as AnthropicRole};
use crate::{count_open_ai_tokens, LanguageModelRequestMessage};
use anthropic::{stream_completion, Request, RequestMessage};
use anyhow::{anyhow, Result};
use editor::{Editor, EditorElement, EditorStyle};
use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
@@ -167,53 +167,37 @@ impl AnthropicCompletionProvider {
.boxed()
}
fn to_anthropic_request(&self, request: LanguageModelRequest) -> Request {
fn to_anthropic_request(&self, mut request: LanguageModelRequest) -> Request {
preprocess_anthropic_request(&mut request);
let model = match request.model {
LanguageModel::Anthropic(model) => model,
_ => self.model(),
};
let mut system_message = String::new();
let mut messages: Vec<RequestMessage> = Vec::new();
for message in request.messages {
if message.content.is_empty() {
continue;
}
match message.role {
Role::User | Role::Assistant => {
let role = match message.role {
Role::User => AnthropicRole::User,
Role::Assistant => AnthropicRole::Assistant,
_ => unreachable!(),
};
if let Some(last_message) = messages.last_mut() {
if last_message.role == role {
last_message.content.push_str("\n\n");
last_message.content.push_str(&message.content);
continue;
}
}
messages.push(RequestMessage {
role,
content: message.content,
});
}
Role::System => {
if !system_message.is_empty() {
system_message.push_str("\n\n");
}
system_message.push_str(&message.content);
}
}
if request
.messages
.first()
.map_or(false, |message| message.role == Role::System)
{
system_message = request.messages.remove(0).content;
}
Request {
model,
messages,
messages: request
.messages
.iter()
.map(|msg| RequestMessage {
role: match msg.role {
Role::User => anthropic::Role::User,
Role::Assistant => anthropic::Role::Assistant,
Role::System => unreachable!("filtered out by preprocess_request"),
},
content: msg.content.clone(),
})
.collect(),
stream: true,
system: system_message,
max_tokens: 4092,
@@ -221,6 +205,49 @@ impl AnthropicCompletionProvider {
}
}
pub fn preprocess_anthropic_request(request: &mut LanguageModelRequest) {
let mut new_messages: Vec<LanguageModelRequestMessage> = Vec::new();
let mut system_message = String::new();
for message in request.messages.drain(..) {
if message.content.is_empty() {
continue;
}
match message.role {
Role::User | Role::Assistant => {
if let Some(last_message) = new_messages.last_mut() {
if last_message.role == message.role {
last_message.content.push_str("\n\n");
last_message.content.push_str(&message.content);
continue;
}
}
new_messages.push(message);
}
Role::System => {
if !system_message.is_empty() {
system_message.push_str("\n\n");
}
system_message.push_str(&message.content);
}
}
}
if !system_message.is_empty() {
request.messages.insert(
0,
LanguageModelRequestMessage {
role: Role::System,
content: system_message,
},
);
}
request.messages = new_messages;
}
struct AuthenticationPrompt {
api_key: View<Editor>,
api_url: String,
crates/assistant/src/completion_provider/zed.rs → crates/assistant/src/completion_provider/cloud.rs
+23 -23
View File
@@ -1,5 +1,5 @@
use crate::{
assistant_settings::ZedDotDevModel, count_open_ai_tokens, CompletionProvider, LanguageModel,
assistant_settings::CloudModel, count_open_ai_tokens, CompletionProvider, LanguageModel,
LanguageModelRequest,
};
use anyhow::{anyhow, Result};
@@ -10,17 +10,17 @@ use std::{future, sync::Arc};
use strum::IntoEnumIterator;
use ui::prelude::*;
pub struct ZedDotDevCompletionProvider {
pub struct CloudCompletionProvider {
client: Arc<Client>,
model: ZedDotDevModel,
model: CloudModel,
settings_version: usize,
status: client::Status,
_maintain_client_status: Task<()>,
}
impl ZedDotDevCompletionProvider {
impl CloudCompletionProvider {
pub fn new(
model: ZedDotDevModel,
model: CloudModel,
client: Arc<Client>,
settings_version: usize,
cx: &mut AppContext,
@@ -30,7 +30,7 @@ impl ZedDotDevCompletionProvider {
let maintain_client_status = cx.spawn(|mut cx| async move {
while let Some(status) = status_rx.next().await {
let _ = cx.update_global::<CompletionProvider, _>(|provider, _cx| {
if let CompletionProvider::ZedDotDev(provider) = provider {
if let CompletionProvider::Cloud(provider) = provider {
provider.status = status;
} else {
unreachable!()
@@ -47,20 +47,20 @@ impl ZedDotDevCompletionProvider {
}
}
pub fn update(&mut self, model: ZedDotDevModel, settings_version: usize) {
pub fn update(&mut self, model: CloudModel, settings_version: usize) {
self.model = model;
self.settings_version = settings_version;
}
pub fn available_models(&self) -> impl Iterator<Item = ZedDotDevModel> {
let mut custom_model = if let ZedDotDevModel::Custom(custom_model) = self.model.clone() {
pub fn available_models(&self) -> impl Iterator<Item = CloudModel> {
let mut custom_model = if let CloudModel::Custom(custom_model) = self.model.clone() {
Some(custom_model)
} else {
None
};
ZedDotDevModel::iter().filter_map(move |model| {
if let ZedDotDevModel::Custom(_) = model {
Some(ZedDotDevModel::Custom(custom_model.take()?))
CloudModel::iter().filter_map(move |model| {
if let CloudModel::Custom(_) = model {
Some(CloudModel::Custom(custom_model.take()?))
} else {
Some(model)
}
@@ -71,7 +71,7 @@ impl ZedDotDevCompletionProvider {
self.settings_version
}
pub fn model(&self) -> ZedDotDevModel {
pub fn model(&self) -> CloudModel {
self.model.clone()
}
@@ -94,21 +94,19 @@ impl ZedDotDevCompletionProvider {
cx: &AppContext,
) -> BoxFuture<'static, Result<usize>> {
match request.model {
LanguageModel::ZedDotDev(ZedDotDevModel::Gpt4)
| LanguageModel::ZedDotDev(ZedDotDevModel::Gpt4Turbo)
| LanguageModel::ZedDotDev(ZedDotDevModel::Gpt4Omni)
| LanguageModel::ZedDotDev(ZedDotDevModel::Gpt3Point5Turbo) => {
LanguageModel::Cloud(CloudModel::Gpt4)
| LanguageModel::Cloud(CloudModel::Gpt4Turbo)
| LanguageModel::Cloud(CloudModel::Gpt4Omni)
| LanguageModel::Cloud(CloudModel::Gpt3Point5Turbo) => {
count_open_ai_tokens(request, cx.background_executor())
}
LanguageModel::ZedDotDev(
ZedDotDevModel::Claude3Opus
| ZedDotDevModel::Claude3Sonnet
| ZedDotDevModel::Claude3Haiku,
LanguageModel::Cloud(
CloudModel::Claude3Opus | CloudModel::Claude3Sonnet | CloudModel::Claude3Haiku,
) => {
// Can't find a tokenizer for Claude 3, so for now just use the same as OpenAI's as an approximation.
count_open_ai_tokens(request, cx.background_executor())
}
LanguageModel::ZedDotDev(ZedDotDevModel::Custom(model)) => {
LanguageModel::Cloud(CloudModel::Custom(model)) => {
let request = self.client.request(proto::CountTokensWithLanguageModel {
model,
messages: request
@@ -129,8 +127,10 @@ impl ZedDotDevCompletionProvider {
pub fn complete(
&self,
request: LanguageModelRequest,
mut request: LanguageModelRequest,
) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
request.preprocess();
let request = proto::CompleteWithLanguageModel {
model: request.model.id().to_string(),
messages: request
crates/assistant/src/completion_provider/open_ai.rs
+4 -4
View File
@@ -1,4 +1,4 @@
use crate::assistant_settings::ZedDotDevModel;
use crate::assistant_settings::CloudModel;
use crate::{
assistant_settings::OpenAiModel, CompletionProvider, LanguageModel, LanguageModelRequest, Role,
};
@@ -210,9 +210,9 @@ pub fn count_open_ai_tokens(
match request.model {
LanguageModel::Anthropic(_)
| LanguageModel::ZedDotDev(ZedDotDevModel::Claude3Opus)
| LanguageModel::ZedDotDev(ZedDotDevModel::Claude3Sonnet)
| LanguageModel::ZedDotDev(ZedDotDevModel::Claude3Haiku) => {
| LanguageModel::Cloud(CloudModel::Claude3Opus)
| LanguageModel::Cloud(CloudModel::Claude3Sonnet)
| LanguageModel::Cloud(CloudModel::Claude3Haiku) => {
// Tiktoken doesn't yet support these models, so we manually use the
// same tokenizer as GPT-4.
tiktoken_rs::num_tokens_from_messages("gpt-4", &messages)