379 lines
15 KiB
Rust
379 lines
15 KiB
Rust
use warp_multi_agent_api as api;
|
|
|
|
use super::convert::{ConversationMessage, MessageContent, MessageRole, ToolDefinition};
|
|
|
|
pub fn extract_messages_from_request(request: &api::Request) -> Vec<ConversationMessage> {
|
|
let mut messages = Vec::new();
|
|
|
|
if let Some(task_context) = &request.task_context {
|
|
for task in &task_context.tasks {
|
|
for msg in &task.messages {
|
|
if let Some(converted) = convert_proto_message(msg) {
|
|
messages.push(converted);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(input) = &request.input {
|
|
if let Some(input_type) = &input.r#type {
|
|
#[allow(deprecated)]
|
|
match input_type {
|
|
api::request::input::Type::UserInputs(user_inputs) => {
|
|
for user_input in &user_inputs.inputs {
|
|
if let Some(input_variant) = &user_input.input {
|
|
match input_variant {
|
|
api::request::input::user_inputs::user_input::Input::UserQuery(
|
|
query,
|
|
) => {
|
|
if !query.query.is_empty() {
|
|
messages.push(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(query.query.clone()),
|
|
});
|
|
}
|
|
}
|
|
api::request::input::user_inputs::user_input::Input::ToolCallResult(
|
|
result,
|
|
) => {
|
|
let content = extract_tool_result_content(result);
|
|
if !result.tool_call_id.is_empty() {
|
|
messages.push(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: result.tool_call_id.clone(),
|
|
content,
|
|
is_error: false,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
api::request::input::Type::UserQuery(query) => {
|
|
if !query.query.is_empty() {
|
|
messages.push(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(query.query.clone()),
|
|
});
|
|
}
|
|
}
|
|
api::request::input::Type::ToolCallResult(result) => {
|
|
let content = extract_tool_result_content(result);
|
|
if !result.tool_call_id.is_empty() {
|
|
messages.push(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: result.tool_call_id.clone(),
|
|
content,
|
|
is_error: false,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
messages
|
|
}
|
|
|
|
pub fn extract_system_prompt(_request: &api::Request) -> Option<String> {
|
|
Some("You are a helpful AI coding assistant. You help users with software engineering tasks including writing code, debugging, and explaining concepts.".to_string())
|
|
}
|
|
|
|
pub fn extract_tools(request: &api::Request) -> Vec<ToolDefinition> {
|
|
let mut tools = Vec::new();
|
|
let mut seen_names = std::collections::HashSet::new();
|
|
|
|
if let Some(task_context) = &request.task_context {
|
|
for task in &task_context.tasks {
|
|
for msg in &task.messages {
|
|
if let Some(api::message::Message::ToolCall(tool_call)) = &msg.message {
|
|
let (name, _) = extract_tool_call_info(tool_call);
|
|
if name != "unknown_tool" && seen_names.insert(name.clone()) {
|
|
tools.push(tool_definition_for_name(&name));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(input) = &request.input {
|
|
if let Some(input_type) = &input.r#type {
|
|
#[allow(deprecated)]
|
|
match input_type {
|
|
api::request::input::Type::UserInputs(user_inputs) => {
|
|
for user_input in &user_inputs.inputs {
|
|
if let Some(api::request::input::user_inputs::user_input::Input::ToolCallResult(_)) = &user_input.input {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
api::request::input::Type::ToolCallResult(_) => {}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
if tools.is_empty() {
|
|
let messages = extract_messages_from_request(request);
|
|
let has_tool_content = messages.iter().any(|m| {
|
|
matches!(
|
|
m.content,
|
|
MessageContent::ToolUse { .. } | MessageContent::ToolResult { .. }
|
|
)
|
|
});
|
|
if has_tool_content {
|
|
tools = default_tool_definitions();
|
|
}
|
|
}
|
|
|
|
tools
|
|
}
|
|
|
|
fn tool_definition_for_name(name: &str) -> ToolDefinition {
|
|
match name {
|
|
"run_shell_command" => ToolDefinition {
|
|
name: "run_shell_command".to_string(),
|
|
description: "Execute a shell command and return its output.".to_string(),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"command": { "type": "string", "description": "The shell command to execute" }
|
|
},
|
|
"required": ["command"]
|
|
}),
|
|
},
|
|
"read_files" => ToolDefinition {
|
|
name: "read_files".to_string(),
|
|
description: "Read the contents of one or more files.".to_string(),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"files": { "type": "array", "items": { "type": "string" }, "description": "File paths to read" }
|
|
},
|
|
"required": ["files"]
|
|
}),
|
|
},
|
|
"apply_file_diffs" => ToolDefinition {
|
|
name: "apply_file_diffs".to_string(),
|
|
description: "Apply search/replace diffs to files.".to_string(),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"diffs": { "type": "array", "items": { "type": "object", "properties": { "file_path": { "type": "string" }, "search": { "type": "string" }, "replace": { "type": "string" } }, "required": ["file_path", "search", "replace"] } }
|
|
},
|
|
"required": ["diffs"]
|
|
}),
|
|
},
|
|
"grep" => ToolDefinition {
|
|
name: "grep".to_string(),
|
|
description: "Search for patterns in files using grep.".to_string(),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"queries": { "type": "array", "items": { "type": "string" }, "description": "Search patterns" },
|
|
"path": { "type": "string", "description": "Directory to search in" }
|
|
},
|
|
"required": ["queries"]
|
|
}),
|
|
},
|
|
"file_glob" => ToolDefinition {
|
|
name: "file_glob".to_string(),
|
|
description: "Find files matching glob patterns.".to_string(),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"patterns": { "type": "array", "items": { "type": "string" }, "description": "Glob patterns to match" }
|
|
},
|
|
"required": ["patterns"]
|
|
}),
|
|
},
|
|
_ => ToolDefinition {
|
|
name: name.to_string(),
|
|
description: format!("Tool: {}", name),
|
|
input_schema: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {}
|
|
}),
|
|
},
|
|
}
|
|
}
|
|
|
|
fn default_tool_definitions() -> Vec<ToolDefinition> {
|
|
vec![
|
|
tool_definition_for_name("run_shell_command"),
|
|
tool_definition_for_name("read_files"),
|
|
tool_definition_for_name("apply_file_diffs"),
|
|
tool_definition_for_name("grep"),
|
|
tool_definition_for_name("file_glob"),
|
|
]
|
|
}
|
|
|
|
fn convert_proto_message(msg: &api::Message) -> Option<ConversationMessage> {
|
|
let message_content = msg.message.as_ref()?;
|
|
|
|
match message_content {
|
|
api::message::Message::UserQuery(query) => Some(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::Text(query.query.clone()),
|
|
}),
|
|
api::message::Message::AgentOutput(output) => Some(ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::Text(output.text.clone()),
|
|
}),
|
|
api::message::Message::ToolCall(tool_call) => {
|
|
let (name, input) = extract_tool_call_info(tool_call);
|
|
Some(ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse {
|
|
tool_use_id: msg.id.clone(),
|
|
name,
|
|
input,
|
|
},
|
|
})
|
|
}
|
|
api::message::Message::ToolCallResult(result) => {
|
|
let content = format_tool_call_result(result);
|
|
Some(ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult {
|
|
tool_use_id: result.tool_call_id.clone(),
|
|
content,
|
|
is_error: false,
|
|
},
|
|
})
|
|
}
|
|
api::message::Message::AgentReasoning(_) => None,
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn extract_tool_call_info(tool_call: &api::message::ToolCall) -> (String, serde_json::Value) {
|
|
if let Some(tool) = &tool_call.tool {
|
|
match tool {
|
|
api::message::tool_call::Tool::RunShellCommand(cmd) => (
|
|
"run_shell_command".to_string(),
|
|
serde_json::json!({ "command": cmd.command }),
|
|
),
|
|
api::message::tool_call::Tool::ReadFiles(read) => (
|
|
"read_files".to_string(),
|
|
serde_json::json!({ "files": read.files.iter().map(|f| &f.name).collect::<Vec<_>>() }),
|
|
),
|
|
api::message::tool_call::Tool::ApplyFileDiffs(diffs) => (
|
|
"apply_file_diffs".to_string(),
|
|
serde_json::json!({ "diffs": diffs.diffs.iter().map(|d| {
|
|
serde_json::json!({
|
|
"file_path": d.file_path,
|
|
"search": d.search,
|
|
"replace": d.replace
|
|
})
|
|
}).collect::<Vec<_>>() }),
|
|
),
|
|
api::message::tool_call::Tool::Grep(grep) => (
|
|
"grep".to_string(),
|
|
serde_json::json!({ "queries": grep.queries, "path": grep.path }),
|
|
),
|
|
#[allow(deprecated)]
|
|
api::message::tool_call::Tool::FileGlob(glob) => (
|
|
"file_glob".to_string(),
|
|
serde_json::json!({ "patterns": glob.patterns }),
|
|
),
|
|
_ => ("unknown_tool".to_string(), serde_json::json!({})),
|
|
}
|
|
} else {
|
|
("unknown_tool".to_string(), serde_json::json!({}))
|
|
}
|
|
}
|
|
|
|
fn extract_tool_result_content(result: &api::request::input::ToolCallResult) -> String {
|
|
if let Some(result_type) = &result.result {
|
|
match result_type {
|
|
api::request::input::tool_call_result::Result::RunShellCommand(cmd_result) => {
|
|
match &cmd_result.result {
|
|
Some(
|
|
api::run_shell_command_result::Result::CommandFinished(finished),
|
|
) => finished.output.clone(),
|
|
Some(
|
|
api::run_shell_command_result::Result::LongRunningCommandSnapshot(
|
|
snapshot,
|
|
),
|
|
) => snapshot.output.clone(),
|
|
_ => "Command completed.".to_string(),
|
|
}
|
|
}
|
|
api::request::input::tool_call_result::Result::ReadFiles(read_result) => {
|
|
match &read_result.result {
|
|
Some(api::read_files_result::Result::TextFilesSuccess(success)) => success
|
|
.files
|
|
.iter()
|
|
.map(|f| format!("{}:\n{}", f.file_path, f.content))
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n"),
|
|
Some(api::read_files_result::Result::AnyFilesSuccess(success)) => success
|
|
.files
|
|
.iter()
|
|
.filter_map(|f| match &f.content {
|
|
Some(api::any_file_content::Content::TextContent(t)) => {
|
|
Some(format!("{}:\n{}", t.file_path, t.content))
|
|
}
|
|
_ => None,
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n"),
|
|
_ => "Failed to read files.".to_string(),
|
|
}
|
|
}
|
|
_ => "Tool completed successfully.".to_string(),
|
|
}
|
|
} else {
|
|
"Tool completed.".to_string()
|
|
}
|
|
}
|
|
|
|
fn format_tool_call_result(result: &api::message::ToolCallResult) -> String {
|
|
if let Some(result_type) = &result.result {
|
|
match result_type {
|
|
api::message::tool_call_result::Result::RunShellCommand(cmd_result) => {
|
|
match &cmd_result.result {
|
|
Some(
|
|
api::run_shell_command_result::Result::CommandFinished(finished),
|
|
) => {
|
|
format!(
|
|
"Exit code: {}\nOutput: {}",
|
|
finished.exit_code, finished.output
|
|
)
|
|
}
|
|
Some(
|
|
api::run_shell_command_result::Result::LongRunningCommandSnapshot(
|
|
snapshot,
|
|
),
|
|
) => {
|
|
format!("Output (running): {}", snapshot.output)
|
|
}
|
|
_ => "Command completed.".to_string(),
|
|
}
|
|
}
|
|
api::message::tool_call_result::Result::ReadFiles(read_result) => {
|
|
match &read_result.result {
|
|
Some(api::read_files_result::Result::TextFilesSuccess(success)) => success
|
|
.files
|
|
.iter()
|
|
.map(|f| format!("{}:\n{}", f.file_path, f.content))
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n"),
|
|
_ => "Read files completed.".to_string(),
|
|
}
|
|
}
|
|
_ => "Tool completed successfully.".to_string(),
|
|
}
|
|
} else {
|
|
"Tool completed.".to_string()
|
|
}
|
|
}
|