use std::collections::HashSet;
use std::fmt;
use ares_llm::{client::TokenUsage, LLMClient};
use ares_types::types::{Result, Source};
use serde::{Deserialize, Serialize};
use tokio::task::JoinSet;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PhaseExecutionMode {
Sequential,
Parallel,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResearchPhase {
pub id: String,
pub query: String,
pub mode: PhaseExecutionMode,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub depends_on: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResearchPlan {
pub root_query: String,
pub phases: Vec<ResearchPhase>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResearchFinding {
pub phase_id: String,
pub content: String,
pub relevance_score: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResearchResult {
pub findings: Vec<ResearchFinding>,
pub sources: Vec<Source>,
pub failed_phase_ids: Vec<String>,
pub partial: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceCandidate {
pub source: Source,
pub published_epoch_secs: Option<u64>,
pub authority_score: f32,
pub content_match_score: f32,
}
impl fmt::Display for ResearchPlan {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "ResearchPlan: {}", self.root_query)?;
for phase in &self.phases {
let deps = if phase.depends_on.is_empty() {
String::new()
} else {
format!(" (after {})", phase.depends_on.join(", "))
};
writeln!(f, " [{} / {:?}] {}{}", phase.id, phase.mode, phase.query, deps)?;
}
Ok(())
}
}
impl fmt::Display for ResearchPhase {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "ResearchPhase {} [{:?}]: {}", self.id, self.mode, self.query)
}
}
impl fmt::Display for ResearchResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "ResearchResult: {} finding(s), {} source(s), partial={}", self.findings.len(), self.sources.len(), self.partial)?;
if !self.failed_phase_ids.is_empty() {
writeln!(f, " failed phases: {}", self.failed_phase_ids.join(", "))?;
}
for finding in &self.findings {
writeln!(f, " - {} (relevance {:.2}): {}", finding.phase_id, finding.relevance_score, finding.content)?;
}
Ok(())
}
}
pub fn plan_research_phases(root_query: &str, sub_queries: &[String], parallel: bool) -> ResearchPlan {
let mode = if parallel { PhaseExecutionMode::Parallel } else { PhaseExecutionMode::Sequential };
let phases = sub_queries.iter().enumerate().map(|(index, query)| {
let id = format!("phase-{}", index + 1);
let depends_on = if parallel || index == 0 { Vec::new() } else { vec![format!("phase-{index}")] };
ResearchPhase { id, query: query.clone(), mode, depends_on }
}).collect();
ResearchPlan { root_query: root_query.to_string(), phases }
}
pub fn execute_phase<F>(phase: &ResearchPhase, executor: &mut F) -> std::result::Result<ResearchFinding, String>
where
F: FnMut(&ResearchPhase) -> std::result::Result<String, String>,
{
let content = executor(phase)?;
Ok(ResearchFinding {
phase_id: phase.id.clone(),
relevance_score: score_content_relevance(&phase.query, &content),
content,
})
}
pub fn run_research_plan<F>(plan: &ResearchPlan, executor: &mut F) -> ResearchResult
where
F: FnMut(&ResearchPhase) -> std::result::Result<String, String>,
{
let mut completed = HashSet::new();
let mut raw_findings = Vec::new();
let mut failed_phase_ids = Vec::new();
loop {
let ready: Vec<&ResearchPhase> = plan.phases.iter().filter(|phase| {
!completed.contains(&phase.id)
&& !failed_phase_ids.contains(&phase.id)
&& phase.depends_on.iter().all(|dep| completed.contains(dep))
}).collect();
if ready.is_empty() { break; }
let parallel = ready[0].mode == PhaseExecutionMode::Parallel;
let batch: Vec<&ResearchPhase> = if parallel { ready } else { vec![ready[0]] };
for phase in batch {
match execute_phase(phase, executor) {
Ok(finding) => { completed.insert(phase.id.clone()); raw_findings.push(finding); }
Err(_) => { failed_phase_ids.push(phase.id.clone()); }
}
}
}
let findings = aggregate_findings(&raw_findings);
let candidates = findings_to_source_candidates(&findings, &plan.root_query);
let sources = rank_sources(&candidates, &plan.root_query, current_epoch_secs());
let partial = !failed_phase_ids.is_empty() && !findings.is_empty();
ResearchResult {
findings,
sources,
failed_phase_ids,
partial,
}
}
pub fn aggregate_findings(findings: &[ResearchFinding]) -> Vec<ResearchFinding> {
use std::collections::HashMap;
let mut by_content: HashMap<String, ResearchFinding> = HashMap::new();
for finding in findings {
let key = normalize_content(&finding.content);
if key.is_empty() {
continue;
}
by_content
.entry(key)
.and_modify(|existing| {
if finding.relevance_score > existing.relevance_score {
*existing = finding.clone();
}
})
.or_insert_with(|| finding.clone());
}
let mut unique: Vec<_> = by_content.into_values().collect();
unique.sort_by(|left, right| {
right
.relevance_score
.partial_cmp(&left.relevance_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
unique
}
pub fn rank_sources(candidates: &[SourceCandidate], query: &str, now_epoch_secs: u64) -> Vec<Source> {
let mut scored: Vec<(f32, &SourceCandidate)> = candidates.iter().map(|candidate| {
let recency = recency_score(candidate.published_epoch_secs, now_epoch_secs);
let content = candidate.content_match_score.max(score_content_relevance(query, &candidate.source.title));
let composite = 0.45 * candidate.authority_score + 0.30 * recency + 0.25 * content;
(composite, candidate)
}).collect();
scored.sort_by(|(left, _), (right, _)| right.partial_cmp(left).unwrap_or(std::cmp::Ordering::Equal));
scored.into_iter().map(|(score, candidate)| {
let mut source = candidate.source.clone();
source.relevance_score = score.clamp(0.0, 1.0);
source
}).collect()
}
pub fn score_content_relevance(query: &str, content: &str) -> f32 {
let query_tokens = tokenize(query);
if query_tokens.is_empty() { return 0.0; }
let content_tokens = tokenize(content);
let overlap = query_tokens.iter().filter(|token| content_tokens.contains(*token)).count();
(overlap as f32 / query_tokens.len() as f32).clamp(0.0, 1.0)
}
fn recency_score(published_epoch_secs: Option<u64>, now_epoch_secs: u64) -> f32 {
published_epoch_secs.map(|published| {
let age_days = now_epoch_secs.saturating_sub(published) / 86_400;
(1.0 - (age_days as f32 / 365.0)).clamp(0.0, 1.0)
}).unwrap_or(0.5)
}
fn score_authority(url: Option<&str>) -> f32 {
let Some(url) = url else { return 0.3; };
let lower = url.to_ascii_lowercase();
if lower.contains(".gov") || lower.contains(".edu") { 1.0 }
else if lower.contains("arxiv.org") || lower.contains("doi.org") || lower.contains(".org") { 0.85 }
else { 0.55 }
}
fn normalize_content(content: &str) -> String {
content.split_whitespace().collect::<Vec<_>>().join(" ").to_ascii_lowercase()
}
fn tokenize(text: &str) -> HashSet<String> {
text.split(|c: char| !c.is_alphanumeric()).filter(|token| token.len() >= 3).map(|token| token.to_ascii_lowercase()).collect()
}
fn extract_url(text: &str) -> Option<String> {
text.split_whitespace().find(|word| word.starts_with("http://") || word.starts_with("https://")).map(str::to_string)
}
fn extract_published_epoch(text: &str) -> Option<u64> {
text.split_whitespace().find_map(|word| word.strip_prefix("published:").and_then(|value| value.parse().ok()))
}
fn findings_to_source_candidates(findings: &[ResearchFinding], root_query: &str) -> Vec<SourceCandidate> {
findings.iter().enumerate().map(|(index, finding)| {
let url = extract_url(&finding.content);
let authority_score = score_authority(url.as_deref());
let content_match_score = score_content_relevance(root_query, &finding.content).max(finding.relevance_score);
SourceCandidate {
source: Source { title: format!("Research Finding {}", index + 1), url, relevance_score: content_match_score },
published_epoch_secs: extract_published_epoch(&finding.content),
authority_score,
content_match_score,
}
}).collect()
}
fn current_epoch_secs() -> u64 {
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|duration| duration.as_secs()).unwrap_or(0)
}
#[derive(Debug, Clone, Default)]
pub struct ResearchUsage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl ResearchUsage {
fn add(&mut self, usage: Option<&TokenUsage>) {
if let Some(usage) = usage {
self.input_tokens += usage.prompt_tokens;
self.output_tokens += usage.completion_tokens;
}
}
}
pub struct ResearchCoordinator {
llm: Box<dyn LLMClient>,
depth: u8,
max_iterations: u8,
}
impl ResearchCoordinator {
pub fn new(llm: Box<dyn LLMClient>, depth: u8, max_iterations: u8) -> Self {
Self {
llm,
depth,
max_iterations,
}
}
pub async fn research(&self, query: &str) -> Result<(String, Vec<Source>)> {
let (synthesis, sources, _) = self.research_with_usage(query).await?;
Ok((synthesis, sources))
}
pub async fn research_with_usage(
&self,
query: &str,
) -> Result<(String, Vec<Source>, ResearchUsage)> {
let mut all_findings = Vec::new();
let mut usage = ResearchUsage::default();
let (questions, question_usage) = self.generate_research_questions(query).await?;
usage.add(question_usage.as_ref());
for iteration in 0..self.max_iterations {
tracing::info!(
"Research iteration {}/{}",
iteration + 1,
self.max_iterations
);
let findings = self.parallel_research(&questions).await?;
all_findings.extend(findings);
if all_findings.len() >= (self.depth as usize * 3) {
break;
}
if iteration < self.max_iterations - 1 {
let (follow_ups, followup_usage) = self
.generate_followup_questions(query, &all_findings)
.await?;
usage.add(followup_usage.as_ref());
if follow_ups.is_empty() {
break;
}
}
}
let (synthesis, synthesis_usage) = self.synthesize_findings(query, &all_findings).await?;
usage.add(synthesis_usage.as_ref());
let all_sources = self.extract_sources(query, &all_findings);
Ok((synthesis, all_sources, usage))
}
async fn generate_research_questions(
&self,
query: &str,
) -> Result<(Vec<String>, Option<TokenUsage>)> {
let prompt = format!(
r#"Generate {} focused research questions to comprehensively answer: {}
Return only the questions, one per line, numbered 1-{}.
Example:
1. [QUESTION 1]
2. [QUESTION 2]
3. [QUESTION 3]
..."#,
self.depth, query, self.depth
);
let response = self
.llm
.generate_with_history(&[("user".to_string(), prompt)])
.await?;
let questions = response
.content
.lines()
.filter(|line| !line.trim().is_empty())
.map(|line| {
line.trim()
.trim_start_matches(|c: char| c.is_numeric() || c == '.' || c == ')')
.trim()
.to_string()
})
.collect();
Ok((questions, response.usage))
}
async fn parallel_research(&self, questions: &[String]) -> Result<Vec<String>> {
let mut set = JoinSet::new();
for question in questions.iter().take(self.depth as usize) {
let question = question.clone();
let _llm_clone = self.llm.model_name().to_string();
set.spawn(async move {
format!("Research findings for: {}", question)
});
}
let mut results = Vec::new();
while let Some(res) = set.join_next().await {
if let Ok(finding) = res {
results.push(finding);
}
}
Ok(results)
}
async fn generate_followup_questions(
&self,
_original_query: &str,
findings: &[String],
) -> Result<(Vec<String>, Option<TokenUsage>)> {
if findings.is_empty() {
return Ok((vec![], None));
}
let prompt = format!(
r#"Based on these findings:
{}
Generate 2-3 follow-up research questions.
ONLY output the questions and nothing else, like this:
<question1>
<question2>
<question3>
"#,
findings.join("\n")
);
let response = self
.llm
.generate_with_history(&[("user".to_string(), prompt)])
.await?;
let questions = response
.content
.lines()
.filter(|line| !line.trim().is_empty())
.take(3)
.map(|s| s.to_string())
.collect();
Ok((questions, response.usage))
}
async fn synthesize_findings(
&self,
query: &str,
findings: &[String],
) -> Result<(String, Option<TokenUsage>)> {
let prompt = format!(
r#"Original query: {}
Research findings:
{}
Synthesize these findings into a comprehensive, well-structured answer. Include:
1. Direct answer to the question
2. Key insights
3. Supporting evidence
4. Caveats or limitations if any
Provide a clear, professional response."#,
query,
findings.join("\n\n")
);
let response = self
.llm
.generate_with_history(&[("user".to_string(), prompt)])
.await?;
Ok((response.content, response.usage))
}
fn extract_sources(&self, query: &str, findings: &[String]) -> Vec<Source> {
let research_findings: Vec<ResearchFinding> = findings
.iter()
.enumerate()
.map(|(index, finding)| ResearchFinding {
phase_id: format!("finding-{}", index + 1),
content: finding.clone(),
relevance_score: score_content_relevance(query, finding),
})
.collect();
let aggregated = aggregate_findings(&research_findings);
let candidates = findings_to_source_candidates(&aggregated, query);
rank_sources(&candidates, query, current_epoch_secs())
}
}
#[cfg(test)]
mod tests {
use super::*;
use ares_llm::{client::TokenUsage, LLMClient, LLMResponse};
use ares_types::types::{AppError, ToolDefinition};
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
struct ResearchMockLlm {
calls: Arc<AtomicUsize>,
fail: bool,
}
impl ResearchMockLlm {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
fail: false,
}
}
fn failing() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
fail: true,
}
}
}
#[async_trait]
impl LLMClient for ResearchMockLlm {
fn model_name(&self) -> &str {
"research-mock"
}
async fn generate(&self, _: &str) -> Result<String> {
Ok(String::new())
}
async fn generate_with_system(&self, _: &str, _: &str) -> Result<String> {
Ok(String::new())
}
async fn generate_with_history(
&self,
messages: &[(String, String)],
) -> Result<LLMResponse> {
if self.fail {
return Err(AppError::Internal("mock llm failure".into()));
}
let _n = self.calls.fetch_add(1, Ordering::SeqCst);
let prompt = messages
.last()
.map(|(_, content)| content.as_str())
.unwrap_or_default();
let (content, usage) = if prompt.contains("follow-up research questions") {
(
"What is the regulatory timeline?\nWhat are adoption barriers?".to_string(),
TokenUsage::new(30, 12),
)
} else if prompt.contains("Synthesize these findings") {
(
"Comprehensive synthesized answer.".to_string(),
TokenUsage::new(50, 25),
)
} else {
(
"1. What is the core technology?\n2. Who are the main vendors?".to_string(),
TokenUsage::new(20, 8),
)
};
Ok(LLMResponse {
content,
tool_calls: vec![],
finish_reason: "stop".to_string(),
usage: Some(usage),
})
}
async fn generate_with_tools(
&self,
_: &str,
_: &[ToolDefinition],
) -> Result<LLMResponse> {
Ok(LLMResponse {
content: String::new(),
tool_calls: vec![],
finish_reason: "stop".to_string(),
usage: None,
})
}
async fn generate_with_tools_and_history(
&self,
_: &[ares_llm::coordinator::ConversationMessage],
_: &[ToolDefinition],
) -> Result<LLMResponse> {
Ok(LLMResponse {
content: String::new(),
tool_calls: vec![],
finish_reason: "stop".to_string(),
usage: None,
})
}
async fn stream(
&self,
_: &str,
) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
Ok(Box::new(futures::stream::empty()))
}
async fn stream_with_system(
&self,
_: &str,
_: &str,
) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
Ok(Box::new(futures::stream::empty()))
}
async fn stream_with_history(
&self,
_: &[(String, String)],
) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
Ok(Box::new(futures::stream::empty()))
}
}
fn sample_sub_queries() -> Vec<String> {
vec![
"What is the core technology?".to_string(),
"Who are the main vendors?".to_string(),
"What is the market size?".to_string(),
]
}
#[test]
fn research_plan_serde_roundtrip() {
let plan = plan_research_phases("quantum networking", &sample_sub_queries(), true);
let json = serde_json::to_string(&plan).expect("serialize plan");
let parsed: ResearchPlan = serde_json::from_str(&json).expect("deserialize plan");
assert_eq!(parsed, plan);
}
#[test]
fn research_phase_serde_roundtrip() {
let phase = ResearchPhase {
id: "phase-1".into(),
query: "vendors".into(),
mode: PhaseExecutionMode::Sequential,
depends_on: vec![],
};
let json = serde_json::to_string(&phase).expect("serialize phase");
let parsed: ResearchPhase = serde_json::from_str(&json).expect("deserialize phase");
assert_eq!(parsed, phase);
}
#[test]
fn research_result_serde_roundtrip() {
let result = ResearchResult {
findings: vec![ResearchFinding {
phase_id: "phase-1".into(),
content: "finding".into(),
relevance_score: 0.9,
}],
sources: vec![Source {
title: "Paper".into(),
url: Some("https://example.edu/paper".into()),
relevance_score: 0.9,
}],
failed_phase_ids: vec!["phase-2".into()],
partial: true,
};
let json = serde_json::to_string(&result).expect("serialize result");
let parsed: ResearchResult = serde_json::from_str(&json).expect("deserialize result");
assert_eq!(parsed.findings, result.findings);
assert_eq!(parsed.failed_phase_ids, result.failed_phase_ids);
assert_eq!(parsed.partial, result.partial);
assert_eq!(parsed.sources.len(), result.sources.len());
}
#[test]
fn phase_execution_mode_serde_roundtrip() {
for mode in [PhaseExecutionMode::Sequential, PhaseExecutionMode::Parallel] {
let json = serde_json::to_string(&mode).expect("serialize mode");
let parsed: PhaseExecutionMode = serde_json::from_str(&json).expect("deserialize mode");
assert_eq!(parsed, mode);
}
}
#[test]
fn plan_research_phases_parallel_has_no_dependencies() {
let plan = plan_research_phases("root", &sample_sub_queries(), true);
assert!(plan.phases.iter().all(|phase| phase.depends_on.is_empty()));
assert!(plan.phases.iter().all(|phase| phase.mode == PhaseExecutionMode::Parallel));
}
#[test]
fn plan_research_phases_sequential_builds_dependency_chain() {
let plan = plan_research_phases("root", &sample_sub_queries(), false);
assert_eq!(plan.phases[0].depends_on, Vec::<String>::new());
assert_eq!(plan.phases[1].depends_on, vec!["phase-1".to_string()]);
assert_eq!(plan.phases[2].depends_on, vec!["phase-2".to_string()]);
assert!(plan.phases.iter().all(|phase| phase.mode == PhaseExecutionMode::Sequential));
}
#[test]
fn plan_research_phases_preserves_sub_queries() {
let queries = sample_sub_queries();
let plan = plan_research_phases("root", &queries, true);
let planned: Vec<_> = plan.phases.iter().map(|phase| phase.query.clone()).collect();
assert_eq!(planned, queries);
}
#[test]
fn execute_phase_success_scores_relevance() {
let phase = ResearchPhase {
id: "phase-1".into(),
query: "quantum error correction".into(),
mode: PhaseExecutionMode::Parallel,
depends_on: vec![],
};
let mut executor = |_: &ResearchPhase| {
Ok("quantum error correction advances published:1700000000 https://arxiv.org/abs/123".into())
};
let finding = execute_phase(&phase, &mut executor).expect("phase should succeed");
assert_eq!(finding.phase_id, "phase-1");
assert!(finding.relevance_score > 0.5);
}
#[test]
fn execute_phase_failure_returns_error() {
let phase = ResearchPhase {
id: "phase-1".into(),
query: "vendors".into(),
mode: PhaseExecutionMode::Parallel,
depends_on: vec![],
};
let mut executor = |_: &ResearchPhase| Err("upstream search unavailable".to_string());
let err = execute_phase(&phase, &mut executor).expect_err("phase should fail");
assert!(err.contains("unavailable"));
}
#[test]
fn aggregate_findings_deduplicates_normalized_content() {
let findings = vec![
ResearchFinding { phase_id: "a".into(), content: "Same content here".into(), relevance_score: 0.4 },
ResearchFinding { phase_id: "b".into(), content: "same content here".into(), relevance_score: 0.9 },
];
let aggregated = aggregate_findings(&findings);
assert_eq!(aggregated.len(), 1);
assert_eq!(aggregated[0].phase_id, "b");
}
#[test]
fn aggregate_findings_sorts_by_relevance_descending() {
let findings = vec![
ResearchFinding { phase_id: "low".into(), content: "alpha".into(), relevance_score: 0.2 },
ResearchFinding { phase_id: "high".into(), content: "beta".into(), relevance_score: 0.95 },
];
let aggregated = aggregate_findings(&findings);
assert_eq!(aggregated[0].phase_id, "high");
assert_eq!(aggregated[1].phase_id, "low");
}
#[test]
fn aggregate_findings_empty_input_returns_empty() {
assert!(aggregate_findings(&[]).is_empty());
}
#[test]
fn score_content_relevance_handles_no_tokens() {
assert_eq!(score_content_relevance("a", "beta gamma delta"), 0.0);
}
#[test]
fn rank_sources_prefers_higher_authority() {
let now = 1_700_000_000;
let candidates = vec![
SourceCandidate {
source: Source { title: "Blog".into(), url: Some("https://example.com/post".into()), relevance_score: 0.0 },
published_epoch_secs: Some(now), authority_score: 0.55, content_match_score: 0.5,
},
SourceCandidate {
source: Source { title: "Government report".into(), url: Some("https://agency.gov/report".into()), relevance_score: 0.0 },
published_epoch_secs: Some(now), authority_score: 1.0, content_match_score: 0.5,
},
];
let ranked = rank_sources(&candidates, "report", now);
assert!(ranked[0].url.as_deref().unwrap().contains(".gov"));
}
#[test]
fn rank_sources_prefers_more_recent_publication() {
let now = 1_700_000_000;
let candidates = vec![
SourceCandidate {
source: Source { title: "Old".into(), url: Some("https://example.com/old".into()), relevance_score: 0.0 },
published_epoch_secs: Some(now - 400 * 86_400), authority_score: 0.55, content_match_score: 0.6,
},
SourceCandidate {
source: Source { title: "New".into(), url: Some("https://example.com/new".into()), relevance_score: 0.0 },
published_epoch_secs: Some(now - 7 * 86_400), authority_score: 0.55, content_match_score: 0.6,
},
];
let ranked = rank_sources(&candidates, "example", now);
assert_eq!(ranked[0].title, "New");
}
#[test]
fn rank_sources_prefers_better_content_match() {
let now = 1_700_000_000;
let candidates = vec![
SourceCandidate {
source: Source { title: "Unrelated".into(), url: None, relevance_score: 0.0 },
published_epoch_secs: Some(now), authority_score: 0.55, content_match_score: 0.1,
},
SourceCandidate {
source: Source { title: "Quantum networking overview".into(), url: None, relevance_score: 0.0 },
published_epoch_secs: Some(now), authority_score: 0.55, content_match_score: 0.95,
},
];
let ranked = rank_sources(&candidates, "quantum networking", now);
assert!(ranked[0].title.contains("Quantum"));
}
#[test]
fn run_research_plan_parallel_collects_all_phases() {
let plan = plan_research_phases("root", &sample_sub_queries(), true);
let mut executor = |phase: &ResearchPhase| Ok(format!("answer for {}", phase.id));
let result = run_research_plan(&plan, &mut executor);
assert_eq!(result.findings.len(), 3);
assert!(!result.partial);
assert!(result.failed_phase_ids.is_empty());
}
#[test]
fn run_research_plan_sequential_stops_after_failure() {
let plan = plan_research_phases("root", &["first".into(), "second".into(), "third".into()], false);
let mut executor = |phase: &ResearchPhase| {
if phase.id == "phase-2" { Err("phase two failed".into()) } else { Ok(format!("ok {}", phase.id)) }
};
let result = run_research_plan(&plan, &mut executor);
assert_eq!(result.findings.len(), 1);
assert!(result.partial);
assert_eq!(result.failed_phase_ids, vec!["phase-2".to_string()]);
}
#[test]
fn run_research_plan_partial_success_flag() {
let plan = plan_research_phases("root", &["only".into()], true);
let mut ok = |_: &ResearchPhase| Ok("content".into());
assert!(!run_research_plan(&plan, &mut ok).partial);
let mut fail = |_: &ResearchPhase| Err("boom".into());
let failure = run_research_plan(&plan, &mut fail);
assert!(!failure.partial);
assert!(failure.findings.is_empty());
}
#[test]
fn research_plan_display_includes_phases() {
let plan = plan_research_phases("quantum", &["vendors".into()], false);
let text = plan.to_string();
assert!(text.contains("ResearchPlan: quantum"));
assert!(text.contains("phase-1"));
}
#[test]
fn research_phase_display_formats_id_and_query() {
let phase = ResearchPhase { id: "phase-9".into(), query: "market size".into(), mode: PhaseExecutionMode::Parallel, depends_on: vec![] };
let text = phase.to_string();
assert!(text.contains("phase-9"));
assert!(text.contains("market size"));
}
#[test]
fn research_result_display_lists_findings() {
let result = ResearchResult {
findings: vec![ResearchFinding { phase_id: "phase-1".into(), content: "data".into(), relevance_score: 0.5 }],
sources: vec![], failed_phase_ids: vec![], partial: false,
};
let text = result.to_string();
assert!(text.contains("1 finding(s)"));
assert!(text.contains("phase-1"));
}
#[test]
fn research_plan_clone_and_debug() {
let plan = plan_research_phases("root", &["q".into()], true);
let cloned = plan.clone();
assert_eq!(format!("{plan:?}"), format!("{cloned:?}"));
}
#[test]
fn research_usage_debug_clone() {
let usage = ResearchUsage { input_tokens: 3, output_tokens: 7 };
let cloned = usage.clone();
assert_eq!(format!("{usage:?}"), format!("{cloned:?}"));
}
#[test]
fn score_authority_boosts_edu_and_gov_domains() {
assert_eq!(score_authority(Some("https://university.edu/paper")), 1.0);
assert!(score_authority(Some("https://blog.example.com")) < 1.0);
}
#[test]
fn findings_to_source_candidates_extracts_url_and_epoch() {
let findings = vec![ResearchFinding {
phase_id: "phase-1".into(),
content: "quantum paper published:1700000000 https://arxiv.org/abs/123".into(),
relevance_score: 0.8,
}];
let candidates = findings_to_source_candidates(&findings, "quantum paper");
assert_eq!(candidates[0].source.url.as_deref(), Some("https://arxiv.org/abs/123"));
assert_eq!(candidates[0].published_epoch_secs, Some(1_700_000_000));
assert!(candidates[0].authority_score >= 0.85);
}
#[tokio::test]
async fn test_research_coordinator_end_to_end() {
let coordinator = ResearchCoordinator::new(Box::new(ResearchMockLlm::new()), 2, 2);
let (report, sources) = coordinator
.research("quantum networking trends")
.await
.expect("research should succeed");
assert!(report.contains("Comprehensive synthesized answer"));
assert_eq!(sources.len(), 2);
assert_eq!(sources[0].title, "Research Finding 1");
assert!(sources[0].relevance_score > 0.0);
}
#[tokio::test]
async fn test_research_with_usage_accumulates_tokens() {
let coordinator = ResearchCoordinator::new(Box::new(ResearchMockLlm::new()), 2, 2);
let (_, _, usage) = coordinator
.research_with_usage("market landscape")
.await
.expect("research_with_usage");
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 45);
}
#[tokio::test]
async fn test_research_propagates_llm_errors() {
let coordinator = ResearchCoordinator::new(Box::new(ResearchMockLlm::failing()), 1, 1);
let err = coordinator
.research("anything")
.await
.expect_err("expected llm failure");
assert!(matches!(err, AppError::Internal(_)));
}
#[test]
fn test_token_usage_serde_roundtrip() {
let usage = TokenUsage::new(11, 7);
let json = serde_json::to_string(&usage).expect("serialize TokenUsage");
let parsed: TokenUsage = serde_json::from_str(&json).expect("deserialize TokenUsage");
assert_eq!(parsed, usage);
}
#[test]
fn test_research_usage_default_is_zero() {
let usage = ResearchUsage::default();
assert_eq!(usage.input_tokens, 0);
assert_eq!(usage.output_tokens, 0);
}
#[tokio::test]
async fn test_generate_followup_questions_empty_findings() {
let coordinator = ResearchCoordinator::new(Box::new(ResearchMockLlm::new()), 2, 2);
let (questions, usage) = coordinator
.generate_followup_questions("original query", &[])
.await
.expect("empty findings should succeed");
assert!(questions.is_empty());
assert!(usage.is_none());
}
}