swiftide_indexing/transformers/
metadata_qa_text.rs

1//! Generates questions and answers from a given text chunk and adds them as metadata.
2//! This module defines the `MetadataQAText` struct and its associated methods,
3//! which are used for generating metadata in the form of questions and answers
4//! from a given text. It interacts with a client (e.g., `OpenAI`) to generate
5//! these questions and answers 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/// `MetadataQAText` is responsible for generating questions and answers
12/// from 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 = "Questions and Answers (text)",
16    default_prompt_file = "prompts/metadata_qa_text.prompt.md"
17)]
18pub struct MetadataQAText {
19    #[builder(default = "5")]
20    num_questions: usize,
21}
22
23#[async_trait]
24impl Transformer for MetadataQAText {
25    type Input = String;
26    type Output = String;
27
28    /// Transforms an `TextNode` by generating questions and answers
29    /// based on the text chunk within the node.
30    ///
31    /// # Arguments
32    ///
33    /// * `node` - The `TextNode` containing the text chunk to process.
34    ///
35    /// # Returns
36    ///
37    /// A `Result` containing the transformed `TextNode` with added metadata,
38    /// or an error if the transformation fails.
39    ///
40    /// # Errors
41    ///
42    /// This function will return an error if the client fails to generate
43    /// questions and answers from the provided prompt.
44    #[tracing::instrument(skip_all, name = "transformers.metadata_qa_text")]
45    async fn transform_node(&self, mut node: TextNode) -> Result<TextNode> {
46        let prompt = self
47            .prompt_template
48            .clone()
49            .with_node(&node)
50            .with_context_value("questions", self.num_questions);
51
52        let response = self.prompt(prompt).await?;
53
54        node.metadata.insert(NAME, response);
55
56        Ok(node)
57    }
58
59    fn concurrency(&self) -> Option<usize> {
60        self.concurrency
61    }
62}
63
64#[cfg(test)]
65mod test {
66    use swiftide_core::MockSimplePrompt;
67
68    use super::*;
69
70    #[tokio::test]
71    async fn test_template() {
72        let template = default_prompt();
73
74        let prompt = template
75            .clone()
76            .with_node(&TextNode::new("test"))
77            .with_context_value("questions", 5);
78        insta::assert_snapshot!(prompt.render().unwrap());
79    }
80
81    #[tokio::test]
82    async fn test_metadata_qacode() {
83        let mut client = MockSimplePrompt::new();
84
85        client
86            .expect_prompt()
87            .returning(|_| Ok("Q1: Hello\nA1: World".to_string()));
88
89        let transformer = MetadataQAText::builder().client(client).build().unwrap();
90        let node = TextNode::new("Some text");
91
92        let result = transformer.transform_node(node).await.unwrap();
93
94        assert_eq!(
95            result.metadata.get("Questions and Answers (text)").unwrap(),
96            "Q1: Hello\nA1: World"
97        );
98    }
99}
swiftide_indexing/transformers/metadata_qa_text.rs

swiftide_indexing/transformers/
metadata_qa_text.rs
