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