use std::sync::Arc;
use swiftide_core::{
indexing::SimplePrompt,
prelude::*,
prompt::PromptTemplate,
querying::{states, Query},
TransformResponse,
};
#[derive(Debug, Clone, Builder)]
pub struct Summary {
#[builder(setter(custom))]
client: Arc<dyn SimplePrompt>,
#[builder(default = "default_prompt()")]
prompt_template: PromptTemplate,
}
impl Summary {
pub fn builder() -> SummaryBuilder {
SummaryBuilder::default()
}
pub fn from_client(client: impl SimplePrompt + 'static) -> Summary {
SummaryBuilder::default()
.client(client)
.to_owned()
.build()
.expect("Failed to build Summary")
}
}
impl SummaryBuilder {
pub fn client(&mut self, client: impl SimplePrompt + 'static) -> &mut Self {
self.client = Some(Arc::new(client));
self
}
}
fn default_prompt() -> PromptTemplate {
indoc::indoc!(
"
Your job is to help a query tool find the right context.
Summarize the following documents.
## Constraints
* Do not add any information that is not available in the documents.
* Summarize comprehensively and ensure no data that might be important is left out.
* Summarize as a single markdown document
## Documents
{% for document in documents -%}
---
{{ document }}
---
{% endfor -%}
"
)
.into()
}
#[async_trait]
impl TransformResponse for Summary {
#[tracing::instrument]
async fn transform_response(
&self,
mut query: Query<states::Retrieved>,
) -> Result<Query<states::Retrieved>> {
let new_response = self
.client
.prompt(
self.prompt_template
.to_prompt()
.with_context_value("documents", query.documents()),
)
.await?;
query.transformed_response(new_response);
Ok(query)
}
}
#[cfg(test)]
mod test {
use super::*;
assert_default_prompt_snapshot!("documents" => vec!["First document", "Second Document"]);
}