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