693 lines
22 KiB
Rust
693 lines
22 KiB
Rust
use std::sync::{Arc, Mutex};
|
|
|
|
use futures::StreamExt;
|
|
use serde_json::json;
|
|
|
|
use super::client::{BedrockClient, BedrockClientConfig};
|
|
use super::convert::{ConversationMessage, MessageContent, MessageRole, ToolDefinition};
|
|
use crate::settings::ai::BedrockAuthMethod;
|
|
|
|
fn get_test_config() -> Option<BedrockClientConfig> {
|
|
if std::env::var("BEDROCK_INTEGRATION_TEST").is_err() {
|
|
return None;
|
|
}
|
|
|
|
let profile =
|
|
std::env::var("BEDROCK_TEST_PROFILE").unwrap_or_else(|_| "coding-assistant".into());
|
|
let region = std::env::var("BEDROCK_TEST_REGION").unwrap_or_else(|_| "us-east-1".into());
|
|
|
|
Some(BedrockClientConfig {
|
|
auth_method: BedrockAuthMethod::Profile,
|
|
profile,
|
|
region,
|
|
access_key_id: String::new(),
|
|
secret_access_key: String::new(),
|
|
cross_region_inference: false,
|
|
})
|
|
}
|
|
|
|
fn get_test_model() -> String {
|
|
std::env::var("BEDROCK_TEST_MODEL").unwrap_or_else(|_| {
|
|
"arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy".into()
|
|
})
|
|
}
|
|
|
|
struct StreamOutput {
|
|
text: String,
|
|
tool_calls: Vec<ToolCallInfo>,
|
|
finished_reason: Option<String>,
|
|
total_tokens: u32,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct ToolCallInfo {
|
|
name: String,
|
|
input_json: String,
|
|
}
|
|
|
|
async fn collect_stream_output(
|
|
client: &BedrockClient,
|
|
model: &str,
|
|
messages: Vec<ConversationMessage>,
|
|
system_prompt: Option<String>,
|
|
tools: Vec<ToolDefinition>,
|
|
) -> StreamOutput {
|
|
let stream = client
|
|
.converse_stream(
|
|
model,
|
|
"test-task-id",
|
|
true,
|
|
messages,
|
|
system_prompt,
|
|
None,
|
|
tools,
|
|
8192,
|
|
None,
|
|
false,
|
|
None,
|
|
None,
|
|
Arc::new(Mutex::new(Vec::new())),
|
|
false,
|
|
)
|
|
.await
|
|
.expect("converse_stream should succeed");
|
|
|
|
let mut text = String::new();
|
|
let tool_calls = Vec::new();
|
|
let mut finished_reason = None;
|
|
let mut total_tokens = 0u32;
|
|
|
|
let mut stream = stream;
|
|
while let Some(event) = stream.next().await {
|
|
let event = event.expect("stream event should be Ok");
|
|
if let Some(event_type) = event.r#type {
|
|
use warp_multi_agent_api::response_event::Type;
|
|
match event_type {
|
|
Type::ClientActions(actions) => {
|
|
for action in actions.actions {
|
|
if let Some(action_type) = action.action {
|
|
use warp_multi_agent_api::client_action::Action;
|
|
match action_type {
|
|
Action::AddMessagesToTask(add) => {
|
|
for msg in add.messages {
|
|
if let Some(msg_content) = msg.message {
|
|
use warp_multi_agent_api::message::Message;
|
|
match msg_content {
|
|
Message::AgentOutput(output) => {
|
|
text.push_str(&output.text);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Action::AppendToMessageContent(append) => {
|
|
if let Some(msg) = append.message {
|
|
if let Some(msg_content) = msg.message {
|
|
use warp_multi_agent_api::message::Message;
|
|
if let Message::AgentOutput(output) = msg_content {
|
|
text.push_str(&output.text);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Type::Finished(finished) => {
|
|
finished_reason = Some(format!("{:?}", finished.reason));
|
|
if let Some(meta) = finished.conversation_usage_metadata {
|
|
if let Some(usage) = meta.byok_token_usage.get("bedrock") {
|
|
total_tokens = usage.total_tokens;
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
StreamOutput {
|
|
text,
|
|
tool_calls,
|
|
finished_reason,
|
|
total_tokens,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_simple_text_response() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("Say exactly: Hello there, how are you?".into()),
|
|
}];
|
|
|
|
let output = collect_stream_output(&client, &model, messages, None, vec![]).await;
|
|
println!("[test] Text output: {:?}", output.text);
|
|
println!("[test] Finished reason: {:?}", output.finished_reason);
|
|
println!("[test] Total tokens: {}", output.total_tokens);
|
|
|
|
assert!(!output.text.is_empty(), "Expected non-empty text response");
|
|
assert!(
|
|
output.finished_reason.is_some(),
|
|
"Expected stream to finish"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_simple_with_system_prompt() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("What is your name?".into()),
|
|
}];
|
|
|
|
let output = collect_stream_output(
|
|
&client,
|
|
&model,
|
|
messages,
|
|
Some("You are a helpful assistant named Warp.".into()),
|
|
vec![],
|
|
)
|
|
.await;
|
|
println!("[test] Text output: {:?}", output.text);
|
|
|
|
assert!(!output.text.is_empty());
|
|
assert!(
|
|
output.text.to_lowercase().contains("warp"),
|
|
"Expected response to mention 'Warp', got: {}",
|
|
&output.text[..output.text.len().min(200)]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_call_round_trip() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let tools = vec![ToolDefinition {
|
|
name: "list_files".into(),
|
|
description: "List files in a directory".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Directory path" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
}];
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("List the files in the /project directory.".into()),
|
|
}];
|
|
|
|
let output = collect_stream_output(&client, &model, messages, None, tools).await;
|
|
println!(
|
|
"[test] Text: {:?}",
|
|
&output.text[..output.text.len().min(200)]
|
|
);
|
|
println!("[test] Tool calls: {:?}", output.tool_calls);
|
|
println!("[test] Finished: {:?}", output.finished_reason);
|
|
|
|
assert!(
|
|
!output.text.is_empty() || !output.tool_calls.is_empty(),
|
|
"Expected either text or a tool call"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_turn_with_tool_result() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let tools = vec![ToolDefinition {
|
|
name: "list_files".into(),
|
|
description: "List files in a directory".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Directory path" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
}];
|
|
|
|
let messages = vec![
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"List files in /project and tell me what you see.".into(),
|
|
),
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: "tool_1".into(),
|
|
name: "list_files".into(),
|
|
input: json!({"path": "/project"}),
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: "tool_1".into(),
|
|
content: "README.md\nsrc/\nCargo.toml\n.gitignore".into(),
|
|
is_error: false,
|
|
},
|
|
},
|
|
];
|
|
|
|
let output = collect_stream_output(&client, &model, messages, None, tools).await;
|
|
println!(
|
|
"[test] Text after tool result: {:?}",
|
|
&output.text[..output.text.len().min(300)]
|
|
);
|
|
|
|
assert!(
|
|
!output.text.is_empty(),
|
|
"Expected text response after tool result"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_turn_conversation_with_tools() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let tools = vec![
|
|
ToolDefinition {
|
|
name: "list_files".into(),
|
|
description: "List files in a directory".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "Directory path" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
},
|
|
ToolDefinition {
|
|
name: "read_file".into(),
|
|
description: "Read contents of a file".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "File path" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
},
|
|
ToolDefinition {
|
|
name: "run_command".into(),
|
|
description: "Run a shell command".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"command": { "type": "string", "description": "Shell command" }
|
|
},
|
|
"required": ["command"]
|
|
}),
|
|
},
|
|
];
|
|
|
|
let system = Some("You are a helpful coding assistant.".into());
|
|
|
|
// Turn 1: Ask model to inspect project
|
|
let turn1_messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"Inspect this project and tell me what it does. Start by listing files in /project."
|
|
.into(),
|
|
),
|
|
}];
|
|
|
|
let turn1 = collect_stream_output(
|
|
&client,
|
|
&model,
|
|
turn1_messages,
|
|
system.clone(),
|
|
tools.clone(),
|
|
)
|
|
.await;
|
|
println!(
|
|
"[test] Turn 1 text: {:?}",
|
|
&turn1.text[..turn1.text.len().min(200)]
|
|
);
|
|
println!("[test] Turn 1 tool_calls: {:?}", turn1.tool_calls);
|
|
|
|
// Turn 2: Provide tool result, continue
|
|
let turn2_messages = vec![
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"Inspect this project and tell me what it does. Start by listing files in /project."
|
|
.into(),
|
|
),
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: "tool_turn1".into(),
|
|
name: "list_files".into(),
|
|
input: json!({"path": "/project"}),
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: "tool_turn1".into(),
|
|
content: "README.md\nsrc/main.rs\nsrc/lib.rs\nCargo.toml\ntests/\n.gitignore".into(),
|
|
is_error: false,
|
|
},
|
|
},
|
|
];
|
|
|
|
let turn2 = collect_stream_output(
|
|
&client,
|
|
&model,
|
|
turn2_messages,
|
|
system.clone(),
|
|
tools.clone(),
|
|
)
|
|
.await;
|
|
println!(
|
|
"[test] Turn 2 text: {:?}",
|
|
&turn2.text[..turn2.text.len().min(200)]
|
|
);
|
|
println!("[test] Turn 2 tool_calls: {:?}", turn2.tool_calls);
|
|
|
|
// Turn 3: Provide README content and ask for summary
|
|
let turn3_messages = vec![
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"Inspect this project and tell me what it does. Start by listing files in /project."
|
|
.into(),
|
|
),
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: "tool_turn1".into(),
|
|
name: "list_files".into(),
|
|
input: json!({"path": "/project"}),
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: "tool_turn1".into(),
|
|
content: "README.md\nsrc/main.rs\nsrc/lib.rs\nCargo.toml\ntests/\n.gitignore".into(),
|
|
is_error: false,
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: "tool_turn2".into(),
|
|
name: "read_file".into(),
|
|
input: json!({"path": "/project/README.md"}),
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: "tool_turn2".into(),
|
|
content: "# My CLI Tool\n\nA Rust command-line tool for managing developer workflows.\n\n## Features\n- Task tracking\n- Git integration\n- Custom scripts\n\n## Usage\n```\ncargo run -- <command>\n```".into(),
|
|
is_error: false,
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"Based on what you've seen, give me a brief summary of this project. Do not use any tools."
|
|
.into(),
|
|
),
|
|
},
|
|
];
|
|
|
|
let turn3 = collect_stream_output(
|
|
&client,
|
|
&model,
|
|
turn3_messages,
|
|
system.clone(),
|
|
tools.clone(),
|
|
)
|
|
.await;
|
|
println!(
|
|
"[test] Turn 3 text: {:?}",
|
|
&turn3.text[..turn3.text.len().min(500)]
|
|
);
|
|
|
|
assert!(
|
|
!turn3.text.is_empty() || !turn3.tool_calls.is_empty(),
|
|
"Expected final summary or tool use after multi-turn conversation"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_error_recovery() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let tools = vec![ToolDefinition {
|
|
name: "read_file".into(),
|
|
description: "Read contents of a file".into(),
|
|
input_schema: json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"path": { "type": "string", "description": "File path" }
|
|
},
|
|
"required": ["path"]
|
|
}),
|
|
}];
|
|
|
|
let messages = vec![
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("Read the file /project/config.yaml".into()),
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: "tool_err".into(),
|
|
name: "read_file".into(),
|
|
input: json!({"path": "/project/config.yaml"}),
|
|
},
|
|
},
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: "tool_err".into(),
|
|
content: "Error: File not found: /project/config.yaml".into(),
|
|
is_error: true,
|
|
},
|
|
},
|
|
];
|
|
|
|
let output = collect_stream_output(&client, &model, messages, None, tools).await;
|
|
println!(
|
|
"[test] Error recovery text: {:?}",
|
|
&output.text[..output.text.len().min(300)]
|
|
);
|
|
|
|
assert!(
|
|
!output.text.is_empty(),
|
|
"Expected model to respond to tool error"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_arn_based_model() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let arn = std::env::var("BEDROCK_TEST_ARN").unwrap_or_else(|_| {
|
|
"arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/1tim45pgo320".into()
|
|
});
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("Reply with the single word: confirmed".into()),
|
|
}];
|
|
|
|
let output = collect_stream_output(&client, &arn, messages, None, vec![]).await;
|
|
println!("[test] ARN model text: {:?}", output.text);
|
|
|
|
assert!(
|
|
!output.text.is_empty(),
|
|
"Expected response from ARN-based model"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_reasoning_model_output() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text("What is 15 * 37? Show your reasoning step by step.".into()),
|
|
}];
|
|
|
|
let output = collect_stream_output(&client, &model, messages, None, vec![]).await;
|
|
println!(
|
|
"[test] Reasoning model text ({} chars): {:?}",
|
|
output.text.len(),
|
|
&output.text[..output.text.len().min(500)]
|
|
);
|
|
println!("[test] Finished: {:?}", output.finished_reason);
|
|
|
|
assert!(!output.text.is_empty(), "Expected reasoning output");
|
|
assert!(
|
|
output.text.contains("555"),
|
|
"Expected correct answer (555) in output, got: {}",
|
|
&output.text[..output.text.len().min(300)]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_all_tools_visible_to_model() {
|
|
let Some(config) = get_test_config() else {
|
|
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
|
|
return;
|
|
};
|
|
|
|
let client = BedrockClient::from_config(config)
|
|
.await
|
|
.expect("client creation");
|
|
let model = get_test_model();
|
|
|
|
let tools = super::request_translator::default_tool_definitions();
|
|
let tool_names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
|
|
|
|
println!("[test] Sending {} tools to Bedrock: {:?}", tools.len(), tool_names);
|
|
|
|
let messages = vec![ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(
|
|
"List every tool you have access to. Output ONLY the tool names, one per line, no descriptions, no formatting, no markdown."
|
|
.into(),
|
|
),
|
|
}];
|
|
|
|
let output = collect_stream_output(
|
|
&client,
|
|
&model,
|
|
messages,
|
|
Some("You are a helpful assistant. When asked about your tools, list them exactly as they appear in your tool configuration.".into()),
|
|
tools.clone(),
|
|
)
|
|
.await;
|
|
|
|
println!("[test] Model's tool list response:\n{}", output.text);
|
|
println!("[test] Total tokens: {}", output.total_tokens);
|
|
|
|
let response_lower = output.text.to_lowercase();
|
|
|
|
let mut missing_tools = Vec::new();
|
|
for tool in &tools {
|
|
if !response_lower.contains(&tool.name.to_lowercase()) {
|
|
missing_tools.push(&tool.name);
|
|
}
|
|
}
|
|
|
|
if !missing_tools.is_empty() {
|
|
println!("[test] WARNING: Model did not mention these tools: {:?}", missing_tools);
|
|
}
|
|
|
|
let expected_core_tools = [
|
|
"run_shell_command",
|
|
"read_files",
|
|
"apply_file_diffs",
|
|
"grep",
|
|
"file_glob",
|
|
"search_codebase",
|
|
"start_agent",
|
|
"ask_user_question",
|
|
];
|
|
|
|
let mut missing_core = Vec::new();
|
|
for name in &expected_core_tools {
|
|
if !response_lower.contains(name) {
|
|
missing_core.push(*name);
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
missing_core.is_empty(),
|
|
"Model failed to list these core tools: {:?}\n\nFull response:\n{}",
|
|
missing_core,
|
|
output.text
|
|
);
|
|
|
|
assert!(
|
|
tools.len() >= 17,
|
|
"Expected at least 17 tool definitions, got {}",
|
|
tools.len()
|
|
);
|
|
}
|