swiftide_query/response_transformers/
summary.rs

1use std::sync::Arc;
2use swiftide_core::{
3    TransformResponse,
4    indexing::SimplePrompt,
5    prelude::*,
6    prompt::Prompt,
7    querying::{Query, states},
8};
9
10#[derive(Debug, Clone, Builder)]
11pub struct Summary {
12    #[builder(setter(custom))]
13    client: Arc<dyn SimplePrompt>,
14    #[builder(default = "default_prompt()")]
15    prompt_template: Prompt,
16}
17
18impl Summary {
19    pub fn builder() -> SummaryBuilder {
20        SummaryBuilder::default()
21    }
22
23    /// Builds a new summary generator from a client that implements [`SimplePrompt`].
24    ///
25    /// Will try to summarize documents using an llm, instructed to preserve as much information as
26    /// possible.
27    ///
28    /// # Panics
29    ///
30    /// Panics if the build failed
31    pub fn from_client(client: impl SimplePrompt + 'static) -> Summary {
32        SummaryBuilder::default()
33            .client(client)
34            .to_owned()
35            .build()
36            .expect("Failed to build Summary")
37    }
38}
39
40impl SummaryBuilder {
41    pub fn client(&mut self, client: impl SimplePrompt + 'static) -> &mut Self {
42        self.client = Some(Arc::new(client) as Arc<dyn SimplePrompt>);
43        self
44    }
45}
46
47fn default_prompt() -> Prompt {
48    indoc::indoc!(
49        "
50    Your job is to help a query tool find the right context.
51
52    Summarize the following documents.
53
54    ## Constraints
55    * Do not add any information that is not available in the documents.
56    * Summarize comprehensively and ensure no data that might be important is left out.
57    * Summarize as a single markdown document
58
59    ## Documents
60
61    {% for document in documents -%}
62    ---
63    {{ document.content }}
64    ---
65    {% endfor -%}
66    "
67    )
68    .into()
69}
70
71#[async_trait]
72impl TransformResponse for Summary {
73    #[tracing::instrument(skip_all)]
74    async fn transform_response(
75        &self,
76        mut query: Query<states::Retrieved>,
77    ) -> Result<Query<states::Retrieved>> {
78        let new_response = self
79            .client
80            .prompt(
81                self.prompt_template
82                    .clone()
83                    .with_context_value("documents", query.documents()),
84            )
85            .await?;
86        query.transformed_response(new_response);
87
88        Ok(query)
89    }
90}
91
92#[cfg(test)]
93mod test {
94    use swiftide_core::document::Document;
95
96    use super::*;
97
98    assert_default_prompt_snapshot!("documents" => vec![Document::from("First document"), Document::from("Second Document")]);
99}
swiftide_query/response_transformers/summary.rs

swiftide_query/response_transformers/
summary.rs
