graphrag_core/query/
planner.rs

1use crate::core::Result;
2use crate::ollama::OllamaClient;
3use serde::{Deserialize, Serialize};
4
5/// Planner for breaking down complex queries
6pub struct QueryPlanner {
7    client: OllamaClient,
8}
9
10#[derive(Debug, Serialize, Deserialize)]
11struct DecompositionResponse {
12    sub_queries: Vec<String>,
13}
14
15impl QueryPlanner {
16    /// Create a new query planner
17    pub fn new(client: OllamaClient) -> Self {
18        Self { client }
19    }
20
21    /// Decompose a complex query into simpler sub-queries
22    pub async fn decompose(&self, query: &str) -> Result<Vec<String>> {
23        let prompt = format!(
24            "You are an expert query planner for a RAG system. \
25            Your task is to decompose the following complex user query into a list of simple, independent sub-queries \
26            that can be answered using vector search. \
27            \
28            Return ONLY a raw JSON object with a single key 'sub_queries' containing the list of strings. \
29            Do not include any explanation, markdown formatting, or preamble. \
30            \
31            Query: '{}' \
32            \
33            JSON Response:",
34            query
35        );
36
37        let response_text = self.client.generate(&prompt).await?;
38
39        // Clean up potential markdown formatting (```json ... ```)
40        let cleaned_json = response_text
41            .trim()
42            .trim_start_matches("```json")
43            .trim_start_matches("```")
44            .trim_end_matches("```")
45            .trim();
46
47        // Parse JSON response
48        let response: DecompositionResponse = serde_json::from_str(cleaned_json).map_err(|e| {
49            crate::core::GraphRAGError::Generation {
50                message: format!(
51                    "Failed to parse planner response: {}. Text: {}",
52                    e, cleaned_json
53                ),
54            }
55        })?;
56
57        Ok(response.sub_queries)
58    }
59}
graphrag_core/query/planner.rs

graphrag_core/query/
planner.rs
