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