Files
galaxy/app/src/ai/agent/api/impl.rs
T

316 lines
13 KiB
Rust

use std::collections::HashMap;
use std::sync::Arc;
use futures_util::StreamExt;
use galaxy_core::features::FeatureFlag;
use warp_multi_agent_api as api;
use super::convert_to::convert_input;
use super::{ConvertToAPITypeError, RequestParams, ResponseStream};
use crate::ai::agent::redaction;
use crate::ai::openai::translator as openai_translator;
use crate::ai::provider::ProviderConfig;
use crate::server::server_api::AIApiError;
use crate::terminal::model::session::SessionType;
pub async fn generate_multi_agent_output(
provider_config: ProviderConfig,
mut params: RequestParams,
cancellation_rx: futures::channel::oneshot::Receiver<()>,
) -> Result<ResponseStream, ConvertToAPITypeError> {
let supported_tools_override = params.supported_tools_override.take();
let supported_tools = supported_tools_override
.clone()
.unwrap_or_else(|| get_supported_tools(&params));
let supported_cli_agent_tools =
supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(&params));
let mut logging_metadata = HashMap::new();
if let Some(ref metadata) = params.metadata {
logging_metadata.insert(
"is_autodetected_user_query".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::BoolValue(
metadata.is_autodetected_user_query,
)),
},
);
logging_metadata.insert(
"entrypoint".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::StringValue(
metadata.entrypoint.entrypoint(),
)),
},
);
logging_metadata.insert(
"is_auto_resume_after_error".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::BoolValue(
metadata.is_auto_resume_after_error,
)),
},
);
}
if params.should_redact_secrets {
redaction::redact_inputs(&mut params.input);
}
let rig_params = matches!(
&provider_config,
ProviderConfig::OpenAI(config) if config.use_rig
)
.then(|| params.clone());
let mut request = api::Request {
task_context: Some(api::request::TaskContext {
tasks: params.tasks,
}),
input: Some(convert_input(params.input)?),
settings: Some(api::request::Settings {
model_config: Some(api::request::settings::ModelConfig {
base: params.model.clone().into(),
cli_agent: params.cli_agent_model.clone().into(),
computer_use_agent: params.computer_use_model.clone().into(),
base_model_context_window_limit: params.context_window_limit.unwrap_or(0),
..Default::default()
}),
rules_enabled: params.is_memory_enabled,
warp_drive_context_enabled: params.warp_drive_context_enabled,
web_context_retrieval_enabled: true,
supports_parallel_tool_calls: true,
use_anthropic_text_editor_tools: false,
planning_enabled: params.planning_enabled,
supports_create_files: true,
supported_tools: supported_tools.into_iter().map(Into::into).collect(),
supports_long_running_commands: true,
should_preserve_file_content_in_history: true,
supports_todos_ui: true,
supports_linked_code_blocks: FeatureFlag::LinkedCodeBlocks.is_enabled(),
supports_started_child_task_message: true,
// Galaxy's direct providers only receive tools with local schemas and
// executors. Hosted-only suggestion/orchestration capability bits must
// remain false so models do not plan around unavailable Warp services.
supports_suggest_prompt: false,
supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
supports_reasoning_message: true,
api_keys: params.api_keys,
autonomy_level: params.autonomy_level.into(),
isolation_level: params.isolation_level.into(),
web_search_enabled: params.web_search_enabled,
supported_cli_agent_tools: supported_cli_agent_tools
.into_iter()
.map(Into::into)
.collect(),
supports_v4a_file_diffs: FeatureFlag::V4AFileDiffs.is_enabled(),
supports_summarization_via_message_replacement:
FeatureFlag::SummarizationViaMessageReplacement.is_enabled(),
supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(),
supports_research_agent: params.research_agent_enabled,
supports_orchestration_v2: false,
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,
conversation_id: params
.conversation_token
.as_ref()
.map(|token| token.as_str().to_string())
.unwrap_or_default(),
ambient_agent_task_id: params
.ambient_agent_task_id
.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())
.unwrap_or_default()
} else {
String::new()
},
parent_agent_id: params.parent_agent_id.unwrap_or_default(),
agent_name: params.agent_name.unwrap_or_default(),
}),
existing_suggestions: params
.existing_suggestions
.map(|suggestions| suggestions.into()),
mcp_context: params.mcp_context.map(Into::into),
};
match provider_config {
ProviderConfig::OpenAI(config) if config.use_rig => {
Ok(crate::ai::runtime::rig_openai_response_stream(
config,
rig_params.expect("Rig request parameters should be retained for a Rig model"),
&mut request,
cancellation_rx,
))
}
ProviderConfig::OpenAI(config) => {
let translator_request = openai_translator::TranslatorRequest {
config,
model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id.clone(),
message_history: params.bedrock_message_history.clone(),
tool_result_archive: params.bedrock_tool_result_archive.clone(),
progressive_summary: params.bedrock_progressive_summary.clone(),
messages_sent: params.bedrock_messages_sent.clone(),
global_rules: params.global_rules.clone(),
};
match openai_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!("[openai] Translator error: {e}");
let err = Arc::new(
crate::server::server_api::AIApiError::Stream {
stream_type: "openai_chat_completions",
source: anyhow::anyhow!("{e}"),
}
.into_quota_limit_if_provider_budget_exhausted(),
);
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
ProviderConfig::Bedrock(config) => {
let translator_request = crate::ai::bedrock::translator::TranslatorRequest {
config,
model_id: params.model.as_str().to_string(),
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(),
global_rules: params.global_rules.clone(),
};
match crate::ai::bedrock::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",
source: anyhow::anyhow!("{e}"),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
ProviderConfig::None => {
// No provider configured — do not fall back to Warp's cloud API.
let err = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "none",
source: anyhow::anyhow!(
"No AI provider configured. Enable Bedrock or OpenAI/LiteLLM in settings."
),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
let mut supported_tools = vec![
api::ToolType::Grep,
api::ToolType::FileGlob,
api::ToolType::FileGlobV2,
api::ToolType::ReadMcpResource,
api::ToolType::CallMcpTool,
api::ToolType::RunShellCommand,
api::ToolType::Subagent,
api::ToolType::WriteToLongRunningShellCommand,
api::ToolType::ReadShellCommandOutput,
api::ToolType::ReadDocuments,
api::ToolType::CreateDocuments,
api::ToolType::EditDocuments,
];
if FeatureFlag::ConversationsAsContext.is_enabled() {
supported_tools.push(api::ToolType::FetchConversation);
}
match params.session_context.session_type() {
None | Some(SessionType::Local) => {
supported_tools.extend(&[
api::ToolType::ReadFiles,
api::ToolType::ApplyFileDiffs,
api::ToolType::SearchCodebase,
]);
}
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 }) => {}
}
if FeatureFlag::ListSkills.is_enabled() {
supported_tools.push(api::ToolType::ReadSkill);
}
if FeatureFlag::AskUserQuestion.is_enabled() && params.ask_user_question_enabled {
supported_tools.push(api::ToolType::AskUserQuestion);
}
supported_tools
}
fn get_supported_cli_agent_tools(params: &RequestParams) -> Vec<api::ToolType> {
let mut supported_cli_agent_tools = vec![
api::ToolType::WriteToLongRunningShellCommand,
api::ToolType::ReadShellCommandOutput,
api::ToolType::Grep,
api::ToolType::FileGlob,
api::ToolType::FileGlobV2,
];
if FeatureFlag::TransferControlTool.is_enabled() {
supported_cli_agent_tools.push(api::ToolType::TransferShellCommandControlToUser);
}
match params.session_context.session_type() {
None | Some(SessionType::Local) => {
supported_cli_agent_tools
.extend(&[api::ToolType::ReadFiles, api::ToolType::SearchCodebase]);
}
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 }) => {}
}
supported_cli_agent_tools
}
#[cfg(test)]
#[path = "impl_tests.rs"]
mod tests;