swiftide_query/query_transformers/
generate_subquestions.rs

1//! Generate subquestions for a query
2//!
3//! Useful for similarity search where you want a wider vector coverage
4use std::sync::Arc;
5use swiftide_core::{
6    indexing::SimplePrompt,
7    prelude::*,
8    prompt::Prompt,
9    querying::{Query, TransformQuery, states},
10};
11
12#[derive(Debug, Clone, Builder)]
13pub struct GenerateSubquestions {
14    #[builder(setter(custom))]
15    client: Arc<dyn SimplePrompt>,
16    #[builder(default = "default_prompt()")]
17    prompt_template: Prompt,
18    #[builder(default = "5")]
19    num_questions: usize,
20}
21
22impl GenerateSubquestions {
23    pub fn builder() -> GenerateSubquestionsBuilder {
24        GenerateSubquestionsBuilder::default()
25    }
26
27    /// Builds a new subquestions generator from a client that implements [`SimplePrompt`]
28    ///
29    /// # Panics
30    ///
31    /// Panics if the build failed
32    pub fn from_client(client: impl SimplePrompt + 'static) -> GenerateSubquestions {
33        GenerateSubquestionsBuilder::default()
34            .client(client)
35            .to_owned()
36            .build()
37            .expect("Failed to build GenerateSubquestions")
38    }
39}
40
41impl GenerateSubquestionsBuilder {
42    pub fn client(&mut self, client: impl SimplePrompt + 'static) -> &mut Self {
43        self.client = Some(Arc::new(client) as Arc<dyn SimplePrompt>);
44        self
45    }
46}
47
48fn default_prompt() -> Prompt {
49    indoc::indoc!("
50    Your job is to help a query tool find the right context.
51
52    Given the following question:
53    {{question}}
54
55    Please think of {{num_questions}}  additional questions that can help answering the original question.
56
57    Especially consider what might be relevant to answer the question, like dependencies, usage and structure of the code.
58
59    Please respond with the original question and the additional questions only.
60
61    ## Example
62
63    - {{question}}
64    - Additional question 1
65    - Additional question 2
66    - Additional question 3
67    - Additional question 4
68    - Additional question 5
69    ").into()
70}
71
72#[async_trait]
73impl TransformQuery for GenerateSubquestions {
74    #[tracing::instrument(skip_self)]
75    async fn transform_query(
76        &self,
77        mut query: Query<states::Pending>,
78    ) -> Result<Query<states::Pending>> {
79        let new_query = self
80            .client
81            .prompt(
82                self.prompt_template
83                    .clone()
84                    .with_context_value("question", query.current())
85                    .with_context_value("num_questions", self.num_questions),
86            )
87            .await?;
88        query.transformed_query(new_query);
89
90        Ok(query)
91    }
92}
93
94#[cfg(test)]
95mod test {
96    use super::*;
97
98    assert_default_prompt_snapshot!("question" => "What is love?", "num_questions" => 5);
99}
swiftide_query/query_transformers/generate_subquestions.rs

swiftide_query/query_transformers/
generate_subquestions.rs
