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