rucora_core/research/
strategies.rs

1//! Deep Research 核心 trait 定义
2//!
3//! 本模块定义了深度研究的核心抽象，包括：
4//! - `DeepResearchEngine` - 研究引擎 trait
5//! - `StrategyTrait` - 搜索策略 trait
6//! - `ResearchQualityAssessor` - 研究质量评估器
7//! - `CitationHandler` - 引用处理器
8
9use crate::provider::LlmProvider;
10use async_trait::async_trait;
11use std::sync::{Arc, OnceLock};
12
13use super::{
14    InfoPiece, ResearchConfig, ResearchPhase, ResearchProgress, ResearchReport, ScoringConfig,
15    SuggestionType,
16};
17
18/// 提取 URL 的正则表达式，使用 OnceLock 避免重复编译。
19fn url_regex() -> &'static regex::Regex {
20    static INSTANCE: OnceLock<regex::Regex> = OnceLock::new();
21    INSTANCE.get_or_init(|| regex::Regex::new(r"https?://[^\s\)]+").unwrap())
22}
23
24/// 深度研究引擎 trait
25///
26/// 负责执行完整的深度研究流程。
27#[async_trait]
28pub trait DeepResearchEngine: Send + Sync {
29    /// 执行研究
30    ///
31    /// # 参数
32    ///
33    /// * `provider` - LLM Provider
34    /// * `topic` - 研究主题
35    ///
36    /// # 返回
37    ///
38    /// 研究报告
39    async fn research(
40        &self,
41        provider: &Arc<dyn LlmProvider>,
42        topic: &str,
43    ) -> Result<ResearchReport, ResearchError>;
44
45    /// 获取研究进度
46    fn progress(&self) -> ResearchProgress;
47}
48
49/// 研究策略 trait
50///
51/// 定义不同研究策略的行为。
52#[async_trait]
53pub trait StrategyTrait: Send + Sync {
54    /// 策略名称
55    fn name(&self) -> &'static str;
56
57    /// 策略描述
58    fn description(&self) -> &'static str;
59
60    /// 执行搜索
61    async fn search(
62        &self,
63        provider: &Arc<dyn LlmProvider>,
64        topic: &str,
65        context: &mut ResearchContext,
66    ) -> Result<StrategyResult, ResearchError>;
67
68    /// 判断是否需要继续搜索
69    fn should_continue(&self, result: &StrategyResult) -> bool;
70
71    /// 获取最大迭代次数
72    fn max_iterations(&self) -> Option<u32>;
73
74    /// 获取配置
75    fn config(&self) -> &ResearchConfig;
76}
77
78/// 研究上下文
79///
80/// 贯穿整个研究流程的共享状态。
81#[derive(Debug)]
82pub struct ResearchContext {
83    /// 研究主题
84    pub topic: String,
85    /// 当前阶段
86    pub phase: ResearchPhase,
87    /// 已收集的信息片段
88    pub collected_info: Vec<super::InfoPiece>,
89    /// 已访问的 URL
90    pub visited_urls: Vec<String>,
91    /// 搜索历史
92    pub search_history: Vec<super::SearchHistory>,
93    /// 引用列表
94    pub citations: Vec<super::Citation>,
95    /// 内部状态（用于策略特定数据）
96    pub state: std::collections::HashMap<String, serde_json::Value>,
97}
98
99impl ResearchContext {
100    pub fn new(topic: &str) -> Self {
101        Self {
102            topic: topic.to_string(),
103            phase: ResearchPhase::Init,
104            collected_info: Vec::new(),
105            visited_urls: Vec::new(),
106            search_history: Vec::new(),
107            citations: Vec::new(),
108            state: std::collections::HashMap::new(),
109        }
110    }
111
112    pub fn add_info(&mut self, info: super::InfoPiece) {
113        if let Some(url) = &info.source_url
114            && !self.visited_urls.contains(url)
115        {
116            self.visited_urls.push(url.clone());
117        }
118        self.collected_info.push(info);
119    }
120
121    pub fn add_citation(&mut self, citation: super::Citation) {
122        if !self.citations.iter().any(|c| c.url == citation.url) {
123            self.citations.push(citation);
124        }
125    }
126
127    pub fn has_visited(&self, url: &str) -> bool {
128        self.visited_urls.contains(&url.to_string())
129    }
130
131    pub fn visit_url(&mut self, url: String) {
132        if !self.visited_urls.contains(&url) {
133            self.visited_urls.push(url);
134        }
135    }
136
137    pub fn set_phase(&mut self, phase: ResearchPhase) {
138        self.phase = phase;
139    }
140
141    pub fn add_search_history(&mut self, query: String, result_count: usize) {
142        self.search_history
143            .push(super::SearchHistory::new(query, result_count));
144    }
145
146    pub fn set_state(&mut self, key: &str, value: serde_json::Value) {
147        self.state.insert(key.to_string(), value);
148    }
149
150    pub fn get_state(&self, key: &str) -> Option<&serde_json::Value> {
151        self.state.get(key)
152    }
153
154    /// 收集的信息总字符数
155    pub fn total_content_length(&self) -> usize {
156        self.collected_info.iter().map(|i| i.content.len()).sum()
157    }
158}
159
160/// 策略执行结果
161#[derive(Debug)]
162pub struct StrategyResult {
163    /// 是否完成
164    pub is_complete: bool,
165    /// 新收集的信息
166    pub new_info: Vec<super::InfoPiece>,
167    /// 新发现的 URL
168    pub discovered_urls: Vec<String>,
169    /// 置信度 (0.0 - 1.0)
170    pub confidence: f32,
171    /// 搜索次数
172    pub search_count: u32,
173    /// 使用的 token 数量
174    pub tokens_used: u32,
175}
176
177impl Default for StrategyResult {
178    fn default() -> Self {
179        Self {
180            is_complete: false,
181            new_info: Vec::new(),
182            discovered_urls: Vec::new(),
183            confidence: 0.0,
184            search_count: 0,
185            tokens_used: 0,
186        }
187    }
188}
189
190impl StrategyResult {
191    pub fn complete() -> Self {
192        Self {
193            is_complete: true,
194            ..Default::default()
195        }
196    }
197
198    pub fn complete_with(mut self, confidence: f32, tokens_used: u32, search_count: u32) -> Self {
199        self.is_complete = true;
200        self.confidence = confidence;
201        self.tokens_used = tokens_used;
202        self.search_count = search_count;
203        self
204    }
205
206    pub fn with_info(mut self, info: Vec<super::InfoPiece>) -> Self {
207        self.new_info = info;
208        self
209    }
210
211    pub fn with_confidence(mut self, confidence: f32) -> Self {
212        self.confidence = confidence;
213        self
214    }
215
216    pub fn with_tokens(mut self, tokens: u32) -> Self {
217        self.tokens_used = tokens;
218        self
219    }
220}
221
222/// 研究错误
223///
224/// 在深度研究过程中可能出现的各类错误。
225#[derive(Debug, thiserror::Error)]
226pub enum ResearchError {
227    /// Provider 错误
228    Provider {
229        message: String,
230        #[source]
231        source: Option<Box<dyn std::error::Error + Send + Sync>>,
232    },
233    /// 工具错误
234    Tool {
235        message: String,
236        #[source]
237        source: Option<Box<dyn std::error::Error + Send + Sync>>,
238    },
239    /// 超时
240    Timeout,
241    /// 无效配置
242    InvalidConfig(String),
243    /// 研究失败
244    Failed {
245        message: String,
246        #[source]
247        source: Option<Box<dyn std::error::Error + Send + Sync>>,
248    },
249    /// 存储错误
250    Storage {
251        message: String,
252        #[source]
253        source: Option<Box<dyn std::error::Error + Send + Sync>>,
254    },
255}
256
257impl std::fmt::Display for ResearchError {
258    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
259        match self {
260            ResearchError::Provider { message, .. } => write!(f, "Provider 错误: {message}"),
261            ResearchError::Tool { message, .. } => write!(f, "工具错误: {message}"),
262            ResearchError::Timeout => write!(f, "超时"),
263            ResearchError::InvalidConfig(msg) => write!(f, "无效配置: {msg}"),
264            ResearchError::Failed { message, .. } => write!(f, "研究失败: {message}"),
265            ResearchError::Storage { message, .. } => write!(f, "存储错误: {message}"),
266        }
267    }
268}
269
270impl From<crate::error::AgentError> for ResearchError {
271    fn from(e: crate::error::AgentError) -> Self {
272        ResearchError::Failed {
273            message: e.to_string(),
274            source: Some(Box::new(e)),
275        }
276    }
277}
278
279/// 研究库 trait
280///
281/// 负责存储和检索研究结果。
282#[async_trait]
283pub trait ResearchLibrary: Send + Sync {
284    /// 保存研究结果
285    async fn save(&self, report: &ResearchReport) -> Result<String, ResearchError>;
286
287    /// 搜索历史研究
288    async fn search(&self, query: &str) -> Result<Vec<ResearchReport>, ResearchError>;
289
290    /// 获取特定研究
291    async fn get(&self, id: &str) -> Result<Option<ResearchReport>, ResearchError>;
292
293    /// 列出所有研究
294    async fn list(&self, limit: usize) -> Result<Vec<ResearchReport>, ResearchError>;
295
296    /// 删除研究
297    async fn delete(&self, id: &str) -> Result<(), ResearchError>;
298}
299
300/// 引用处理器 trait
301///
302/// 负责处理和格式化引用。
303pub trait CitationHandler: Send + Sync {
304    /// 从内容中提取引用
305    fn extract_citations(&self, content: &str) -> Vec<super::Citation>;
306
307    /// 格式化单个引用
308    fn format_citation(&self, citation: &super::Citation) -> String;
309
310    /// 格式化引用列表
311    fn format_reference_list(&self, citations: &[super::Citation]) -> String;
312}
313
314/// 默认引用处理器
315///
316/// 提供基于正则表达式的 URL 提取和引用格式化功能。
317pub struct DefaultCitationHandler;
318
319impl DefaultCitationHandler {
320    /// 创建新的默认引用处理器实例
321    pub fn new() -> Self {
322        Self
323    }
324}
325
326impl Default for DefaultCitationHandler {
327    fn default() -> Self {
328        Self::new()
329    }
330}
331
332impl CitationHandler for DefaultCitationHandler {
333    fn extract_citations(&self, content: &str) -> Vec<super::Citation> {
334        let mut citations = Vec::new();
335
336        for cap in url_regex().find_iter(content) {
337            let url = cap.as_str().to_string();
338            // 避免重复
339            if !citations.iter().any(|c: &super::Citation| c.url == url) {
340                citations.push(super::Citation::new(url, "".to_string(), "".to_string()));
341            }
342        }
343
344        citations
345    }
346
347    fn format_citation(&self, citation: &super::Citation) -> String {
348        citation.format_apa()
349    }
350
351    fn format_reference_list(&self, citations: &[super::Citation]) -> String {
352        if citations.is_empty() {
353            return "无可用引用".to_string();
354        }
355
356        let mut result = String::from("## 参考来源\n\n");
357        for (i, citation) in citations.iter().enumerate() {
358            result.push_str(&format!("{}. {}\n\n", i + 1, citation.format_apa()));
359        }
360        result
361    }
362}
363
364/// 搜索策略工厂
365///
366/// 用于创建不同类型的搜索策略实例。
367pub trait StrategyFactory: Send + Sync {
368    /// 创建策略
369    fn create(&self, config: &ResearchConfig) -> Box<dyn StrategyTrait>;
370}
371
372/// 研究质量评估器
373///
374/// 用于评估研究质量、生成改进建议。
375pub struct ResearchQualityAssessor {
376    /// 评分配置
377    pub config: ScoringConfig,
378    /// 主题关键词
379    pub topic_keywords: Vec<String>,
380}
381
382impl ResearchQualityAssessor {
383    pub fn new(config: ScoringConfig, topic_keywords: Vec<String>) -> Self {
384        Self {
385            config,
386            topic_keywords,
387        }
388    }
389
390    pub fn with_default(topic: &str) -> Self {
391        Self {
392            config: ScoringConfig::default(),
393            topic_keywords: extract_keywords(topic),
394        }
395    }
396
397    /// 评估研究质量
398    pub fn assess(
399        &self,
400        info_pieces: &[InfoPiece],
401        citations: &[super::Citation],
402        search_rounds: usize,
403    ) -> super::ResearchQualityScore {
404        super::ResearchQualityScore::calculate(
405            info_pieces,
406            citations,
407            search_rounds,
408            &self.topic_keywords,
409        )
410    }
411
412    /// 生成改进建议
413    pub fn suggest(&self, score: &super::ResearchQualityScore) -> super::ResearchSuggestion {
414        // 检查是否已达到目标
415        if score.is_sufficient(self.config.quality_threshold) {
416            return super::ResearchSuggestion::sufficient();
417        }
418
419        // 检查信息数量
420        if score.details.info_count < self.config.min_info_count {
421            return super::ResearchSuggestion::need_more_info(
422                score.details.info_count,
423                self.config.min_info_count,
424            );
425        }
426
427        // 检查重复率
428        if score.details.duplicate_ratio > self.config.duplicate_threshold {
429            return super::ResearchSuggestion::need_new_angle(score.details.duplicate_ratio);
430        }
431
432        // 检查来源多样性
433        if score.details.source_diversity < self.config.min_source_diversity {
434            return super::ResearchSuggestion::need_more_sources(score.details.source_diversity);
435        }
436
437        // 检查置信度
438        if score.confidence < self.config.confidence_threshold {
439            return super::ResearchSuggestion::need_validation(score.confidence);
440        }
441
442        // 默认：信息不足
443        super::ResearchSuggestion::need_more_info(
444            score.details.info_count,
445            self.config.min_info_count,
446        )
447    }
448
449    /// 检查是否应该继续搜索
450    pub fn should_continue(
451        &self,
452        score: &super::ResearchQualityScore,
453        current_round: u32,
454        max_rounds: u32,
455    ) -> bool {
456        // 达到最大轮次，不再继续
457        if current_round >= max_rounds {
458            return false;
459        }
460
461        // 已达到质量阈值，可以停止
462        if score.is_sufficient(self.config.quality_threshold) {
463            return false;
464        }
465
466        // 质量不足但还有轮次，继续
467        true
468    }
469
470    /// 根据评分获取下一轮搜索建议
471    pub fn get_next_search_hint(
472        &self,
473        score: &super::ResearchQualityScore,
474        current_topic: &str,
475    ) -> String {
476        let suggestion = self.suggest(score);
477
478        let mut hints = vec![current_topic.to_string()];
479
480        // 添加主题关键词
481        if !self.topic_keywords.is_empty() {
482            hints.extend(self.topic_keywords.iter().take(2).cloned());
483        }
484
485        // 根据建议类型添加关键词
486        match suggestion.suggestion_type {
487            SuggestionType::NeedMoreInfo => {
488                hints.push("详细介绍".to_string());
489                hints.push("详细说明".to_string());
490            }
491            SuggestionType::NeedMoreSources => {
492                hints.extend(suggestion.suggested_keywords);
493            }
494            SuggestionType::NeedNewAngle => {
495                hints.extend(suggestion.suggested_keywords);
496            }
497            SuggestionType::NeedValidation => {
498                hints.push("官方".to_string());
499                hints.push("验证".to_string());
500            }
501            SuggestionType::Sufficient => {
502                return "研究已完成".to_string();
503            }
504        }
505
506        hints.join(" ")
507    }
508}
509
510/// 提取中文/英文关键词用于主题分析。
511///
512/// 按常见分隔符（空格、逗号、中文标点等）分割主题文本，
513/// 过滤短词（<=1字符），最多返回 5 个关键词。
514///
515/// # 参数
516///
517/// - `topic`: 研究主题文本
518///
519/// # 返回
520///
521/// 关键词列表，最多 5 个
522///
523/// # 示例
524///
525/// ```ignore
526/// let keywords = extract_keywords("Rust 异步编程");
527/// assert!(keywords.len() <= 5);
528/// ```
529pub(crate) fn extract_keywords(topic: &str) -> Vec<String> {
530    topic
531        .split(&[' ', ',', '，', '。', '、', '？', '?'][..])
532        .filter(|s| !s.is_empty() && s.len() > 1)
533        .take(5)
534        .map(|s| s.to_string())
535        .collect()
536}
rucora_core/research/strategies.rs

rucora_core/research/
strategies.rs
