use dragen::{Agent, AgentConfig};
use littrs::{PyValue, ToolInfo};
use serde::{Deserialize, Serialize};
use std::env;
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ExaSearchRequest {
query: String,
num_results: u32,
#[serde(rename = "type")]
search_type: String,
contents: ExaContents,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ExaContents {
text: ExaTextConfig,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ExaTextConfig {
max_characters: u32,
}
#[derive(Deserialize)]
struct ExaSearchResponse {
results: Vec<ExaResult>,
}
#[derive(Deserialize)]
struct ExaResult {
title: String,
url: String,
text: Option<String>,
}
fn search_web(query: String, num_results: i64) -> PyValue {
let api_key = match env::var("EXA_API_KEY") {
Ok(key) => key,
Err(_) => return PyValue::Str("Error: EXA_API_KEY not set".to_string()),
};
let request = ExaSearchRequest {
query: query.clone(),
num_results: num_results.max(1).min(10) as u32,
search_type: "auto".to_string(),
contents: ExaContents {
text: ExaTextConfig { max_characters: 1500 },
},
};
let response = ureq::post("https://api.exa.ai/search")
.header("x-api-key", &api_key)
.header("Content-Type", "application/json")
.send_json(&request);
match response {
Ok(mut resp) => match resp.body_mut().read_json::<ExaSearchResponse>() {
Ok(data) => {
let results: Vec<PyValue> = data
.results
.into_iter()
.map(|r| {
PyValue::Dict(vec![
(PyValue::Str("title".to_string()), PyValue::Str(r.title)),
(PyValue::Str("url".to_string()), PyValue::Str(r.url)),
(
PyValue::Str("snippet".to_string()),
PyValue::Str(r.text.unwrap_or_default()),
),
])
})
.collect();
PyValue::List(results)
}
Err(e) => PyValue::Str(format!("Error parsing response: {}", e)),
},
Err(ureq::Error::StatusCode(code)) => PyValue::Str(format!("HTTP error {}", code)),
Err(e) => PyValue::Str(format!("Request error: {:?}", e)),
}
}
fn register_search_tool(agent: &mut Agent) {
let search_info = ToolInfo::new("search", "Search the web for information")
.arg("query", "str", "The search query")
.arg_opt("num_results", "int", "Number of results (1-10, default 5)")
.returns("list");
agent.register_tool(search_info, |args| {
let query = args
.get(0)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let num_results = args.get(1).and_then(|v| v.as_int()).unwrap_or(5);
search_web(query, num_results)
});
}
const PLANNER_SYSTEM: &str = r#"You are a research planner. Your job is to:
1. Do a broad search on the given topic
2. Identify 3-4 key sections that should be researched in depth
3. Return a structured plan as a dictionary
FORMAT: Write Python code in ```python blocks.
Process:
1. Search for broad overview of the topic
2. Based on results, identify key themes/sections
3. Call finish() with a dictionary of sections
Example output format:
```python
results = search("topic overview", 5)
print(results)
```
Then after seeing results:
```python
finish({
"section1_title": "Brief description of what to research",
"section2_title": "Brief description of what to research",
"section3_title": "Brief description of what to research"
})
```
Keep section titles concise (3-6 words). Descriptions should guide the research agent.
Do exactly 1-2 searches, then call finish() with your sections dictionary."#;
const RESEARCHER_SYSTEM: &str = r#"You are a research specialist. You will research a specific section of a larger report.
CRITICAL: You MUST do exactly 2 searches before calling finish(). Do NOT skip searches.
FORMAT: Write Python code in ```python blocks. ONE action per code block.
REQUIRED STEPS (follow exactly):
STEP 1 - First search:
```python
results = search("your first query about the section", 5)
print(results)
```
STEP 2 - Save note from first search:
```python
note("Key facts and figures from search results")
```
STEP 3 - Second search (different angle):
```python
results = search("different aspect or more specific detail", 5)
print(results)
```
STEP 4 - Save note and finish with structured output:
```python
note("Additional facts from second search")
finish({
"content": "Your 2-3 paragraph section content here. Include specific facts and figures but do NOT include URLs in the content.",
"sources": [
"https://example.com/source1 - Description of source",
"https://example.com/source2 - Description of source"
]
})
```
IMPORTANT:
- ONE action per code block
- Make second search query DIFFERENT from first
- finish() must be a dict with "content" and "sources" keys
- Put URLs ONLY in sources list, not in content
- Do NOT repeat information from previous sections"#;
fn create_planner_agent() -> Agent {
let config = AgentConfig::new("llama-3.3-70b-versatile")
.max_iterations(8)
.system(PLANNER_SYSTEM);
let mut agent = Agent::new(config);
register_search_tool(&mut agent);
let finish_info = ToolInfo::new("finish", "Return the research sections as a structured plan")
.arg("sections", "dict", "Dictionary mapping section titles to research descriptions")
.returns("dict");
agent.register_finish(finish_info, |args| {
args.get(0).cloned().unwrap_or(PyValue::None)
});
agent
}
fn create_researcher_agent() -> Agent {
let config = AgentConfig::new("llama-3.3-70b-versatile")
.max_iterations(12)
.system(RESEARCHER_SYSTEM);
let mut agent = Agent::new(config);
register_search_tool(&mut agent);
let notes_info = ToolInfo::new("note", "Save a research note")
.arg("content", "str", "The note content")
.returns("str");
let notes: std::sync::Arc<std::sync::Mutex<Vec<String>>> =
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
agent.register_tool(notes_info, move |args| {
let content = args
.get(0)
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
if let Ok(mut n) = notes.lock() {
n.push(content);
PyValue::Str(format!("Note saved. Total: {}", n.len()))
} else {
PyValue::Str("Error saving note".to_string())
}
});
let finish_info = ToolInfo::new("finish", "Return section content with sources")
.arg("result", "dict", "Dict with 'content' (str) and 'sources' (list of str)")
.returns("dict");
agent.register_finish(finish_info, |args| {
args.get(0).cloned().unwrap_or(PyValue::None)
});
agent
}
struct SectionResult {
title: String,
content: String,
sources: Vec<String>,
}
fn extract_section_from_pyvalue(value: &PyValue) -> (String, Vec<String>) {
match value {
PyValue::Dict(pairs) => {
let mut content = String::new();
let mut sources = Vec::new();
for (key, val) in pairs {
match key.as_str() {
Some("content") => {
if let PyValue::Str(s) = val {
content = s.clone();
}
}
Some("sources") => {
if let PyValue::List(items) = val {
for item in items {
if let PyValue::Str(s) = item {
sources.push(s.clone());
}
}
}
}
_ => {}
}
}
(content, sources)
}
PyValue::Str(s) => {
(s.clone(), vec![])
}
_ => (String::new(), vec![])
}
}
fn print_separator(title: &str) {
println!("\n{}", "═".repeat(70));
println!(" {}", title);
println!("{}\n", "═".repeat(70));
}
fn print_subseparator(title: &str) {
println!("\n{}", "─".repeat(60));
println!(" {}", title);
println!("{}\n", "─".repeat(60));
}
fn extract_code_block(text: &str) -> Option<String> {
if let Some(start) = text.find("```") {
let after_start = &text[start + 3..];
if let Some(newline) = after_start.find('\n') {
let code_start = &after_start[newline + 1..];
if let Some(end) = code_start.find("```") {
return Some(code_start[..end].trim().to_string());
}
}
}
None
}
fn print_agent_step(agent_name: &str, step_num: &mut usize, content: &str) {
if let Some(code) = extract_code_block(content) {
*step_num += 1;
if code.contains("search(") {
if let Some(query_start) = code.find('"') {
let after_quote = &code[query_start + 1..];
if let Some(query_end) = after_quote.find('"') {
let query = &after_quote[..query_end];
println!(" [{}] Step {}: 🔍 SEARCH \"{}\"", agent_name, step_num, query);
return;
}
}
if let Some(query_start) = code.find('\'') {
let after_quote = &code[query_start + 1..];
if let Some(query_end) = after_quote.find('\'') {
let query = &after_quote[..query_end];
println!(" [{}] Step {}: 🔍 SEARCH \"{}\"", agent_name, step_num, query);
return;
}
}
}
if code.contains("note(") {
println!(" [{}] Step {}: 📝 SAVING NOTE", agent_name, step_num);
return;
}
if code.contains("finish(") {
println!(" [{}] Step {}: ✅ FINISHING", agent_name, step_num);
return;
}
let truncated = if code.len() > 50 { &code[..50] } else { &code };
println!(" [{}] Step {}: ⚡ {}", agent_name, step_num, truncated);
}
}
fn parse_sections(result: &str) -> Vec<(String, String)> {
let mut sections = Vec::new();
let mut remaining = result;
while let Some(key_start) = remaining.find('"') {
let after_key_start = &remaining[key_start + 1..];
if let Some(key_end) = after_key_start.find('"') {
let key = &after_key_start[..key_end];
let after_key = &after_key_start[key_end + 1..];
if let Some(val_start) = after_key.find('"') {
let after_val_start = &after_key[val_start + 1..];
if let Some(val_end) = after_val_start.find('"') {
let value = &after_val_start[..val_end];
sections.push((key.to_string(), value.to_string()));
remaining = &after_val_start[val_end + 1..];
continue;
}
}
}
break;
}
if sections.is_empty() {
for line in result.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
if let Some(colon_pos) = line.find(':') {
let title = line[..colon_pos].trim();
let desc = line[colon_pos + 1..].trim();
if !title.is_empty() && !desc.is_empty() && title.len() < 50 {
sections.push((title.to_string(), desc.to_string()));
}
}
}
}
sections
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
if env::var("EXA_API_KEY").is_err() {
eprintln!("Error: EXA_API_KEY environment variable not set");
std::process::exit(1);
}
let topic = env::args().nth(1).unwrap_or_else(|| {
"The current state and future of quantum computing".to_string()
});
print_separator(&format!("MULTI-AGENT DEEP RESEARCH: {}", topic));
print_separator("PHASE 1: PLANNER AGENT");
println!("The planner will research the topic and identify key sections to explore.\n");
let mut planner = create_planner_agent();
let planner_task = format!(
"Research this topic and identify 3-4 key sections for a comprehensive report: {}",
topic
);
println!("📋 Task: {}\n", planner_task);
let planner_result = match planner.run::<String>(&planner_task).await {
Ok(result) => {
let mut step_num = 0;
for msg in planner.messages().iter().skip(1) {
let role = format!("{:?}", msg.role);
if role == "Assistant" {
print_agent_step("Planner", &mut step_num, &msg.content);
}
}
println!("\n📊 Planner Output:");
println!("{}", result);
result
}
Err(e) => {
eprintln!("Planner error: {}", e);
let mut step_num = 0;
for msg in planner.messages().iter().skip(1) {
let role = format!("{:?}", msg.role);
if role == "Assistant" {
print_agent_step("Planner", &mut step_num, &msg.content);
}
}
return Err(e.into());
}
};
let sections = parse_sections(&planner_result);
if sections.is_empty() {
eprintln!("Error: Could not parse sections from planner output");
eprintln!("Raw output: {}", planner_result);
return Ok(());
}
println!("\n📑 Identified {} sections to research:", sections.len());
for (i, (title, desc)) in sections.iter().enumerate() {
println!(" {}. {} - {}", i + 1, title, desc);
}
print_separator("PHASE 2: RESEARCH AGENTS");
let mut section_results: Vec<SectionResult> = Vec::new();
for (i, (section_title, section_desc)) in sections.iter().enumerate() {
print_subseparator(&format!("Section {}/{}: {}", i + 1, sections.len(), section_title));
let mut researcher = create_researcher_agent();
let previous_context = if let Some(prev) = section_results.last() {
format!(
"\n\nPrevious section for context:\n## {}\n{}\n\nBuild on this context but focus on your assigned section.",
prev.title, prev.content
)
} else {
String::new()
};
let research_task = format!(
"Topic: {}\n\nSection to research: {}\nGuidance: {}{}\n\nResearch this section thoroughly and provide detailed, well-sourced content.",
topic, section_title, section_desc, previous_context
);
println!("📋 Research Task: {} - {}\n", section_title, section_desc);
match researcher.run::<String>(&research_task).await {
Ok(result) => {
let mut step_num = 0;
for msg in researcher.messages().iter().skip(1) {
let role = format!("{:?}", msg.role);
if role == "Assistant" {
print_agent_step(&format!("Researcher-{}", i + 1), &mut step_num, &msg.content);
}
}
let (content, sources) = if let Some(value) = researcher.finish_value() {
extract_section_from_pyvalue(&value)
} else {
(result, vec![])
};
println!("\n📄 Section Content Preview:");
let preview: String = content.chars().take(300).collect();
println!("{}...\n", preview);
if !sources.is_empty() {
println!("📚 Sources: {}", sources.len());
}
section_results.push(SectionResult {
title: section_title.clone(),
content,
sources,
});
}
Err(e) => {
eprintln!("Research error for section '{}': {}", section_title, e);
let mut step_num = 0;
for msg in researcher.messages().iter().skip(1) {
let role = format!("{:?}", msg.role);
if role == "Assistant" {
print_agent_step(&format!("Researcher-{}", i + 1), &mut step_num, &msg.content);
}
}
section_results.push(SectionResult {
title: section_title.clone(),
content: format!("[Research incomplete: {}]", e),
sources: vec![],
});
}
}
}
print_separator("FINAL RESEARCH REPORT");
println!("# {}\n", topic);
for section in §ion_results {
println!("## {}\n", section.title);
println!("{}\n", section.content);
println!("{}\n", "─".repeat(50));
}
let all_sources: Vec<&String> = section_results
.iter()
.flat_map(|s| &s.sources)
.collect();
if !all_sources.is_empty() {
println!("\n## Sources\n");
for (i, source) in all_sources.iter().enumerate() {
println!("{}. {}", i + 1, source);
}
println!();
}
print_separator("RESEARCH COMPLETE");
println!("Generated {} sections with {} sources for topic: {}", section_results.len(), all_sources.len(), topic);
Ok(())
}