lc_agents/deep_research/
planner.rs

1// src/agents/deep_research/planner.rs
2//! Planner - decomposes a research topic into sub-topics and generates
3//! search queries for each sub-topic using the LLM.
4
5use lc_core::language_models::BaseChatModel;
6use lc_schema::Message;
7
8use super::ResearchError;
9
10/// A sub-topic with its associated search queries.
11#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
12pub struct SubTopic {
13    /// Short name for the sub-topic.
14    pub name: String,
15    /// Search queries to investigate this sub-topic.
16    pub queries: Vec<String>,
17}
18
19/// A research plan containing decomposed sub-topics.
20#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
21pub struct ResearchPlan {
22    /// The original research topic.
23    pub topic: String,
24    /// Decomposed sub-topics with their queries.
25    pub subtopics: Vec<SubTopic>,
26}
27
28impl ResearchPlan {
29    /// Collects all search queries across all sub-topics.
30    pub fn all_queries(&self) -> Vec<String> {
31        self.subtopics
32            .iter()
33            .flat_map(|st| st.queries.clone())
34            .collect()
35    }
36}
37
38/// Uses the LLM to decompose a topic into sub-topics with search queries.
39pub async fn plan<M: BaseChatModel>(
40    llm: &M,
41    topic: &str,
42    max_subtopics: usize,
43) -> Result<ResearchPlan, ResearchError> {
44    let prompt = format!(
45        "Decompose the following research topic into at most {} sub-topics. \
46         For each sub-topic, generate 1-3 specific search queries.\n\n\
47         IMPORTANT: Ensure sub-topics are non-overlapping and collectively exhaustive (MECE principle). \
48         Each sub-topic should cover a distinct aspect of the topic without redundancy, \
49         and together they should comprehensively cover the entire topic.\n\n\
50         Example:\n\
51         Topic: Impact of AI on healthcare\n\
52         Output:\n\
53         [{{\"name\": \"AI in Diagnostics\", \"queries\": [\"AI diagnostic accuracy healthcare\", \"deep learning medical imaging\"]}}, \
54          {{\"name\": \"AI in Treatment Planning\", \"queries\": [\"AI treatment recommendation systems\"]}}, \
55          {{\"name\": \"Ethics and Regulation\", \"queries\": [\"AI healthcare regulation FDA\", \"medical AI bias ethics\"]}}]\n\n\
56         Topic: {}\n\n\
57         Output a JSON array of objects with \"name\" and \"queries\" fields. \
58         Only output the JSON array, nothing else.",
59        max_subtopics, topic,
60    );
61    let messages = vec![
62        Message::system("You are a research planning assistant. Output only valid JSON."),
63        Message::human(prompt),
64    ];
65
66    let response =
67        crate::retry::retry_chat(llm, messages, None, &crate::retry::RetryConfig::default())
68            .await
69            .map_err(|e| ResearchError::Llm(format!("{:?}", e)))?;
70
71    let subtopics = parse_subtopics(&response.content)?;
72    Ok(ResearchPlan {
73        topic: topic.to_string(),
74        subtopics,
75    })
76}
77
78/// Parses the LLM output into a list of `SubTopic` structs.
79///
80/// Uses tolerant JSON parsing that handles markdown code fences,
81/// trailing commas, and surrounding text.
82fn parse_subtopics(content: &str) -> Result<Vec<SubTopic>, ResearchError> {
83    lc_core::json_parse::parse_llm_json::<Vec<SubTopic>>(content)
84        .map_err(|e| ResearchError::Llm(format!("failed to parse sub-topics: {}", e)))
85}
86
87/// Parses a JSON array of strings from LLM output.
88///
89/// Uses tolerant JSON parsing. Kept for backward compatibility;
90/// `generate_follow_ups` now uses the gap→query mapping format.
91pub fn parse_json_array(content: &str) -> Result<Vec<String>, ResearchError> {
92    lc_core::json_parse::parse_llm_json::<Vec<String>>(content)
93        .map_err(|e| ResearchError::Llm(format!("failed to parse JSON array: {}", e)))
94}
95
96/// Extracts JSON from LLM output, tolerating markdown code fences
97/// and surrounding text.
98pub fn extract_json(content: &str) -> String {
99    let trimmed = content.trim();
100
101    // Strip markdown code fences
102    let stripped = if trimmed.starts_with("```") {
103        trimmed
104            .strip_prefix("```json")
105            .or_else(|| trimmed.strip_prefix("```"))
106            .unwrap_or(trimmed)
107            .strip_suffix("```")
108            .unwrap_or(trimmed)
109            .trim()
110    } else {
111        trimmed
112    };
113
114    // Find the first [ or { bracket, whichever comes first
115    let start_bracket = stripped.find(['[', '{']);
116
117    if let Some(start) = start_bracket {
118        let open_char = stripped.as_bytes()[start];
119        let close_char = if open_char == b'[' { b']' } else { b'}' };
120
121        // Find the matching closing bracket
122        let mut depth = 0i32;
123        let mut in_string = false;
124        let mut escape_next = false;
125        let bytes = stripped.as_bytes();
126
127        for i in start..bytes.len() {
128            let ch = bytes[i];
129            if escape_next {
130                escape_next = false;
131                continue;
132            }
133            if ch == b'\\' {
134                if in_string {
135                    escape_next = true;
136                }
137                continue;
138            }
139            if ch == b'"' {
140                in_string = !in_string;
141                continue;
142            }
143            if in_string {
144                continue;
145            }
146            if ch == open_char {
147                depth += 1;
148            } else if ch == close_char {
149                depth -= 1;
150                if depth == 0 {
151                    return stripped[start..=i].to_string();
152                }
153            }
154        }
155    }
156
157    stripped.to_string()
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn test_extract_json_plain_array() {
166        let input = r#"[{"name": "A", "queries": ["q1"]}]"#;
167        assert_eq!(extract_json(input), input);
168    }
169
170    #[test]
171    fn test_extract_json_markdown_fenced() {
172        let input = "```json\n[{\"name\": \"A\", \"queries\": [\"q1\"]}]\n```";
173        let expected = r#"[{"name": "A", "queries": ["q1"]}]"#;
174        assert_eq!(extract_json(input), expected);
175    }
176
177    #[test]
178    fn test_extract_json_with_surrounding_text() {
179        let input = r#"Here is the plan: [{"name": "A", "queries": ["q1"]}] Done."#;
180        let expected = r#"[{"name": "A", "queries": ["q1"]}]"#;
181        assert_eq!(extract_json(input), expected);
182    }
183
184    #[test]
185    fn test_extract_json_nested_brackets() {
186        let input = r#"[{"name": "A", "queries": ["q1", "q2"]}, {"name": "B", "queries": ["q3"]}]"#;
187        assert_eq!(extract_json(input), input);
188    }
189
190    #[test]
191    fn test_extract_json_object() {
192        let input = r#"Result: {"report": "text", "gaps": []} end"#;
193        let expected = r#"{"report": "text", "gaps": []}"#;
194        assert_eq!(extract_json(input), expected);
195    }
196
197    #[test]
198    fn test_parse_subtopics_valid() {
199        let content = r#"[{"name": "AI Ethics", "queries": ["AI ethics policy"]}]"#;
200        let subtopics = parse_subtopics(content).unwrap();
201        assert_eq!(subtopics.len(), 1);
202        assert_eq!(subtopics[0].name, "AI Ethics");
203        assert_eq!(subtopics[0].queries, vec!["AI ethics policy"]);
204    }
205
206    #[test]
207    fn test_parse_subtopics_invalid() {
208        let content = "not json at all";
209        let result = parse_subtopics(content);
210        assert!(result.is_err());
211    }
212
213    #[test]
214    fn test_parse_json_array_valid() {
215        let content = r#"["query1", "query2"]"#;
216        let queries = parse_json_array(content).unwrap();
217        assert_eq!(queries, vec!["query1", "query2"]);
218    }
219
220    #[test]
221    fn test_parse_json_array_markdown() {
222        let content = "```json\n[\"q1\", \"q2\"]\n```";
223        let queries = parse_json_array(content).unwrap();
224        assert_eq!(queries, vec!["q1", "q2"]);
225    }
226
227    #[test]
228    fn test_research_plan_all_queries() {
229        let plan = ResearchPlan {
230            topic: "test".to_string(),
231            subtopics: vec![
232                SubTopic {
233                    name: "A".to_string(),
234                    queries: vec!["q1".to_string(), "q2".to_string()],
235                },
236                SubTopic {
237                    name: "B".to_string(),
238                    queries: vec!["q3".to_string()],
239                },
240            ],
241        };
242        assert_eq!(plan.all_queries(), vec!["q1", "q2", "q3"]);
243    }
244
245    #[test]
246    fn test_extract_json_with_brackets_in_strings() {
247        let input = r#"[{"name": "A [1]", "queries": ["q1"]}]"#;
248        assert_eq!(extract_json(input), input);
249    }
250
251    /// Verify the planner prompt contains the MECE constraint and few-shot example.
252    #[test]
253    fn test_plan_prompt_contains_mece_and_example() {
254        // The prompt is built in the `plan` function via format!.
255        // We verify the static parts are present by checking the format string.
256        let max_subtopics = 3;
257        let topic = "Test Topic";
258        let prompt = format!(
259            "Decompose the following research topic into at most {} sub-topics. \
260             For each sub-topic, generate 1-3 specific search queries.\n\n\
261             IMPORTANT: Ensure sub-topics are non-overlapping and collectively exhaustive (MECE principle). \
262             Each sub-topic should cover a distinct aspect of the topic without redundancy, \
263             and together they should comprehensively cover the entire topic.\n\n\
264             Example:\n\
265             Topic: Impact of AI on healthcare\n\
266             Output:\n\
267             [{{\"name\": \"AI in Diagnostics\", \"queries\": [\"AI diagnostic accuracy healthcare\", \"deep learning medical imaging\"]}}, \
268              {{\"name\": \"AI in Treatment Planning\", \"queries\": [\"AI treatment recommendation systems\"]}}, \
269              {{\"name\": \"Ethics and Regulation\", \"queries\": [\"AI healthcare regulation FDA\", \"medical AI bias ethics\"]}}]\n\n\
270             Topic: {}\n\n\
271             Output a JSON array of objects with \"name\" and \"queries\" fields. \
272             Only output the JSON array, nothing else.",
273            max_subtopics, topic,
274        );
275        assert!(
276            prompt.contains("MECE"),
277            "prompt should contain MECE principle"
278        );
279        assert!(
280            prompt.contains("non-overlapping"),
281            "prompt should contain non-overlapping constraint"
282        );
283        assert!(
284            prompt.contains("Example:"),
285            "prompt should contain few-shot example"
286        );
287        assert!(
288            prompt.contains("AI in Diagnostics"),
289            "prompt example should contain sub-topic example"
290        );
291    }
292}
lc_agents/deep_research/planner.rs

lc_agents/deep_research/
planner.rs
