Surface upstream rate limits from Anthropic (#16118)

This PR makes it so hitting upstream rate limits from Anthropic result in an HTTP 429 response instead of an HTTP 500. To do this we need to surface structured errors out of the `anthropic` crate. Release Notes: - N/A

Marshall Bowers committed 2 years ago

Commit ebdb7551 parent fbb533b6 files · +145 −48Browse files

Unified · Split

6 files

Cargo.lock +1 −0
crates/anthropic/Cargo.toml +1 −0
crates/anthropic/src/anthropic.rs +102 −38
crates/collab/src/llm.rs +14 −1
crates/language_model/src/provider/anthropic.rs +15 −5
crates/language_model/src/provider/cloud.rs +12 −4

Cargo.lock+1−0

@@ -255,6 +255,7 @@ dependencies = [ · 254 unmodified lines
255	"serde",	255	"serde",
256	"serde_json",	256	"serde_json",
257	"strum",	257	"strum",
		258	"thiserror",
258	"tokio",	259	"tokio",
259	]	260	]
260		261

crates/anthropic/Cargo.toml+1−0

@@ -24,6 +24,7 @@ schemars = { workspace = true, optional = true } · 23 unmodified lines
24	serde.workspace = true	24	serde.workspace = true
25	serde_json.workspace = true	25	serde_json.workspace = true
26	strum.workspace = true	26	strum.workspace = true
		27	thiserror.workspace = true
27		28
28	[dev-dependencies]	29	[dev-dependencies]
29	tokio.workspace = true	30	tokio.workspace = true

crates/anthropic/src/anthropic.rs+102−38

@@ -1,12 +1,14 @@
1	mod supported_countries;	1	mod supported_countries;
2		2
3	use anyhow::{anyhow, Result};	3	use anyhow::{anyhow, Context, Result};
4	use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};	4	use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
5	use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};	5	use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest};
6	use isahc::config::Configurable;	6	use isahc::config::Configurable;
7	use serde::{Deserialize, Serialize};	7	use serde::{Deserialize, Serialize};
		8	use std::str::FromStr;
8	use std::time::Duration;	9	use std::time::Duration;
9	use strum::EnumIter;	10	use strum::{EnumIter, EnumString};
		11	use thiserror::Error;
10		12
11	pub use supported_countries::*;	13	pub use supported_countries::*;
12		14
@@ -96,7 +98,7 @@ pub async fn complete( · 83 unmodified lines
96	api_url: &str,	98	api_url: &str,
97	api_key: &str,	99	api_key: &str,
98	request: Request,	100	request: Request,
99	) -> Result<Response> {	101	) -> Result<Response, AnthropicError> {
100	let uri = format!("{api_url}/v1/messages");	102	let uri = format!("{api_url}/v1/messages");
101	let request_builder = HttpRequest::builder()	103	let request_builder = HttpRequest::builder()
102	.method(Method::POST)	104	.method(Method::POST)
@@ -106,24 +108,40 @@ pub async fn complete( · 3 unmodified lines
106	.header("X-Api-Key", api_key)	108	.header("X-Api-Key", api_key)
107	.header("Content-Type", "application/json");	109	.header("Content-Type", "application/json");
108		110
109	let serialized_request = serde_json::to_string(&request)?;	111	let serialized_request =
110	let request = request_builder.body(AsyncBody::from(serialized_request))?;	112	serde_json::to_string(&request).context("failed to serialize request")?;
		113	let request = request_builder
		114	.body(AsyncBody::from(serialized_request))
		115	.context("failed to construct request body")?;
111		116
112	let mut response = client.send(request).await?;	117	let mut response = client
		118	.send(request)
		119	.await
		120	.context("failed to send request to Anthropic")?;
113	if response.status().is_success() {	121	if response.status().is_success() {
114	let mut body = Vec::new();	122	let mut body = Vec::new();
115	response.body_mut().read_to_end(&mut body).await?;	123	response
116	let response_message: Response = serde_json::from_slice(&body)?;	124	.body_mut()
		125	.read_to_end(&mut body)
		126	.await
		127	.context("failed to read response body")?;
		128	let response_message: Response =
		129	serde_json::from_slice(&body).context("failed to deserialize response body")?;
117	Ok(response_message)	130	Ok(response_message)
118	} else {	131	} else {
119	let mut body = Vec::new();	132	let mut body = Vec::new();
120	response.body_mut().read_to_end(&mut body).await?;	133	response
121	let body_str = std::str::from_utf8(&body)?;	134	.body_mut()
122	Err(anyhow!(	135	.read_to_end(&mut body)
		136	.await
		137	.context("failed to read response body")?;
		138	let body_str =
		139	std::str::from_utf8(&body).context("failed to parse response body as UTF-8")?;
		140	Err(AnthropicError::Other(anyhow!(
123	"Failed to connect to API: {} {}",	141	"Failed to connect to API: {} {}",
124	response.status(),	142	response.status(),
125	body_str	143	body_str
126	))	144	)))
127	}	145	}
128	}	146	}
129		147
@@ -133,7 +151,7 @@ pub async fn stream_completion( · 3 unmodified lines
133	api_key: &str,	151	api_key: &str,
134	request: Request,	152	request: Request,
135	low_speed_timeout: Option<Duration>,	153	low_speed_timeout: Option<Duration>,
136	) -> Result<BoxStream<'static, Result<Event>>> {	154	) -> Result<BoxStream<'static, Result<Event, AnthropicError>>, AnthropicError> {
137	let request = StreamingRequest {	155	let request = StreamingRequest {
138	base: request,	156	base: request,
139	stream: true,	157	stream: true,
@@ -149,10 +167,16 @@ pub async fn stream_completion( · 9 unmodified lines
149	if let Some(low_speed_timeout) = low_speed_timeout {	167	if let Some(low_speed_timeout) = low_speed_timeout {
150	request_builder = request_builder.low_speed_timeout(100, low_speed_timeout);	168	request_builder = request_builder.low_speed_timeout(100, low_speed_timeout);
151	}	169	}
152	let serialized_request = serde_json::to_string(&request)?;	170	let serialized_request =
153	let request = request_builder.body(AsyncBody::from(serialized_request))?;	171	serde_json::to_string(&request).context("failed to serialize request")?;
154		172	let request = request_builder
155	let mut response = client.send(request).await?;	173	.body(AsyncBody::from(serialized_request))
		174	.context("failed to construct request body")?;
		175
		176	let mut response = client
		177	.send(request)
		178	.await
		179	.context("failed to send request to Anthropic")?;
156	if response.status().is_success() {	180	if response.status().is_success() {
157	let reader = BufReader::new(response.into_body());	181	let reader = BufReader::new(response.into_body());
158	Ok(reader	182	Ok(reader
@@ -163,36 +187,41 @@ pub async fn stream_completion( · 4 unmodified lines
163	let line = line.strip_prefix("data: ")?;	187	let line = line.strip_prefix("data: ")?;
164	match serde_json::from_str(line) {	188	match serde_json::from_str(line) {
165	Ok(response) => Some(Ok(response)),	189	Ok(response) => Some(Ok(response)),
166	Err(error) => Some(Err(anyhow!(error))),	190	Err(error) => Some(Err(AnthropicError::Other(anyhow!(error)))),
167	}	191	}
168	}	192	}
169	Err(error) => Some(Err(anyhow!(error))),	193	Err(error) => Some(Err(AnthropicError::Other(anyhow!(error)))),
170	}	194	}
171	})	195	})
172	.boxed())	196	.boxed())
173	} else {	197	} else {
174	let mut body = Vec::new();	198	let mut body = Vec::new();
175	response.body_mut().read_to_end(&mut body).await?;	199	response
		200	.body_mut()
		201	.read_to_end(&mut body)
		202	.await
		203	.context("failed to read response body")?;
176		204
177	let body_str = std::str::from_utf8(&body)?;	205	let body_str =
		206	std::str::from_utf8(&body).context("failed to parse response body as UTF-8")?;
178		207
179	match serde_json::from_str::<Event>(body_str) {	208	match serde_json::from_str::<Event>(body_str) {
180	Ok(Event::Error { error }) => Err(api_error_to_err(error)),	209	Ok(Event::Error { error }) => Err(AnthropicError::ApiError(error)),
181	Ok(_) => Err(anyhow!(	210	Ok(_) => Err(AnthropicError::Other(anyhow!(
182	"Unexpected success response while expecting an error: '{body_str}'",	211	"Unexpected success response while expecting an error: '{body_str}'",
183	)),	212	))),
184	Err(_) => Err(anyhow!(	213	Err(_) => Err(AnthropicError::Other(anyhow!(
185	"Failed to connect to API: {} {}",	214	"Failed to connect to API: {} {}",
186	response.status(),	215	response.status(),
187	body_str,	216	body_str,
188	)),	217	))),
189	}	218	}
190	}	219	}
191	}	220	}
192		221
193	pub fn extract_text_from_events(	222	pub fn extract_text_from_events(
194	response: impl Stream<Item = Result<Event>>,	223	response: impl Stream<Item = Result<Event, AnthropicError>>,
195	) -> impl Stream<Item = Result<String>> {	224	) -> impl Stream<Item = Result<String, AnthropicError>> {
196	response.filter_map(\|response\| async move {	225	response.filter_map(\|response\| async move {
197	match response {	226	match response {
198	Ok(response) => match response {	227	Ok(response) => match response {
@@ -204,7 +233,7 @@ pub fn extract_text_from_events( · 5 unmodified lines
204	ContentDelta::TextDelta { text } => Some(Ok(text)),	233	ContentDelta::TextDelta { text } => Some(Ok(text)),
205	_ => None,	234	_ => None,
206	},	235	},
207	Event::Error { error } => Some(Err(api_error_to_err(error))),	236	Event::Error { error } => Some(Err(AnthropicError::ApiError(error))),
208	_ => None,	237	_ => None,
209	},	238	},
210	Err(error) => Some(Err(error)),	239	Err(error) => Some(Err(error)),
@@ -212,15 +241,6 @@ pub fn extract_text_from_events( · 1 unmodified line
212	})	241	})
213	}	242	}
214		243
215	fn api_error_to_err(
216	ApiError {
217	error_type,
218	message,
219	}: ApiError,
220	) -> anyhow::Error {
221	anyhow!("API error. Type: '{error_type}', message: '{message}'",)
222	}
223
224	#[derive(Debug, Serialize, Deserialize)]	244	#[derive(Debug, Serialize, Deserialize)]
225	pub struct Message {	245	pub struct Message {
226	pub role: Role,	246	pub role: Role,
@@ -374,9 +394,53 @@ pub struct MessageDelta { · 147 unmodified lines
374	pub stop_sequence: Option<String>,	394	pub stop_sequence: Option<String>,
375	}	395	}
376		396
		397	#[derive(Error, Debug)]
		398	pub enum AnthropicError {
		399	#[error("an error occurred while interacting with the Anthropic API: {error_type}: {message}", error_type = .0.error_type, message = .0.message)]
		400	ApiError(ApiError),
		401	#[error("{0}")]
		402	Other(#[from] anyhow::Error),
		403	}
		404
377	#[derive(Debug, Serialize, Deserialize)]	405	#[derive(Debug, Serialize, Deserialize)]
378	pub struct ApiError {	406	pub struct ApiError {
379	#[serde(rename = "type")]	407	#[serde(rename = "type")]
380	pub error_type: String,	408	pub error_type: String,
381	pub message: String,	409	pub message: String,
382	}	410	}
		411
		412	/// An Anthropic API error code.
		413	/// https://docs.anthropic.com/en/api/errors#http-errors
		414	#[derive(Debug, PartialEq, Eq, Clone, Copy, EnumString)]
		415	#[strum(serialize_all = "snake_case")]
		416	pub enum ApiErrorCode {
		417	/// 400 - `invalid_request_error`: There was an issue with the format or content of your request.
		418	InvalidRequestError,
		419	/// 401 - `authentication_error`: There's an issue with your API key.
		420	AuthenticationError,
		421	/// 403 - `permission_error`: Your API key does not have permission to use the specified resource.
		422	PermissionError,
		423	/// 404 - `not_found_error`: The requested resource was not found.
		424	NotFoundError,
		425	/// 413 - `request_too_large`: Request exceeds the maximum allowed number of bytes.
		426	RequestTooLarge,
		427	/// 429 - `rate_limit_error`: Your account has hit a rate limit.
		428	RateLimitError,
		429	/// 500 - `api_error`: An unexpected error has occurred internal to Anthropic's systems.
		430	ApiError,
		431	/// 529 - `overloaded_error`: Anthropic's API is temporarily overloaded.
		432	OverloadedError,
		433	}
		434
		435	impl ApiError {
		436	pub fn code(&self) -> Option<ApiErrorCode> {
		437	ApiErrorCode::from_str(&self.error_type).ok()
		438	}
		439
		440	pub fn is_rate_limit_error(&self) -> bool {
		441	match self.error_type.as_str() {
		442	"rate_limit_error" => true,
		443	_ => false,
		444	}
		445	}
		446	}

crates/collab/src/llm.rs+14−1

@@ -197,7 +197,20 @@ async fn perform_completion( · 196 unmodified lines
197	request,	197	request,
198	None,	198	None,
199	)	199	)
200	.await?;	200	.await
		201	.map_err(\|err\| match err {
		202	anthropic::AnthropicError::ApiError(ref api_error) => {
		203	if api_error.code() == Some(anthropic::ApiErrorCode::RateLimitError) {
		204	return Error::http(
		205	StatusCode::TOO_MANY_REQUESTS,
		206	"Upstream Anthropic rate limit exceeded.".to_string(),
		207	);
		208	}
		209
		210	Error::Internal(anyhow!(err))
		211	}
		212	anthropic::AnthropicError::Other(err) => Error::Internal(err),
		213	})?;
201		214
202	chunks	215	chunks
203	.map(move \|event\| {	216	.map(move \|event\| {

crates/language_model/src/provider/anthropic.rs+15−5

     LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
     LanguageModelProviderState, LanguageModelRequest, RateLimiter, Role,
 };
+use anthropic::AnthropicError;
 use anyhow::{anyhow, Context as _, Result};
 use collections::BTreeMap;
 use editor::{Editor, EditorElement, EditorStyle};
         async move {
             let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
-            anthropic::complete(http_client.as_ref(), &api_url, &api_key, request).await
+            anthropic::complete(http_client.as_ref(), &api_url, &api_key, request)
+                .await
+                .context("failed to retrieve completion")
         }
         .boxed()
     }
         &self,
         request: anthropic::Request,
         cx: &AsyncAppContext,
-    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<anthropic::Event>>>> {
+    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<anthropic::Event, AnthropicError>>>>
+    {
         let http_client = self.http_client.clone();
         let Ok((api_key, api_url, low_speed_timeout)) = cx.read_model(&self.state, |state, cx| {
                 request,
                 low_speed_timeout,
             );
-            request.await
+            request.await.context("failed to stream completion")
         }
         .boxed()
     }
         let request = request.into_anthropic(self.model.id().into());
         let request = self.stream_completion(request, cx);
         let future = self.request_limiter.stream(async move {
-            let response = request.await?;
+            let response = request.await.map_err(|err| anyhow!(err))?;
             Ok(anthropic::extract_text_from_events(response))
         });
-        async move { Ok(future.await?.boxed()) }.boxed()
+        async move {
+            Ok(future
+                .await?
+                .map(|result| result.map_err(|err| anyhow!(err)))
+                .boxed())
+        }
+        .boxed()
     }
     fn use_any_tool(

crates/language_model/src/provider/cloud.rs+12−4

     LanguageModelName, LanguageModelProviderId, LanguageModelProviderName,
     LanguageModelProviderState, LanguageModelRequest, RateLimiter, ZedModel,
 };
-use anyhow::{anyhow, bail, Result};
+use anthropic::AnthropicError;
+use anyhow::{anyhow, bail, Context as _, Result};
 use client::{Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME};
 use collections::BTreeMap;
 use feature_flags::{FeatureFlagAppExt, LanguageModels};
                         match body.read_line(&mut buffer).await {
                             Ok(0) => Ok(None),
                             Ok(_) => {
-                                let event: anthropic::Event = serde_json::from_str(&buffer)?;
+                                let event: anthropic::Event = serde_json::from_str(&buffer)
+                                    .context("failed to parse Anthropic event")?;
                                 Ok(Some((event, body)))
                             }
-                            Err(e) => Err(e.into()),
+                            Err(err) => Err(AnthropicError::Other(err.into())),
                         }
                     });
                     Ok(anthropic::extract_text_from_events(stream))
                 });
-                async move { Ok(future.await?.boxed()) }.boxed()
+                async move {
+                    Ok(future
+                        .await?
+                        .map(|result| result.map_err(|err| anyhow!(err)))
+                        .boxed())
+                }
+                .boxed()
             }
             CloudModel::OpenAi(model) => {
                 let client = self.client.clone();