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