swiftide_indexing/transformers/
metadata_summary.rs

1//! Generate a summary and adds it as metadata
2//! This module defines the `MetadataSummary` struct and its associated methods,
3//! which are used for generating metadata in the form of a summary
4//! for a given text. It interacts with a client (e.g., `OpenAI`) to generate
5//! the summary based on the text chunk in an `TextNode`.
6
7use anyhow::Result;
8use async_trait::async_trait;
9use swiftide_core::{Transformer, indexing::TextNode};
10
11/// `MetadataSummary` is responsible for generating a summary
12/// for a given text chunk. It uses a templated prompt to interact with a client
13/// that implements the `SimplePrompt` trait.
14#[swiftide_macros::indexing_transformer(
15    metadata_field_name = "Summary",
16    default_prompt_file = "prompts/metadata_summary.prompt.md"
17)]
18pub struct MetadataSummary {}
19
20#[async_trait]
21impl Transformer for MetadataSummary {
22    type Input = String;
23    type Output = String;
24    /// Transforms an `TextNode` by extracting a summary
25    /// based on the text chunk within the node.
26    ///
27    /// # Arguments
28    ///
29    /// * `node` - The `TextNode` containing the text chunk to process.
30    ///
31    /// # Returns
32    ///
33    /// A `Result` containing the transformed `TextNode` with added metadata,
34    /// or an error if the transformation fails.
35    ///
36    /// # Errors
37    ///
38    /// This function will return an error if the client fails to generate
39    /// a summary from the provided prompt.
40    #[tracing::instrument(skip_all, name = "transformers.metadata_summary")]
41    async fn transform_node(&self, mut node: TextNode) -> Result<TextNode> {
42        let prompt = self.prompt_template.clone().with_node(&node);
43
44        let response = self.prompt(prompt).await?;
45
46        node.metadata.insert(NAME, response);
47
48        Ok(node)
49    }
50
51    fn concurrency(&self) -> Option<usize> {
52        self.concurrency
53    }
54}
55
56#[cfg(test)]
57mod test {
58    use swiftide_core::MockSimplePrompt;
59
60    use super::*;
61
62    #[tokio::test]
63    async fn test_template() {
64        let template = default_prompt();
65
66        let prompt = template.clone().with_node(&TextNode::new("test"));
67        insta::assert_snapshot!(prompt.render().unwrap());
68    }
69
70    #[tokio::test]
71    async fn test_metadata_summary() {
72        let mut client = MockSimplePrompt::new();
73
74        client
75            .expect_prompt()
76            .returning(|_| Ok("A Summary".to_string()));
77
78        let transformer = MetadataSummary::builder().client(client).build().unwrap();
79        let node = TextNode::new("Some text");
80
81        let result = transformer.transform_node(node).await.unwrap();
82
83        assert_eq!(result.metadata.get("Summary").unwrap(), "A Summary");
84    }
85}
swiftide_indexing/transformers/metadata_summary.rs

swiftide_indexing/transformers/
metadata_summary.rs
