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