first pass of merging in warp (doesn't build)
This commit is contained in:
@@ -1,15 +1,15 @@
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{ai::agent::redaction, terminal::model::session::SessionType};
|
||||
use futures_util::StreamExt;
|
||||
use galaxy_core::features::FeatureFlag;
|
||||
use warp_multi_agent_api as api;
|
||||
|
||||
use crate::ai::bedrock::translator::{self, TranslatorRequest};
|
||||
use crate::ai::openai::translator as openai_translator;
|
||||
use crate::ai::provider::ProviderConfig;
|
||||
|
||||
use super::{convert_to::convert_input, ConvertToAPITypeError, RequestParams, ResponseStream};
|
||||
use super::convert_to::convert_input;
|
||||
use super::{ConvertToAPITypeError, RequestParams, ResponseStream};
|
||||
use crate::ai::agent::redaction;
|
||||
use crate::server::server_api::{AIApiError, ServerApi};
|
||||
use crate::terminal::model::session::SessionType;
|
||||
|
||||
pub async fn generate_multi_agent_output(
|
||||
provider_config: ProviderConfig,
|
||||
@@ -53,10 +53,10 @@ pub async fn generate_multi_agent_output(
|
||||
redaction::redact_inputs(&mut params.input);
|
||||
}
|
||||
|
||||
let mut api_keys = params.api_keys;
|
||||
if let Some(api_keys) = &mut api_keys {
|
||||
api_keys.allow_use_of_warp_credits = params.allow_use_of_warp_credits_with_byok;
|
||||
}
|
||||
let api_keys = api_keys_with_warp_credit_fallback_setting(
|
||||
params.api_keys,
|
||||
params.allow_use_of_warp_credits,
|
||||
);
|
||||
|
||||
let mut request = api::Request {
|
||||
task_context: Some(api::request::TaskContext {
|
||||
@@ -68,6 +68,7 @@ pub async fn generate_multi_agent_output(
|
||||
base: params.model.into(),
|
||||
cli_agent: params.cli_agent_model.into(),
|
||||
computer_use_agent: params.computer_use_model.into(),
|
||||
base_model_context_window_limit: params.context_window_limit.unwrap_or(0),
|
||||
..Default::default()
|
||||
}),
|
||||
rules_enabled: params.is_memory_enabled,
|
||||
@@ -99,7 +100,11 @@ pub async fn generate_multi_agent_output(
|
||||
FeatureFlag::SummarizationViaMessageReplacement.is_enabled(),
|
||||
supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(),
|
||||
supports_research_agent: params.research_agent_enabled,
|
||||
supports_orchestration_v2: FeatureFlag::OrchestrationV2.is_enabled(),
|
||||
supports_orchestration_v2: supports_orchestration_v2(params.orchestration_enabled),
|
||||
supports_background_computer_use: FeatureFlag::BackgroundComputerUse.is_enabled()
|
||||
&& computer_use::background_supported(),
|
||||
custom_model_providers: params.custom_model_providers,
|
||||
custom_model_routers: params.custom_model_routers,
|
||||
}),
|
||||
metadata: Some(api::request::Metadata {
|
||||
logging: logging_metadata,
|
||||
@@ -113,6 +118,8 @@ pub async fn generate_multi_agent_output(
|
||||
.map(|id| id.to_string())
|
||||
.unwrap_or_default(),
|
||||
forked_from_conversation_id: if params.conversation_token.is_none() {
|
||||
// We only include this param on our initial request to the server
|
||||
// (when the forked conversation has not been assigned a new id yet).
|
||||
params
|
||||
.forked_from_conversation_token
|
||||
.map(|token| token.as_str().to_string())
|
||||
@@ -129,41 +136,19 @@ pub async fn generate_multi_agent_output(
|
||||
mcp_context: params.mcp_context.map(Into::into),
|
||||
};
|
||||
|
||||
let model_id = request
|
||||
.settings
|
||||
.as_ref()
|
||||
.and_then(|s| s.model_config.as_ref())
|
||||
.map(|mc| mc.base.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
match provider_config {
|
||||
ProviderConfig::Bedrock(config) => {
|
||||
let translator_request = TranslatorRequest {
|
||||
config,
|
||||
model_id,
|
||||
root_task_id: params.root_task_id.clone(),
|
||||
bedrock_message_history: params.bedrock_message_history.clone(),
|
||||
bedrock_tool_result_archive: params.bedrock_tool_result_archive.clone(),
|
||||
bedrock_progressive_summary: params.bedrock_progressive_summary.clone(),
|
||||
bedrock_messages_sent: params.bedrock_messages_sent.clone(),
|
||||
};
|
||||
|
||||
match translator::execute(translator_request, &mut request).await {
|
||||
Ok(stream) => {
|
||||
let output_stream = stream.take_until(cancellation_rx);
|
||||
Ok(Box::pin(output_stream))
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!("[bedrock] Translator error: {e}");
|
||||
let err = Arc::new(crate::server::server_api::AIApiError::Stream {
|
||||
stream_type: "bedrock_converse",
|
||||
source: anyhow::anyhow!("{e}"),
|
||||
});
|
||||
let (tx, rx) = async_channel::unbounded();
|
||||
let _ = tx.send(Err(err)).await;
|
||||
Ok(Box::pin(rx))
|
||||
}
|
||||
}
|
||||
let response_stream =
|
||||
warp_multi_agent_client::generate_multi_agent_output(server_api.as_ref(), &request).await;
|
||||
match response_stream {
|
||||
Ok(stream) => {
|
||||
let output_stream = stream
|
||||
.then(|result| async {
|
||||
match result {
|
||||
Ok(event) => Ok(event),
|
||||
Err(error) => Err(convert_multi_agent_client_error(error).await),
|
||||
}
|
||||
})
|
||||
.take_until(cancellation_rx);
|
||||
Ok(Box::pin(output_stream))
|
||||
}
|
||||
ProviderConfig::OpenAI(config) => {
|
||||
let translator_request = openai_translator::TranslatorRequest {
|
||||
@@ -202,12 +187,53 @@ pub async fn generate_multi_agent_output(
|
||||
),
|
||||
});
|
||||
let (tx, rx) = async_channel::unbounded();
|
||||
let _ = tx.send(Err(err)).await;
|
||||
let _ = tx
|
||||
.send(Err(convert_multi_agent_client_error(e).await))
|
||||
.await;
|
||||
Ok(Box::pin(rx))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn convert_multi_agent_client_error(
|
||||
error: warp_multi_agent_client::Error,
|
||||
) -> Arc<AIApiError> {
|
||||
let error = match error {
|
||||
warp_multi_agent_client::Error::Authentication(error)
|
||||
| warp_multi_agent_client::Error::AmbientHeaders(error) => AIApiError::Other(error),
|
||||
warp_multi_agent_client::Error::Base64Decode(error) => {
|
||||
AIApiError::Other(anyhow::Error::from(error))
|
||||
}
|
||||
warp_multi_agent_client::Error::ProtobufDecode(error) => {
|
||||
AIApiError::Other(anyhow::Error::from(error))
|
||||
}
|
||||
warp_multi_agent_client::Error::EventSource(error) => {
|
||||
AIApiError::from_stream_error("GenerateMultiAgentOutput", *error).await
|
||||
}
|
||||
};
|
||||
Arc::new(error)
|
||||
}
|
||||
|
||||
fn api_keys_with_warp_credit_fallback_setting(
|
||||
api_keys: Option<api::request::settings::ApiKeys>,
|
||||
allow_use_of_warp_credits: bool,
|
||||
) -> Option<api::request::settings::ApiKeys> {
|
||||
match api_keys {
|
||||
Some(mut api_keys) => {
|
||||
api_keys.allow_use_of_warp_credits = allow_use_of_warp_credits;
|
||||
Some(api_keys)
|
||||
}
|
||||
None if allow_use_of_warp_credits => Some(api::request::settings::ApiKeys {
|
||||
allow_use_of_warp_credits: true,
|
||||
..Default::default()
|
||||
}),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn supports_orchestration_v2(orchestration_enabled: bool) -> bool {
|
||||
orchestration_enabled
|
||||
}
|
||||
fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
let mut supported_tools = vec![
|
||||
api::ToolType::Grep,
|
||||
@@ -245,7 +271,14 @@ fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
}
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
|
||||
// Remote session with a known host — enable tools that route
|
||||
// through RemoteServerClient. The host_id is only populated
|
||||
// after a successful connection handshake, so its presence is a
|
||||
// sufficient proxy for client availability.
|
||||
supported_tools.extend(&[api::ToolType::ReadFiles, api::ToolType::ApplyFileDiffs]);
|
||||
if FeatureFlag::RemoteCodebaseIndexing.is_enabled() {
|
||||
supported_tools.push(api::ToolType::SearchCodebase);
|
||||
}
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: None }) => {}
|
||||
}
|
||||
@@ -264,12 +297,10 @@ fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
}
|
||||
|
||||
if params.orchestration_enabled {
|
||||
supported_tools.push(if FeatureFlag::OrchestrationV2.is_enabled() {
|
||||
api::ToolType::StartAgentV2
|
||||
} else {
|
||||
api::ToolType::StartAgent
|
||||
});
|
||||
supported_tools.push(api::ToolType::SendMessageToAgent);
|
||||
supported_tools.extend([api::ToolType::RunAgents, api::ToolType::SendMessageToAgent]);
|
||||
// Declare client-handled wait_for_events so the server doesn't
|
||||
// fall back to the legacy server-handled form.
|
||||
supported_tools.push(api::ToolType::WaitForEvents);
|
||||
}
|
||||
|
||||
if FeatureFlag::AskUserQuestion.is_enabled() && params.ask_user_question_enabled {
|
||||
@@ -299,6 +330,9 @@ fn get_supported_cli_agent_tools(params: &RequestParams) -> Vec<api::ToolType> {
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: Some(_) }) => {
|
||||
supported_cli_agent_tools.push(api::ToolType::ReadFiles);
|
||||
if FeatureFlag::RemoteCodebaseIndexing.is_enabled() {
|
||||
supported_cli_agent_tools.push(api::ToolType::SearchCodebase);
|
||||
}
|
||||
}
|
||||
Some(SessionType::WarpifiedRemote { host_id: None }) => {}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user