Bedrock improvements

This commit is contained in:
Ryan Ward
2026-05-07 16:19:39 -05:00
parent 2d5bc55939
commit 9e0af95953
13 changed files with 68 additions and 18 deletions
+1 -1
View File
@@ -252,7 +252,7 @@ pub async fn generate_multi_agent_output(
messages, messages,
system_prompt, system_prompt,
tools, tools,
8192, 64000,
None, None,
true, true,
diagnostic_logger, diagnostic_logger,
+40 -17
View File
@@ -50,13 +50,11 @@ pub fn bedrock_stream_to_response_events(
let mut current_tool_use_id = String::new(); let mut current_tool_use_id = String::new();
let mut current_tool_name = String::new(); let mut current_tool_name = String::new();
let mut current_tool_input_json = String::new(); let mut current_tool_input_json = String::new();
let mut has_tool_calls = false; let mut _has_tool_calls = false;
let mut input_tokens: i32 = 0; let mut input_tokens: i32 = 0;
let mut output_tokens: i32 = 0; let mut output_tokens: i32 = 0;
let mut stop_reason = stream_finished::Reason::Done(stream_finished::Done {}); let mut stop_reason = stream_finished::Reason::Done(stream_finished::Done {});
const TEXT_FLUSH_THRESHOLD: usize = 20;
loop { loop {
match output.stream.recv().await { match output.stream.recv().await {
Ok(Some(event)) => match event { Ok(Some(event)) => match event {
@@ -65,10 +63,12 @@ pub fn bedrock_stream_to_response_events(
if let Some(start) = block_start.start() { if let Some(start) = block_start.start() {
match start { match start {
ContentBlockStart::ToolUse(tool_start) => { ContentBlockStart::ToolUse(tool_start) => {
has_tool_calls = true; _has_tool_calls = true;
if !text_flushed && !buffered_text.is_empty() { if !buffered_text.is_empty() {
if buffered_text.len() >= TEXT_FLUSH_THRESHOLD { let msg_id = current_text_message_id
let msg_id = Uuid::new_v4().to_string(); .clone()
.unwrap_or_else(|| Uuid::new_v4().to_string());
if !text_flushed {
current_text_message_id = Some(msg_id.clone()); current_text_message_id = Some(msg_id.clone());
text_flushed = true; text_flushed = true;
log::debug!("[bedrock] Flushing buffered text ({} chars) before tool call", buffered_text.len()); log::debug!("[bedrock] Flushing buffered text ({} chars) before tool call", buffered_text.len());
@@ -79,7 +79,13 @@ pub fn bedrock_stream_to_response_events(
); );
yield Ok(add_msg); yield Ok(add_msg);
} else { } else {
log::debug!("[bedrock] Discarding short text fragment ({} chars) before tool call: {:?}", buffered_text.len(), &buffered_text); log::debug!("[bedrock] Flushing remaining buffered text ({} chars) as append before tool call", buffered_text.len());
let append = build_append_text(
&task_id,
&msg_id,
&buffered_text,
);
yield Ok(append);
} }
buffered_text.clear(); buffered_text.clear();
} }
@@ -95,7 +101,7 @@ pub fn bedrock_stream_to_response_events(
if let Some(d) = delta.delta() { if let Some(d) = delta.delta() {
match d { match d {
ContentBlockDelta::Text(text) => { ContentBlockDelta::Text(text) => {
log::trace!("[bedrock] Text delta ({} chars): {:?}", text.len(), &text[..text.len().min(100)]); log::trace!("[bedrock] Text delta ({} chars)", text.len());
if text_flushed { if text_flushed {
let msg_id = current_text_message_id.as_ref().unwrap(); let msg_id = current_text_message_id.as_ref().unwrap();
let append = build_append_text( let append = build_append_text(
@@ -106,11 +112,10 @@ pub fn bedrock_stream_to_response_events(
yield Ok(append); yield Ok(append);
} else { } else {
buffered_text.push_str(text); buffered_text.push_str(text);
if buffered_text.len() >= TEXT_FLUSH_THRESHOLD { if buffered_text.len() >= 1 {
let msg_id = Uuid::new_v4().to_string(); let msg_id = Uuid::new_v4().to_string();
current_text_message_id = Some(msg_id.clone()); current_text_message_id = Some(msg_id.clone());
text_flushed = true; text_flushed = true;
log::debug!("[bedrock] Text reached flush threshold, creating message msg_id={msg_id}");
let add_msg = build_add_agent_output_message( let add_msg = build_add_agent_output_message(
&task_id, &task_id,
&msg_id, &msg_id,
@@ -187,6 +192,19 @@ pub fn bedrock_stream_to_response_events(
if let Some(ref logger) = diagnostic_logger { if let Some(ref logger) = diagnostic_logger {
logger.log_stream_error(&format!("{e}")); logger.log_stream_error(&format!("{e}"));
} }
if !buffered_text.is_empty() {
let msg_id = current_text_message_id
.clone()
.unwrap_or_else(|| Uuid::new_v4().to_string());
if !text_flushed {
let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text);
yield Ok(add_msg);
} else {
let append = build_append_text(&task_id, &msg_id, &buffered_text);
yield Ok(append);
}
buffered_text.clear();
}
yield Err(Arc::new(AIApiError::Stream { yield Err(Arc::new(AIApiError::Stream {
stream_type: "bedrock_converse", stream_type: "bedrock_converse",
source: anyhow::anyhow!("Bedrock stream error: {}", e), source: anyhow::anyhow!("Bedrock stream error: {}", e),
@@ -196,13 +214,18 @@ pub fn bedrock_stream_to_response_events(
} }
} }
if !text_flushed && !buffered_text.is_empty() && !has_tool_calls { if !buffered_text.is_empty() {
let msg_id = Uuid::new_v4().to_string(); let msg_id = current_text_message_id
.clone()
.unwrap_or_else(|| Uuid::new_v4().to_string());
log::debug!("[bedrock] Flushing remaining buffered text ({} chars) at stream end", buffered_text.len()); log::debug!("[bedrock] Flushing remaining buffered text ({} chars) at stream end", buffered_text.len());
let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text); if !text_flushed {
yield Ok(add_msg); let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text);
} else if !text_flushed && !buffered_text.is_empty() && has_tool_calls { yield Ok(add_msg);
log::debug!("[bedrock] Discarding short unflushed text ({} chars) - stream ended with tool calls", buffered_text.len()); } else {
let append = build_append_text(&task_id, &msg_id, &buffered_text);
yield Ok(append);
}
} }
log::info!("[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}"); log::info!("[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}");
+5
View File
@@ -1,3 +1,4 @@
#[allow(dead_code)]
pub(super) mod user_persistence; pub(super) mod user_persistence;
use std::result::Result as StdResult; use std::result::Result as StdResult;
@@ -9,9 +10,11 @@ use settings::Setting as _;
use uuid::Uuid; use uuid::Uuid;
use galaxy_core::channel::ChannelState; use galaxy_core::channel::ChannelState;
use galaxy_core::features::FeatureFlag; use galaxy_core::features::FeatureFlag;
#[allow(unused_imports)]
use galaxy_graphql::mutations::create_anonymous_user::{ use galaxy_graphql::mutations::create_anonymous_user::{
AnonymousUserType, CreateAnonymousUserResult, AnonymousUserType, CreateAnonymousUserResult,
}; };
#[allow(unused_imports)]
use galaxyui::{clipboard::ClipboardContent, Entity, ModelContext, SingletonEntity, UpdateModel}; use galaxyui::{clipboard::ClipboardContent, Entity, ModelContext, SingletonEntity, UpdateModel};
use super::auth_state::{AuthState, PersistAction}; use super::auth_state::{AuthState, PersistAction};
@@ -55,6 +58,7 @@ use url::Url;
use user_persistence::PersistedUser; use user_persistence::PersistedUser;
#[derive(Debug)] #[derive(Debug)]
#[allow(dead_code)]
pub enum AuthManagerEvent { pub enum AuthManagerEvent {
/// Successfully authenticated a user with no errors. /// Successfully authenticated a user with no errors.
AuthComplete, AuthComplete,
@@ -102,6 +106,7 @@ pub struct AuthManager {
pending_auth_state: Option<String>, pending_auth_state: Option<String>,
} }
#[allow(dead_code)]
impl AuthManager { impl AuthManager {
/// Creates a new instance of the AuthManager. The auth state must already be initialized through /// Creates a new instance of the AuthManager. The auth state must already be initialized through
/// [`AuthStateProvider`]. /// [`AuthStateProvider`].
@@ -9,6 +9,7 @@ pub enum AuthOverrideWarningModalVariant {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[allow(dead_code)]
pub enum AuthOverrideWarningModalEvent { pub enum AuthOverrideWarningModalEvent {
Close, Close,
BulkExport, BulkExport,
+3
View File
@@ -87,6 +87,7 @@ impl AuthState {
state state
} }
#[allow(dead_code)]
fn should_use_test_user() -> bool { fn should_use_test_user() -> bool {
cfg!(any(test, feature = "skip_login")) || ChannelState::channel() == Channel::Integration cfg!(any(test, feature = "skip_login")) || ChannelState::channel() == Channel::Integration
} }
@@ -131,6 +132,7 @@ impl AuthState {
} }
/// Applies a deserialized PersistedUser, splitting it into User and Credentials. /// Applies a deserialized PersistedUser, splitting it into User and Credentials.
#[allow(dead_code)]
fn apply_persisted_user(&self, persisted: PersistedUser) { fn apply_persisted_user(&self, persisted: PersistedUser) {
let user = User { let user = User {
is_onboarded: persisted.is_onboarded, is_onboarded: persisted.is_onboarded,
@@ -171,6 +173,7 @@ impl AuthState {
/// Updates the Firebase auth tokens within the current credentials. /// Updates the Firebase auth tokens within the current credentials.
/// Reports an error if the current credentials are not Firebase. /// Reports an error if the current credentials are not Firebase.
#[allow(dead_code)]
pub(crate) fn update_firebase_tokens(&self, new_auth_tokens: FirebaseAuthTokens) { pub(crate) fn update_firebase_tokens(&self, new_auth_tokens: FirebaseAuthTokens) {
let mut write_lock = self.credentials.write(); let mut write_lock = self.credentials.write();
if let Some(Credentials::Firebase(tokens)) = write_lock.as_mut() { if let Some(Credentials::Firebase(tokens)) = write_lock.as_mut() {
+3
View File
@@ -27,6 +27,7 @@ pub enum AuthViewVariant {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[allow(dead_code)]
pub enum AuthViewEvent { pub enum AuthViewEvent {
Close, Close,
} }
@@ -67,8 +68,10 @@ impl TypedActionView for AuthView {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[allow(dead_code)]
pub enum LoginFailureReason { pub enum LoginFailureReason {
InvalidRedirectUrl { was_pasted: bool }, InvalidRedirectUrl { was_pasted: bool },
} }
#[allow(dead_code)]
pub fn init(_app: &mut AppContext) {} pub fn init(_app: &mut AppContext) {}
+1
View File
@@ -9,6 +9,7 @@ pub enum LoginSlideSource {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[allow(dead_code)]
pub enum LoginSlideEvent { pub enum LoginSlideEvent {
BackToOnboarding, BackToOnboarding,
LoginLaterConfirmed, LoginLaterConfirmed,
+2
View File
@@ -51,9 +51,11 @@ use crate::{persistence, GlobalResourceHandlesProvider};
use crate::{report_if_error, send_telemetry_sync_from_app_ctx}; use crate::{report_if_error, send_telemetry_sync_from_app_ctx};
/// Prefix for API keys used in authentication /// Prefix for API keys used in authentication
#[allow(dead_code)]
#[cfg_attr(target_family = "wasm", allow(dead_code))] #[cfg_attr(target_family = "wasm", allow(dead_code))]
pub const API_KEY_PREFIX: &str = "wk-"; pub const API_KEY_PREFIX: &str = "wk-";
#[allow(dead_code)]
pub fn init(_app: &mut AppContext) {} pub fn init(_app: &mut AppContext) {}
/// If the app has running processes or dirty objects, we'll show a confirmation modal before logging out. /// If the app has running processes or dirty objects, we'll show a confirmation modal before logging out.
+1
View File
@@ -1,6 +1,7 @@
use galaxyui::{AppContext, Element, Entity, TypedActionView, View, ViewContext}; use galaxyui::{AppContext, Element, Entity, TypedActionView, View, ViewContext};
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[allow(dead_code)]
pub enum PasteAuthTokenModalEvent { pub enum PasteAuthTokenModalEvent {
Cancelled, Cancelled,
} }
+1
View File
@@ -11,6 +11,7 @@ mod rendering;
pub use block_onboarding_layer::{BlockOnboarding, BLOCK_ONBOARDING_LAYER}; pub use block_onboarding_layer::{BlockOnboarding, BLOCK_ONBOARDING_LAYER};
pub use free_tier_default_model_layer::{FreeTierDefaultModel, FREE_TIER_DEFAULT_MODEL_LAYER}; pub use free_tier_default_model_layer::{FreeTierDefaultModel, FREE_TIER_DEFAULT_MODEL_LAYER};
pub use improved_palette_search_layer::{ImprovedPaletteSearch, IMPROVED_PALETTE_SEARCH_LAYER}; pub use improved_palette_search_layer::{ImprovedPaletteSearch, IMPROVED_PALETTE_SEARCH_LAYER};
#[allow(unused_imports)]
pub use login_layer::{AuthFlowInstructions, LOGIN_LAYER}; pub use login_layer::{AuthFlowInstructions, LOGIN_LAYER};
use galaxy_core::user_preferences::GetUserPreferences as _; use galaxy_core::user_preferences::GetUserPreferences as _;
+1
View File
@@ -135,6 +135,7 @@ lazy_static! {
/// that this is hard-coded for the default Dark theme. This is because it is only used by the /// that this is hard-coded for the default Dark theme. This is because it is only used by the
/// AuthView and OnboardingSurveyModal which do not respect the chosen theme. So, do not use this for Views /// AuthView and OnboardingSurveyModal which do not respect the chosen theme. So, do not use this for Views
/// which respect themes. /// which respect themes.
#[allow(dead_code)]
pub(crate) fn unthemed_window_border() -> Border { pub(crate) fn unthemed_window_border() -> Border {
if cfg!(all(not(target_os = "macos"), not(target_family = "wasm"))) { if cfg!(all(not(target_os = "macos"), not(target_family = "wasm"))) {
// The 15% blend of fg into bg is the "ui surface" color. // The 15% blend of fg into bg is the "ui surface" color.
+5
View File
@@ -18,8 +18,10 @@ use crate::ai::predict::generate_am_query_suggestions;
use crate::ai::predict::generate_am_query_suggestions::GenerateAMQuerySuggestionsRequest; use crate::ai::predict::generate_am_query_suggestions::GenerateAMQuerySuggestionsRequest;
use crate::ai::predict::predict_am_queries::{PredictAMQueriesRequest, PredictAMQueriesResponse}; use crate::ai::predict::predict_am_queries::{PredictAMQueriesRequest, PredictAMQueriesResponse};
use crate::ai::voice::transcribe::{TranscribeRequest, TranscribeResponse}; use crate::ai::voice::transcribe::{TranscribeRequest, TranscribeResponse};
#[allow(unused_imports)]
use crate::auth::auth_manager::AuthManager; use crate::auth::auth_manager::AuthManager;
use crate::auth::auth_state::AuthState; use crate::auth::auth_state::AuthState;
#[allow(unused_imports)]
use crate::server::graphql::default_request_options; use crate::server::graphql::default_request_options;
use crate::server::server_api::presigned_upload::HttpStatusError; use crate::server::server_api::presigned_upload::HttpStatusError;
use ai::AIClient; use ai::AIClient;
@@ -357,6 +359,7 @@ cfg_if::cfg_if! {
/// Most errors should be handled in callbacks to individual APIs, rather than sent over the /// Most errors should be handled in callbacks to individual APIs, rather than sent over the
/// server API channel. /// server API channel.
#[derive(Clone)] #[derive(Clone)]
#[allow(dead_code)]
pub enum ServerApiEvent { pub enum ServerApiEvent {
/// We made a staging API call that was blocked, which may indicate a firewall misconfiguration. /// We made a staging API call that was blocked, which may indicate a firewall misconfiguration.
StagingAccessBlocked, StagingAccessBlocked,
@@ -399,8 +402,10 @@ pub struct ServerApi {
telemetry_api: TelemetryApi, telemetry_api: TelemetryApi,
last_server_time: Arc<Mutex<Option<ServerTime>>>, last_server_time: Arc<Mutex<Option<ServerTime>>>,
// We technically use OAuth2 for headless device authentication. // We technically use OAuth2 for headless device authentication.
#[allow(dead_code)]
oauth_client: self::auth::OAuth2Client, oauth_client: self::auth::OAuth2Client,
/// Cached ambient workload token for requests from ambient agents. /// Cached ambient workload token for requests from ambient agents.
#[allow(dead_code)]
ambient_workload_token: Arc<Mutex<Option<galaxy_isolation_platform::WorkloadToken>>>, ambient_workload_token: Arc<Mutex<Option<galaxy_isolation_platform::WorkloadToken>>>,
/// The ambient agent task ID for requests from cloud agents. /// The ambient agent task ID for requests from cloud agents.
ambient_agent_task_id: Arc<RwLock<Option<AmbientAgentTaskId>>>, ambient_agent_task_id: Arc<RwLock<Option<AmbientAgentTaskId>>>,
+4
View File
@@ -54,6 +54,7 @@ pub struct FetchUserResult {
pub llms: crate::ai::llms::ModelsByFeature, pub llms: crate::ai::llms::ModelsByFeature,
} }
#[allow(dead_code)]
#[cfg_attr(test, automock)] #[cfg_attr(test, automock)]
#[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(not(target_family = "wasm"), async_trait)]
#[cfg_attr(target_family = "wasm", async_trait(?Send))] #[cfg_attr(target_family = "wasm", async_trait(?Send))]
@@ -272,6 +273,7 @@ pub type OAuth2Client = oauth2::basic::BasicClient<
>; >;
#[derive(Error, Debug)] #[derive(Error, Debug)]
#[allow(dead_code)]
/// Error type when retrieving a user and validating it against Firebase. /// Error type when retrieving a user and validating it against Firebase.
pub enum UserAuthenticationError { pub enum UserAuthenticationError {
/// The user's refresh token is invalid. This could occur if the user authed through /// The user's refresh token is invalid. This could occur if the user authed through
@@ -330,6 +332,7 @@ impl From<FirebaseError> for UserAuthenticationError {
} }
#[derive(Error, Debug)] #[derive(Error, Debug)]
#[allow(dead_code)]
/// Error type when creating anonymous users /// Error type when creating anonymous users
pub enum AnonymousUserCreationError { pub enum AnonymousUserCreationError {
#[error("The network request to create the anonymous user failed")] #[error("The network request to create the anonymous user failed")]
@@ -347,6 +350,7 @@ pub enum AnonymousUserCreationError {
} }
#[derive(Error, Debug)] #[derive(Error, Debug)]
#[allow(dead_code)]
/// Error type when minting a new custom token for an anonymous user /// Error type when minting a new custom token for an anonymous user
pub enum MintCustomTokenError { pub enum MintCustomTokenError {
#[error("Received a user facing error: {0}")] #[error("Received a user facing error: {0}")]