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