language_models: Fix beta_headers for Anthropic custom models (#37306)

Umesh Yadav created 2 months ago

Closes #37289

The current implementation has a problem. The **`from_id` method** in
the Anthropic crate works well for predefined models, but not for custom
models that are defined in the settings. This is because it fallbacks to
using default beta headers, which are incorrect for custom models.

The issue is that the model instance for custom models lives within the
`language_models` provider, so I've updated the **`stream_completion`**
method to explicitly accept beta headers from its caller. Now, the beta
headers are passed from the `language_models` provider all the way to
`anthropic.stream_completion`, which resolves the issue.

Release Notes:

- Fixed a bug where extra_beta_headers defined in settings for Anthropic
custom models were being ignored.

---------

Signed-off-by: Umesh Yadav <[e-mail]>

Change summary

crates/anthropic/src/anthropic.rs                | 12 +++++-------
crates/language_models/src/provider/anthropic.rs | 11 +++++++++--
2 files changed, 14 insertions(+), 9 deletions(-)

Detailed changes

crates/anthropic/src/anthropic.rs

@@ -363,11 +363,9 @@ pub async fn complete(
     api_url: &str,
     api_key: &str,
     request: Request,
+    beta_headers: String,
 ) -> Result<Response, AnthropicError> {
     let uri = format!("{api_url}/v1/messages");
-    let beta_headers = Model::from_id(&request.model)
-        .map(|model| model.beta_headers())
-        .unwrap_or_else(|_| Model::DEFAULT_BETA_HEADERS.join(","));
     let request_builder = HttpRequest::builder()
         .method(Method::POST)
         .uri(uri)
@@ -409,8 +407,9 @@ pub async fn stream_completion(
     api_url: &str,
     api_key: &str,
     request: Request,
+    beta_headers: String,
 ) -> Result<BoxStream<'static, Result<Event, AnthropicError>>, AnthropicError> {
-    stream_completion_with_rate_limit_info(client, api_url, api_key, request)
+    stream_completion_with_rate_limit_info(client, api_url, api_key, request, beta_headers)
         .await
         .map(|output| output.0)
 }
@@ -506,6 +505,7 @@ pub async fn stream_completion_with_rate_limit_info(
     api_url: &str,
     api_key: &str,
     request: Request,
+    beta_headers: String,
 ) -> Result<
     (
         BoxStream<'static, Result<Event, AnthropicError>>,
@@ -518,9 +518,7 @@ pub async fn stream_completion_with_rate_limit_info(
         stream: true,
     };
     let uri = format!("{api_url}/v1/messages");
-    let beta_headers = Model::from_id(&request.base.model)
-        .map(|model| model.beta_headers())
-        .unwrap_or_else(|_| Model::DEFAULT_BETA_HEADERS.join(","));
+
     let request_builder = HttpRequest::builder()
         .method(Method::POST)
         .uri(uri)

crates/language_models/src/provider/anthropic.rs

@@ -424,14 +424,21 @@ impl AnthropicModel {
             return futures::future::ready(Err(anyhow!("App state dropped").into())).boxed();
         };
 
+        let beta_headers = self.model.beta_headers();
+
         async move {
             let Some(api_key) = api_key else {
                 return Err(LanguageModelCompletionError::NoApiKey {
                     provider: PROVIDER_NAME,
                 });
             };
-            let request =
-                anthropic::stream_completion(http_client.as_ref(), &api_url, &api_key, request);
+            let request = anthropic::stream_completion(
+                http_client.as_ref(),
+                &api_url,
+                &api_key,
+                request,
+                beta_headers,
+            );
             request.await.map_err(Into::into)
         }
         .boxed()