1use std::pin::Pin;
2use std::str::FromStr;
3use std::sync::Arc;
4
5use anyhow::{Context as _, Result, anyhow};
6use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;
7use aws_config::{BehaviorVersion, Region};
8use aws_credential_types::Credentials;
9use aws_http_client::AwsHttpClient;
10use bedrock::bedrock_client::Client as BedrockClient;
11use bedrock::bedrock_client::config::timeout::TimeoutConfig;
12use bedrock::bedrock_client::types::{
13 CachePointBlock, CachePointType, ContentBlockDelta, ContentBlockStart, ConverseStreamOutput,
14 ReasoningContentBlockDelta, StopReason,
15};
16use bedrock::{
17 BedrockAnyToolChoice, BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockInnerContent,
18 BedrockMessage, BedrockModelMode, BedrockStreamingResponse, BedrockThinkingBlock,
19 BedrockThinkingTextBlock, BedrockTool, BedrockToolChoice, BedrockToolConfig,
20 BedrockToolInputSchema, BedrockToolResultBlock, BedrockToolResultContentBlock,
21 BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
22};
23use collections::{BTreeMap, HashMap};
24use credentials_provider::CredentialsProvider;
25use futures::{FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream};
26use gpui::{
27 AnyView, App, AsyncApp, Context, Entity, FocusHandle, FontWeight, Subscription, Task, Window,
28 actions,
29};
30use gpui_tokio::Tokio;
31use http_client::HttpClient;
32use language_model::{
33 AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
34 LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
35 LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
36 LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
37 LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, RateLimiter, Role,
38 TokenUsage,
39};
40use schemars::JsonSchema;
41use serde::{Deserialize, Serialize};
42use serde_json::Value;
43use settings::{BedrockAvailableModel as AvailableModel, Settings, SettingsStore};
44use smol::lock::OnceCell;
45use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
46use ui::{ButtonLink, ConfiguredApiCard, List, ListBulletItem, prelude::*};
47use ui_input::InputField;
48use util::ResultExt;
49
50use crate::AllLanguageModelSettings;
51
52actions!(bedrock, [Tab, TabPrev]);
53
54const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("amazon-bedrock");
55const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Amazon Bedrock");
56
57#[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]
58pub struct BedrockCredentials {
59 pub access_key_id: String,
60 pub secret_access_key: String,
61 pub session_token: Option<String>,
62 pub region: String,
63}
64
65#[derive(Default, Clone, Debug, PartialEq)]
66pub struct AmazonBedrockSettings {
67 pub available_models: Vec<AvailableModel>,
68 pub region: Option<String>,
69 pub endpoint: Option<String>,
70 pub profile_name: Option<String>,
71 pub role_arn: Option<String>,
72 pub authentication_method: Option<BedrockAuthMethod>,
73 pub allow_global: Option<bool>,
74}
75
76#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, EnumIter, IntoStaticStr, JsonSchema)]
77pub enum BedrockAuthMethod {
78 #[serde(rename = "named_profile")]
79 NamedProfile,
80 #[serde(rename = "sso")]
81 SingleSignOn,
82 /// IMDSv2, PodIdentity, env vars, etc.
83 #[serde(rename = "default")]
84 Automatic,
85}
86
87impl From<settings::BedrockAuthMethodContent> for BedrockAuthMethod {
88 fn from(value: settings::BedrockAuthMethodContent) -> Self {
89 match value {
90 settings::BedrockAuthMethodContent::SingleSignOn => BedrockAuthMethod::SingleSignOn,
91 settings::BedrockAuthMethodContent::Automatic => BedrockAuthMethod::Automatic,
92 settings::BedrockAuthMethodContent::NamedProfile => BedrockAuthMethod::NamedProfile,
93 }
94 }
95}
96
97#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
98#[serde(tag = "type", rename_all = "lowercase")]
99pub enum ModelMode {
100 #[default]
101 Default,
102 Thinking {
103 /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
104 budget_tokens: Option<u64>,
105 },
106}
107
108impl From<ModelMode> for BedrockModelMode {
109 fn from(value: ModelMode) -> Self {
110 match value {
111 ModelMode::Default => BedrockModelMode::Default,
112 ModelMode::Thinking { budget_tokens } => BedrockModelMode::Thinking { budget_tokens },
113 }
114 }
115}
116
117impl From<BedrockModelMode> for ModelMode {
118 fn from(value: BedrockModelMode) -> Self {
119 match value {
120 BedrockModelMode::Default => ModelMode::Default,
121 BedrockModelMode::Thinking { budget_tokens } => ModelMode::Thinking { budget_tokens },
122 }
123 }
124}
125
126/// The URL of the base AWS service.
127///
128/// Right now we're just using this as the key to store the AWS credentials
129/// under in the keychain.
130const AMAZON_AWS_URL: &str = "https://amazonaws.com";
131
132// These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials.
133const ZED_BEDROCK_ACCESS_KEY_ID_VAR: &str = "ZED_ACCESS_KEY_ID";
134const ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: &str = "ZED_SECRET_ACCESS_KEY";
135const ZED_BEDROCK_SESSION_TOKEN_VAR: &str = "ZED_SESSION_TOKEN";
136const ZED_AWS_PROFILE_VAR: &str = "ZED_AWS_PROFILE";
137const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION";
138const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS";
139const ZED_AWS_ENDPOINT_VAR: &str = "ZED_AWS_ENDPOINT";
140
141pub struct State {
142 credentials: Option<BedrockCredentials>,
143 settings: Option<AmazonBedrockSettings>,
144 credentials_from_env: bool,
145 _subscription: Subscription,
146}
147
148impl State {
149 fn reset_credentials(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
150 let credentials_provider = <dyn CredentialsProvider>::global(cx);
151 cx.spawn(async move |this, cx| {
152 credentials_provider
153 .delete_credentials(AMAZON_AWS_URL, cx)
154 .await
155 .log_err();
156 this.update(cx, |this, cx| {
157 this.credentials = None;
158 this.credentials_from_env = false;
159 this.settings = None;
160 cx.notify();
161 })
162 })
163 }
164
165 fn set_credentials(
166 &mut self,
167 credentials: BedrockCredentials,
168 cx: &mut Context<Self>,
169 ) -> Task<Result<()>> {
170 let credentials_provider = <dyn CredentialsProvider>::global(cx);
171 cx.spawn(async move |this, cx| {
172 credentials_provider
173 .write_credentials(
174 AMAZON_AWS_URL,
175 "Bearer",
176 &serde_json::to_vec(&credentials)?,
177 cx,
178 )
179 .await?;
180 this.update(cx, |this, cx| {
181 this.credentials = Some(credentials);
182 cx.notify();
183 })
184 })
185 }
186
187 fn is_authenticated(&self) -> bool {
188 let derived = self
189 .settings
190 .as_ref()
191 .and_then(|s| s.authentication_method.as_ref());
192 let creds = self.credentials.as_ref();
193
194 derived.is_some() || creds.is_some()
195 }
196
197 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
198 if self.is_authenticated() {
199 return Task::ready(Ok(()));
200 }
201
202 let credentials_provider = <dyn CredentialsProvider>::global(cx);
203 cx.spawn(async move |this, cx| {
204 let (credentials, from_env) =
205 if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {
206 (credentials, true)
207 } else {
208 let (_, credentials) = credentials_provider
209 .read_credentials(AMAZON_AWS_URL, cx)
210 .await?
211 .ok_or_else(|| AuthenticateError::CredentialsNotFound)?;
212 (
213 String::from_utf8(credentials)
214 .context("invalid {PROVIDER_NAME} credentials")?,
215 false,
216 )
217 };
218
219 let credentials: BedrockCredentials =
220 serde_json::from_str(&credentials).context("failed to parse credentials")?;
221
222 this.update(cx, |this, cx| {
223 this.credentials = Some(credentials);
224 this.credentials_from_env = from_env;
225 cx.notify();
226 })?;
227
228 Ok(())
229 })
230 }
231
232 fn get_region(&self) -> String {
233 // Get region - from credentials or directly from settings
234 let credentials_region = self.credentials.as_ref().map(|s| s.region.clone());
235 let settings_region = self.settings.as_ref().and_then(|s| s.region.clone());
236
237 // Use credentials region if available, otherwise use settings region, finally fall back to default
238 credentials_region
239 .or(settings_region)
240 .unwrap_or(String::from("us-east-1"))
241 }
242
243 fn get_allow_global(&self) -> bool {
244 self.settings
245 .as_ref()
246 .and_then(|s| s.allow_global)
247 .unwrap_or(false)
248 }
249}
250
251pub struct BedrockLanguageModelProvider {
252 http_client: AwsHttpClient,
253 handle: tokio::runtime::Handle,
254 state: Entity<State>,
255}
256
257impl BedrockLanguageModelProvider {
258 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
259 let state = cx.new(|cx| State {
260 credentials: None,
261 settings: Some(AllLanguageModelSettings::get_global(cx).bedrock.clone()),
262 credentials_from_env: false,
263 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
264 cx.notify();
265 }),
266 });
267
268 Self {
269 http_client: AwsHttpClient::new(http_client.clone()),
270 handle: Tokio::handle(cx),
271 state,
272 }
273 }
274
275 fn create_language_model(&self, model: bedrock::Model) -> Arc<dyn LanguageModel> {
276 Arc::new(BedrockModel {
277 id: LanguageModelId::from(model.id().to_string()),
278 model,
279 http_client: self.http_client.clone(),
280 handle: self.handle.clone(),
281 state: self.state.clone(),
282 client: OnceCell::new(),
283 request_limiter: RateLimiter::new(4),
284 })
285 }
286}
287
288impl LanguageModelProvider for BedrockLanguageModelProvider {
289 fn id(&self) -> LanguageModelProviderId {
290 PROVIDER_ID
291 }
292
293 fn name(&self) -> LanguageModelProviderName {
294 PROVIDER_NAME
295 }
296
297 fn icon(&self) -> IconName {
298 IconName::AiBedrock
299 }
300
301 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
302 Some(self.create_language_model(bedrock::Model::default()))
303 }
304
305 fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
306 let region = self.state.read(cx).get_region();
307 Some(self.create_language_model(bedrock::Model::default_fast(region.as_str())))
308 }
309
310 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
311 let mut models = BTreeMap::default();
312
313 for model in bedrock::Model::iter() {
314 if !matches!(model, bedrock::Model::Custom { .. }) {
315 // TODO: Sonnet 3.7 vs. 3.7 Thinking bug is here.
316 models.insert(model.id().to_string(), model);
317 }
318 }
319
320 // Override with available models from settings
321 for model in AllLanguageModelSettings::get_global(cx)
322 .bedrock
323 .available_models
324 .iter()
325 {
326 models.insert(
327 model.name.clone(),
328 bedrock::Model::Custom {
329 name: model.name.clone(),
330 display_name: model.display_name.clone(),
331 max_tokens: model.max_tokens,
332 max_output_tokens: model.max_output_tokens,
333 default_temperature: model.default_temperature,
334 cache_configuration: model.cache_configuration.as_ref().map(|config| {
335 bedrock::BedrockModelCacheConfiguration {
336 max_cache_anchors: config.max_cache_anchors,
337 min_total_token: config.min_total_token,
338 }
339 }),
340 },
341 );
342 }
343
344 models
345 .into_values()
346 .map(|model| self.create_language_model(model))
347 .collect()
348 }
349
350 fn is_authenticated(&self, cx: &App) -> bool {
351 self.state.read(cx).is_authenticated()
352 }
353
354 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
355 self.state.update(cx, |state, cx| state.authenticate(cx))
356 }
357
358 fn configuration_view(
359 &self,
360 _target_agent: language_model::ConfigurationViewTargetAgent,
361 window: &mut Window,
362 cx: &mut App,
363 ) -> AnyView {
364 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
365 .into()
366 }
367
368 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
369 self.state
370 .update(cx, |state, cx| state.reset_credentials(cx))
371 }
372}
373
374impl LanguageModelProviderState for BedrockLanguageModelProvider {
375 type ObservableEntity = State;
376
377 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
378 Some(self.state.clone())
379 }
380}
381
382struct BedrockModel {
383 id: LanguageModelId,
384 model: Model,
385 http_client: AwsHttpClient,
386 handle: tokio::runtime::Handle,
387 client: OnceCell<BedrockClient>,
388 state: Entity<State>,
389 request_limiter: RateLimiter,
390}
391
392impl BedrockModel {
393 fn get_or_init_client(&self, cx: &AsyncApp) -> anyhow::Result<&BedrockClient> {
394 self.client
395 .get_or_try_init_blocking(|| {
396 let (auth_method, credentials, endpoint, region, settings) =
397 cx.read_entity(&self.state, |state, _cx| {
398 let auth_method = state
399 .settings
400 .as_ref()
401 .and_then(|s| s.authentication_method.clone());
402
403 let endpoint = state.settings.as_ref().and_then(|s| s.endpoint.clone());
404
405 let region = state.get_region();
406
407 (
408 auth_method,
409 state.credentials.clone(),
410 endpoint,
411 region,
412 state.settings.clone(),
413 )
414 })?;
415
416 let mut config_builder = aws_config::defaults(BehaviorVersion::latest())
417 .stalled_stream_protection(StalledStreamProtectionConfig::disabled())
418 .http_client(self.http_client.clone())
419 .region(Region::new(region))
420 .timeout_config(TimeoutConfig::disabled());
421
422 if let Some(endpoint_url) = endpoint
423 && !endpoint_url.is_empty()
424 {
425 config_builder = config_builder.endpoint_url(endpoint_url);
426 }
427
428 match auth_method {
429 None => {
430 if let Some(creds) = credentials {
431 let aws_creds = Credentials::new(
432 creds.access_key_id,
433 creds.secret_access_key,
434 creds.session_token,
435 None,
436 "zed-bedrock-provider",
437 );
438 config_builder = config_builder.credentials_provider(aws_creds);
439 }
440 }
441 Some(BedrockAuthMethod::NamedProfile)
442 | Some(BedrockAuthMethod::SingleSignOn) => {
443 // Currently NamedProfile and SSO behave the same way but only the instructions change
444 // Until we support BearerAuth through SSO, this will not change.
445 let profile_name = settings
446 .and_then(|s| s.profile_name)
447 .unwrap_or_else(|| "default".to_string());
448
449 if !profile_name.is_empty() {
450 config_builder = config_builder.profile_name(profile_name);
451 }
452 }
453 Some(BedrockAuthMethod::Automatic) => {
454 // Use default credential provider chain
455 }
456 }
457
458 let config = self.handle.block_on(config_builder.load());
459 anyhow::Ok(BedrockClient::new(&config))
460 })
461 .context("initializing Bedrock client")?;
462
463 self.client.get().context("Bedrock client not initialized")
464 }
465
466 fn stream_completion(
467 &self,
468 request: bedrock::Request,
469 cx: &AsyncApp,
470 ) -> BoxFuture<
471 'static,
472 Result<BoxStream<'static, Result<BedrockStreamingResponse, BedrockError>>>,
473 > {
474 let Ok(runtime_client) = self
475 .get_or_init_client(cx)
476 .cloned()
477 .context("Bedrock client not initialized")
478 else {
479 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
480 };
481
482 match Tokio::spawn(cx, bedrock::stream_completion(runtime_client, request)) {
483 Ok(res) => async { res.await.map_err(|err| anyhow!(err))? }.boxed(),
484 Err(err) => futures::future::ready(Err(anyhow!(err))).boxed(),
485 }
486 }
487}
488
489impl LanguageModel for BedrockModel {
490 fn id(&self) -> LanguageModelId {
491 self.id.clone()
492 }
493
494 fn name(&self) -> LanguageModelName {
495 LanguageModelName::from(self.model.display_name().to_string())
496 }
497
498 fn provider_id(&self) -> LanguageModelProviderId {
499 PROVIDER_ID
500 }
501
502 fn provider_name(&self) -> LanguageModelProviderName {
503 PROVIDER_NAME
504 }
505
506 fn supports_tools(&self) -> bool {
507 self.model.supports_tool_use()
508 }
509
510 fn supports_images(&self) -> bool {
511 false
512 }
513
514 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
515 match choice {
516 LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
517 self.model.supports_tool_use()
518 }
519 // Add support for None - we'll filter tool calls at response
520 LanguageModelToolChoice::None => self.model.supports_tool_use(),
521 }
522 }
523
524 fn telemetry_id(&self) -> String {
525 format!("bedrock/{}", self.model.id())
526 }
527
528 fn max_token_count(&self) -> u64 {
529 self.model.max_token_count()
530 }
531
532 fn max_output_tokens(&self) -> Option<u64> {
533 Some(self.model.max_output_tokens())
534 }
535
536 fn count_tokens(
537 &self,
538 request: LanguageModelRequest,
539 cx: &App,
540 ) -> BoxFuture<'static, Result<u64>> {
541 get_bedrock_tokens(request, cx)
542 }
543
544 fn stream_completion(
545 &self,
546 request: LanguageModelRequest,
547 cx: &AsyncApp,
548 ) -> BoxFuture<
549 'static,
550 Result<
551 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
552 LanguageModelCompletionError,
553 >,
554 > {
555 let Ok((region, allow_global)) = cx.read_entity(&self.state, |state, _cx| {
556 (state.get_region(), state.get_allow_global())
557 }) else {
558 return async move { Err(anyhow::anyhow!("App State Dropped").into()) }.boxed();
559 };
560
561 let model_id = match self.model.cross_region_inference_id(®ion, allow_global) {
562 Ok(s) => s,
563 Err(e) => {
564 return async move { Err(e.into()) }.boxed();
565 }
566 };
567
568 let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None);
569
570 let request = match into_bedrock(
571 request,
572 model_id,
573 self.model.default_temperature(),
574 self.model.max_output_tokens(),
575 self.model.mode(),
576 self.model.supports_caching(),
577 ) {
578 Ok(request) => request,
579 Err(err) => return futures::future::ready(Err(err.into())).boxed(),
580 };
581
582 let request = self.stream_completion(request, cx);
583 let future = self.request_limiter.stream(async move {
584 let response = request.await.map_err(|err| anyhow!(err))?;
585 let events = map_to_language_model_completion_events(response);
586
587 if deny_tool_calls {
588 Ok(deny_tool_use_events(events).boxed())
589 } else {
590 Ok(events.boxed())
591 }
592 });
593
594 async move { Ok(future.await?.boxed()) }.boxed()
595 }
596
597 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
598 self.model
599 .cache_configuration()
600 .map(|config| LanguageModelCacheConfiguration {
601 max_cache_anchors: config.max_cache_anchors,
602 should_speculate: false,
603 min_total_token: config.min_total_token,
604 })
605 }
606}
607
608fn deny_tool_use_events(
609 events: impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
610) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
611 events.map(|event| {
612 match event {
613 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
614 // Convert tool use to an error message if model decided to call it
615 Ok(LanguageModelCompletionEvent::Text(format!(
616 "\n\n[Error: Tool calls are disabled in this context. Attempted to call '{}']",
617 tool_use.name
618 )))
619 }
620 other => other,
621 }
622 })
623}
624
625pub fn into_bedrock(
626 request: LanguageModelRequest,
627 model: String,
628 default_temperature: f32,
629 max_output_tokens: u64,
630 mode: BedrockModelMode,
631 supports_caching: bool,
632) -> Result<bedrock::Request> {
633 let mut new_messages: Vec<BedrockMessage> = Vec::new();
634 let mut system_message = String::new();
635
636 for message in request.messages {
637 if message.contents_empty() {
638 continue;
639 }
640
641 match message.role {
642 Role::User | Role::Assistant => {
643 let mut bedrock_message_content: Vec<BedrockInnerContent> = message
644 .content
645 .into_iter()
646 .filter_map(|content| match content {
647 MessageContent::Text(text) => {
648 if !text.is_empty() {
649 Some(BedrockInnerContent::Text(text))
650 } else {
651 None
652 }
653 }
654 MessageContent::Thinking { text, signature } => {
655 if model.contains(Model::DeepSeekR1.request_id()) {
656 // DeepSeekR1 doesn't support thinking blocks
657 // And the AWS API demands that you strip them
658 return None;
659 }
660 let thinking = BedrockThinkingTextBlock::builder()
661 .text(text)
662 .set_signature(signature)
663 .build()
664 .context("failed to build reasoning block")
665 .log_err()?;
666
667 Some(BedrockInnerContent::ReasoningContent(
668 BedrockThinkingBlock::ReasoningText(thinking),
669 ))
670 }
671 MessageContent::RedactedThinking(blob) => {
672 if model.contains(Model::DeepSeekR1.request_id()) {
673 // DeepSeekR1 doesn't support thinking blocks
674 // And the AWS API demands that you strip them
675 return None;
676 }
677 let redacted =
678 BedrockThinkingBlock::RedactedContent(BedrockBlob::new(blob));
679
680 Some(BedrockInnerContent::ReasoningContent(redacted))
681 }
682 MessageContent::ToolUse(tool_use) => {
683 let input = if tool_use.input.is_null() {
684 // Bedrock API requires valid JsonValue, not null, for tool use input
685 value_to_aws_document(&serde_json::json!({}))
686 } else {
687 value_to_aws_document(&tool_use.input)
688 };
689 BedrockToolUseBlock::builder()
690 .name(tool_use.name.to_string())
691 .tool_use_id(tool_use.id.to_string())
692 .input(input)
693 .build()
694 .context("failed to build Bedrock tool use block")
695 .log_err()
696 .map(BedrockInnerContent::ToolUse)
697 },
698 MessageContent::ToolResult(tool_result) => {
699 BedrockToolResultBlock::builder()
700 .tool_use_id(tool_result.tool_use_id.to_string())
701 .content(match tool_result.content {
702 LanguageModelToolResultContent::Text(text) => {
703 BedrockToolResultContentBlock::Text(text.to_string())
704 }
705 LanguageModelToolResultContent::Image(_) => {
706 BedrockToolResultContentBlock::Text(
707 // TODO: Bedrock image support
708 "[Tool responded with an image, but Zed doesn't support these in Bedrock models yet]".to_string()
709 )
710 }
711 })
712 .status({
713 if tool_result.is_error {
714 BedrockToolResultStatus::Error
715 } else {
716 BedrockToolResultStatus::Success
717 }
718 })
719 .build()
720 .context("failed to build Bedrock tool result block")
721 .log_err()
722 .map(BedrockInnerContent::ToolResult)
723 }
724 _ => None,
725 })
726 .collect();
727 if message.cache && supports_caching {
728 bedrock_message_content.push(BedrockInnerContent::CachePoint(
729 CachePointBlock::builder()
730 .r#type(CachePointType::Default)
731 .build()
732 .context("failed to build cache point block")?,
733 ));
734 }
735 let bedrock_role = match message.role {
736 Role::User => bedrock::BedrockRole::User,
737 Role::Assistant => bedrock::BedrockRole::Assistant,
738 Role::System => unreachable!("System role should never occur here"),
739 };
740 if let Some(last_message) = new_messages.last_mut()
741 && last_message.role == bedrock_role
742 {
743 last_message.content.extend(bedrock_message_content);
744 continue;
745 }
746 new_messages.push(
747 BedrockMessage::builder()
748 .role(bedrock_role)
749 .set_content(Some(bedrock_message_content))
750 .build()
751 .context("failed to build Bedrock message")?,
752 );
753 }
754 Role::System => {
755 if !system_message.is_empty() {
756 system_message.push_str("\n\n");
757 }
758 system_message.push_str(&message.string_contents());
759 }
760 }
761 }
762
763 let mut tool_spec: Vec<BedrockTool> = request
764 .tools
765 .iter()
766 .filter_map(|tool| {
767 Some(BedrockTool::ToolSpec(
768 BedrockToolSpec::builder()
769 .name(tool.name.clone())
770 .description(tool.description.clone())
771 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
772 &tool.input_schema,
773 )))
774 .build()
775 .log_err()?,
776 ))
777 })
778 .collect();
779
780 if !tool_spec.is_empty() && supports_caching {
781 tool_spec.push(BedrockTool::CachePoint(
782 CachePointBlock::builder()
783 .r#type(CachePointType::Default)
784 .build()
785 .context("failed to build cache point block")?,
786 ));
787 }
788
789 let tool_choice = match request.tool_choice {
790 Some(LanguageModelToolChoice::Auto) | None => {
791 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
792 }
793 Some(LanguageModelToolChoice::Any) => {
794 BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
795 }
796 Some(LanguageModelToolChoice::None) => {
797 // For None, we still use Auto but will filter out tool calls in the response
798 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
799 }
800 };
801 let tool_config: BedrockToolConfig = BedrockToolConfig::builder()
802 .set_tools(Some(tool_spec))
803 .tool_choice(tool_choice)
804 .build()?;
805
806 Ok(bedrock::Request {
807 model,
808 messages: new_messages,
809 max_tokens: max_output_tokens,
810 system: Some(system_message),
811 tools: Some(tool_config),
812 thinking: if request.thinking_allowed
813 && let BedrockModelMode::Thinking { budget_tokens } = mode
814 {
815 Some(bedrock::Thinking::Enabled { budget_tokens })
816 } else {
817 None
818 },
819 metadata: None,
820 stop_sequences: Vec::new(),
821 temperature: request.temperature.or(Some(default_temperature)),
822 top_k: None,
823 top_p: None,
824 })
825}
826
827// TODO: just call the ConverseOutput.usage() method:
828// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
829pub fn get_bedrock_tokens(
830 request: LanguageModelRequest,
831 cx: &App,
832) -> BoxFuture<'static, Result<u64>> {
833 cx.background_executor()
834 .spawn(async move {
835 let messages = request.messages;
836 let mut tokens_from_images = 0;
837 let mut string_messages = Vec::with_capacity(messages.len());
838
839 for message in messages {
840 use language_model::MessageContent;
841
842 let mut string_contents = String::new();
843
844 for content in message.content {
845 match content {
846 MessageContent::Text(text) | MessageContent::Thinking { text, .. } => {
847 string_contents.push_str(&text);
848 }
849 MessageContent::RedactedThinking(_) => {}
850 MessageContent::Image(image) => {
851 tokens_from_images += image.estimate_tokens();
852 }
853 MessageContent::ToolUse(_tool_use) => {
854 // TODO: Estimate token usage from tool uses.
855 }
856 MessageContent::ToolResult(tool_result) => match tool_result.content {
857 LanguageModelToolResultContent::Text(text) => {
858 string_contents.push_str(&text);
859 }
860 LanguageModelToolResultContent::Image(image) => {
861 tokens_from_images += image.estimate_tokens();
862 }
863 },
864 }
865 }
866
867 if !string_contents.is_empty() {
868 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
869 role: match message.role {
870 Role::User => "user".into(),
871 Role::Assistant => "assistant".into(),
872 Role::System => "system".into(),
873 },
874 content: Some(string_contents),
875 name: None,
876 function_call: None,
877 });
878 }
879 }
880
881 // Tiktoken doesn't yet support these models, so we manually use the
882 // same tokenizer as GPT-4.
883 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
884 .map(|tokens| (tokens + tokens_from_images) as u64)
885 })
886 .boxed()
887}
888
889pub fn map_to_language_model_completion_events(
890 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
891) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
892 struct RawToolUse {
893 id: String,
894 name: String,
895 input_json: String,
896 }
897
898 struct State {
899 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
900 tool_uses_by_index: HashMap<i32, RawToolUse>,
901 }
902
903 let initial_state = State {
904 events,
905 tool_uses_by_index: HashMap::default(),
906 };
907
908 futures::stream::unfold(initial_state, |mut state| async move {
909 match state.events.next().await {
910 Some(event_result) => match event_result {
911 Ok(event) => {
912 let result = match event {
913 ConverseStreamOutput::ContentBlockDelta(cb_delta) => match cb_delta.delta {
914 Some(ContentBlockDelta::Text(text)) => {
915 Some(Ok(LanguageModelCompletionEvent::Text(text)))
916 }
917 Some(ContentBlockDelta::ToolUse(tool_output)) => {
918 if let Some(tool_use) = state
919 .tool_uses_by_index
920 .get_mut(&cb_delta.content_block_index)
921 {
922 tool_use.input_json.push_str(tool_output.input());
923 }
924 None
925 }
926 Some(ContentBlockDelta::ReasoningContent(thinking)) => match thinking {
927 ReasoningContentBlockDelta::Text(thoughts) => {
928 Some(Ok(LanguageModelCompletionEvent::Thinking {
929 text: thoughts,
930 signature: None,
931 }))
932 }
933 ReasoningContentBlockDelta::Signature(sig) => {
934 Some(Ok(LanguageModelCompletionEvent::Thinking {
935 text: "".into(),
936 signature: Some(sig),
937 }))
938 }
939 ReasoningContentBlockDelta::RedactedContent(redacted) => {
940 let content = String::from_utf8(redacted.into_inner())
941 .unwrap_or("REDACTED".to_string());
942 Some(Ok(LanguageModelCompletionEvent::Thinking {
943 text: content,
944 signature: None,
945 }))
946 }
947 _ => None,
948 },
949 _ => None,
950 },
951 ConverseStreamOutput::ContentBlockStart(cb_start) => {
952 if let Some(ContentBlockStart::ToolUse(tool_start)) = cb_start.start {
953 state.tool_uses_by_index.insert(
954 cb_start.content_block_index,
955 RawToolUse {
956 id: tool_start.tool_use_id,
957 name: tool_start.name,
958 input_json: String::new(),
959 },
960 );
961 }
962 None
963 }
964 ConverseStreamOutput::ContentBlockStop(cb_stop) => state
965 .tool_uses_by_index
966 .remove(&cb_stop.content_block_index)
967 .map(|tool_use| {
968 let input = if tool_use.input_json.is_empty() {
969 Value::Null
970 } else {
971 serde_json::Value::from_str(&tool_use.input_json)
972 .unwrap_or(Value::Null)
973 };
974
975 Ok(LanguageModelCompletionEvent::ToolUse(
976 LanguageModelToolUse {
977 id: tool_use.id.into(),
978 name: tool_use.name.into(),
979 is_input_complete: true,
980 raw_input: tool_use.input_json,
981 input,
982 thought_signature: None,
983 },
984 ))
985 }),
986 ConverseStreamOutput::Metadata(cb_meta) => cb_meta.usage.map(|metadata| {
987 Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
988 input_tokens: metadata.input_tokens as u64,
989 output_tokens: metadata.output_tokens as u64,
990 cache_creation_input_tokens: metadata
991 .cache_write_input_tokens
992 .unwrap_or_default()
993 as u64,
994 cache_read_input_tokens: metadata
995 .cache_read_input_tokens
996 .unwrap_or_default()
997 as u64,
998 }))
999 }),
1000 ConverseStreamOutput::MessageStop(message_stop) => {
1001 let stop_reason = match message_stop.stop_reason {
1002 StopReason::ToolUse => language_model::StopReason::ToolUse,
1003 _ => language_model::StopReason::EndTurn,
1004 };
1005 Some(Ok(LanguageModelCompletionEvent::Stop(stop_reason)))
1006 }
1007 _ => None,
1008 };
1009
1010 Some((result, state))
1011 }
1012 Err(err) => Some((
1013 Some(Err(LanguageModelCompletionError::Other(anyhow!(err)))),
1014 state,
1015 )),
1016 },
1017 None => None,
1018 }
1019 })
1020 .filter_map(|result| async move { result })
1021}
1022
1023struct ConfigurationView {
1024 access_key_id_editor: Entity<InputField>,
1025 secret_access_key_editor: Entity<InputField>,
1026 session_token_editor: Entity<InputField>,
1027 region_editor: Entity<InputField>,
1028 state: Entity<State>,
1029 load_credentials_task: Option<Task<()>>,
1030 focus_handle: FocusHandle,
1031}
1032
1033impl ConfigurationView {
1034 const PLACEHOLDER_ACCESS_KEY_ID_TEXT: &'static str = "XXXXXXXXXXXXXXXX";
1035 const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
1036 "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1037 const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1038 const PLACEHOLDER_REGION: &'static str = "us-east-1";
1039
1040 fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
1041 let focus_handle = cx.focus_handle();
1042
1043 cx.observe(&state, |_, _, cx| {
1044 cx.notify();
1045 })
1046 .detach();
1047
1048 let access_key_id_editor = cx.new(|cx| {
1049 InputField::new(window, cx, Self::PLACEHOLDER_ACCESS_KEY_ID_TEXT)
1050 .label("Access Key ID")
1051 .tab_index(0)
1052 .tab_stop(true)
1053 });
1054
1055 let secret_access_key_editor = cx.new(|cx| {
1056 InputField::new(window, cx, Self::PLACEHOLDER_SECRET_ACCESS_KEY_TEXT)
1057 .label("Secret Access Key")
1058 .tab_index(1)
1059 .tab_stop(true)
1060 });
1061
1062 let session_token_editor = cx.new(|cx| {
1063 InputField::new(window, cx, Self::PLACEHOLDER_SESSION_TOKEN_TEXT)
1064 .label("Session Token (Optional)")
1065 .tab_index(2)
1066 .tab_stop(true)
1067 });
1068
1069 let region_editor = cx.new(|cx| {
1070 InputField::new(window, cx, Self::PLACEHOLDER_REGION)
1071 .label("Region")
1072 .tab_index(3)
1073 .tab_stop(true)
1074 });
1075
1076 let load_credentials_task = Some(cx.spawn({
1077 let state = state.clone();
1078 async move |this, cx| {
1079 if let Some(task) = state
1080 .update(cx, |state, cx| state.authenticate(cx))
1081 .log_err()
1082 {
1083 // We don't log an error, because "not signed in" is also an error.
1084 let _ = task.await;
1085 }
1086 this.update(cx, |this, cx| {
1087 this.load_credentials_task = None;
1088 cx.notify();
1089 })
1090 .log_err();
1091 }
1092 }));
1093
1094 Self {
1095 access_key_id_editor,
1096 secret_access_key_editor,
1097 session_token_editor,
1098 region_editor,
1099 state,
1100 load_credentials_task,
1101 focus_handle,
1102 }
1103 }
1104
1105 fn save_credentials(
1106 &mut self,
1107 _: &menu::Confirm,
1108 _window: &mut Window,
1109 cx: &mut Context<Self>,
1110 ) {
1111 let access_key_id = self
1112 .access_key_id_editor
1113 .read(cx)
1114 .text(cx)
1115 .trim()
1116 .to_string();
1117 let secret_access_key = self
1118 .secret_access_key_editor
1119 .read(cx)
1120 .text(cx)
1121 .trim()
1122 .to_string();
1123 let session_token = self
1124 .session_token_editor
1125 .read(cx)
1126 .text(cx)
1127 .trim()
1128 .to_string();
1129 let session_token = if session_token.is_empty() {
1130 None
1131 } else {
1132 Some(session_token)
1133 };
1134 let region = self.region_editor.read(cx).text(cx).trim().to_string();
1135 let region = if region.is_empty() {
1136 "us-east-1".to_string()
1137 } else {
1138 region
1139 };
1140
1141 let state = self.state.clone();
1142 cx.spawn(async move |_, cx| {
1143 state
1144 .update(cx, |state, cx| {
1145 let credentials: BedrockCredentials = BedrockCredentials {
1146 region: region.clone(),
1147 access_key_id: access_key_id.clone(),
1148 secret_access_key: secret_access_key.clone(),
1149 session_token: session_token.clone(),
1150 };
1151
1152 state.set_credentials(credentials, cx)
1153 })?
1154 .await
1155 })
1156 .detach_and_log_err(cx);
1157 }
1158
1159 fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1160 self.access_key_id_editor
1161 .update(cx, |editor, cx| editor.set_text("", window, cx));
1162 self.secret_access_key_editor
1163 .update(cx, |editor, cx| editor.set_text("", window, cx));
1164 self.session_token_editor
1165 .update(cx, |editor, cx| editor.set_text("", window, cx));
1166 self.region_editor
1167 .update(cx, |editor, cx| editor.set_text("", window, cx));
1168
1169 let state = self.state.clone();
1170 cx.spawn(async move |_, cx| {
1171 state
1172 .update(cx, |state, cx| state.reset_credentials(cx))?
1173 .await
1174 })
1175 .detach_and_log_err(cx);
1176 }
1177
1178 fn should_render_editor(&self, cx: &Context<Self>) -> bool {
1179 self.state.read(cx).is_authenticated()
1180 }
1181
1182 fn on_tab(&mut self, _: &menu::SelectNext, window: &mut Window, _: &mut Context<Self>) {
1183 window.focus_next();
1184 }
1185
1186 fn on_tab_prev(
1187 &mut self,
1188 _: &menu::SelectPrevious,
1189 window: &mut Window,
1190 _: &mut Context<Self>,
1191 ) {
1192 window.focus_prev();
1193 }
1194}
1195
1196impl Render for ConfigurationView {
1197 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1198 let env_var_set = self.state.read(cx).credentials_from_env;
1199 let bedrock_settings = self.state.read(cx).settings.as_ref();
1200 let bedrock_method = bedrock_settings
1201 .as_ref()
1202 .and_then(|s| s.authentication_method.clone());
1203
1204 if self.load_credentials_task.is_some() {
1205 return div().child(Label::new("Loading credentials...")).into_any();
1206 }
1207
1208 let configured_label = if env_var_set {
1209 format!(
1210 "Access Key ID is set in {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, Secret Key is set in {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, Region is set in {ZED_BEDROCK_REGION_VAR} environment variables."
1211 )
1212 } else {
1213 match bedrock_method {
1214 Some(BedrockAuthMethod::Automatic) => "You are using automatic credentials.".into(),
1215 Some(BedrockAuthMethod::NamedProfile) => "You are using named profile.".into(),
1216 Some(BedrockAuthMethod::SingleSignOn) => {
1217 "You are using a single sign on profile.".into()
1218 }
1219 None => "You are using static credentials.".into(),
1220 }
1221 };
1222
1223 let tooltip_label = if env_var_set {
1224 Some(format!(
1225 "To reset your credentials, unset the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, and {ZED_BEDROCK_REGION_VAR} environment variables."
1226 ))
1227 } else if bedrock_method.is_some() {
1228 Some("You cannot reset credentials as they're being derived, check Zed settings to understand how.".to_string())
1229 } else {
1230 None
1231 };
1232
1233 if self.should_render_editor(cx) {
1234 return ConfiguredApiCard::new(configured_label)
1235 .disabled(env_var_set || bedrock_method.is_some())
1236 .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx)))
1237 .when_some(tooltip_label, |this, label| this.tooltip_label(label))
1238 .into_any_element();
1239 }
1240
1241 v_flex()
1242 .size_full()
1243 .track_focus(&self.focus_handle)
1244 .on_action(cx.listener(Self::on_tab))
1245 .on_action(cx.listener(Self::on_tab_prev))
1246 .on_action(cx.listener(ConfigurationView::save_credentials))
1247 .child(Label::new("To use Zed's agent with Bedrock, you can set a custom authentication strategy through the settings.json, or use static credentials."))
1248 .child(Label::new("But, to access models on AWS, you need to:").mt_1())
1249 .child(
1250 List::new()
1251 .child(
1252 ListBulletItem::new("")
1253 .child(Label::new("Grant permissions to the strategy you'll use according to the:"))
1254 .child(ButtonLink::new("Prerequisites", "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html"))
1255 )
1256 .child(
1257 ListBulletItem::new("")
1258 .child(Label::new("Select the models you would like access to:"))
1259 .child(ButtonLink::new("Bedrock Model Catalog", "https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess"))
1260 )
1261 )
1262 .child(self.render_static_credentials_ui())
1263 .child(
1264 Label::new(
1265 format!("You can also assign the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR} AND {ZED_BEDROCK_REGION_VAR} environment variables and restart Zed."),
1266 )
1267 .size(LabelSize::Small)
1268 .color(Color::Muted)
1269 .my_1(),
1270 )
1271 .child(
1272 Label::new(
1273 format!("Optionally, if your environment uses AWS CLI profiles, you can set {ZED_AWS_PROFILE_VAR}; if it requires a custom endpoint, you can set {ZED_AWS_ENDPOINT_VAR}; and if it requires a Session Token, you can set {ZED_BEDROCK_SESSION_TOKEN_VAR}."),
1274 )
1275 .size(LabelSize::Small)
1276 .color(Color::Muted),
1277 )
1278 .into_any()
1279 }
1280}
1281
1282impl ConfigurationView {
1283 fn render_static_credentials_ui(&self) -> impl IntoElement {
1284 v_flex()
1285 .my_2()
1286 .tab_group()
1287 .gap_1p5()
1288 .child(
1289 Label::new("Static Keys")
1290 .size(LabelSize::Default)
1291 .weight(FontWeight::BOLD),
1292 )
1293 .child(
1294 Label::new(
1295 "This method uses your AWS access key ID and secret access key directly.",
1296 )
1297 )
1298 .child(
1299 List::new()
1300 .child(
1301 ListBulletItem::new("")
1302 .child(Label::new("Create an IAM user in the AWS console with programmatic access"))
1303 .child(ButtonLink::new("IAM Console", "https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users"))
1304 )
1305 .child(
1306 ListBulletItem::new("")
1307 .child(Label::new("Attach the necessary Bedrock permissions to this"))
1308 .child(ButtonLink::new("user", "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html"))
1309 )
1310 .child(
1311 ListBulletItem::new("Copy the access key ID and secret access key when provided")
1312 )
1313 .child(
1314 ListBulletItem::new("Enter these credentials below")
1315 )
1316 )
1317 .child(self.access_key_id_editor.clone())
1318 .child(self.secret_access_key_editor.clone())
1319 .child(self.session_token_editor.clone())
1320 .child(self.region_editor.clone())
1321 }
1322}