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 `Node`.
6
7use anyhow::Result;
8use async_trait::async_trait;
9use swiftide_core::{Transformer, indexing::Node};
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    /// Transforms an `Node` by generating questions and answers
26    /// based on the text chunk within the node.
27    ///
28    /// # Arguments
29    ///
30    /// * `node` - The `Node` containing the text chunk to process.
31    ///
32    /// # Returns
33    ///
34    /// A `Result` containing the transformed `Node` with added metadata,
35    /// or an error if the transformation fails.
36    ///
37    /// # Errors
38    ///
39    /// This function will return an error if the client fails to generate
40    /// questions and answers from the provided prompt.
41    #[tracing::instrument(skip_all, name = "transformers.metadata_qa_text")]
42    async fn transform_node(&self, mut node: Node) -> Result<Node> {
43        let prompt = self
44            .prompt_template
45            .clone()
46            .with_node(&node)
47            .with_context_value("questions", self.num_questions);
48
49        let response = self.prompt(prompt).await?;
50
51        node.metadata.insert(NAME, response);
52
53        Ok(node)
54    }
55
56    fn concurrency(&self) -> Option<usize> {
57        self.concurrency
58    }
59}
60
61#[cfg(test)]
62mod test {
63    use swiftide_core::MockSimplePrompt;
64
65    use super::*;
66
67    #[tokio::test]
68    async fn test_template() {
69        let template = default_prompt();
70
71        let prompt = template
72            .clone()
73            .with_node(&Node::new("test"))
74            .with_context_value("questions", 5);
75        insta::assert_snapshot!(prompt.render().unwrap());
76    }
77
78    #[tokio::test]
79    async fn test_metadata_qacode() {
80        let mut client = MockSimplePrompt::new();
81
82        client
83            .expect_prompt()
84            .returning(|_| Ok("Q1: Hello\nA1: World".to_string()));
85
86        let transformer = MetadataQAText::builder().client(client).build().unwrap();
87        let node = Node::new("Some text");
88
89        let result = transformer.transform_node(node).await.unwrap();
90
91        assert_eq!(
92            result.metadata.get("Questions and Answers (text)").unwrap(),
93            "Q1: Hello\nA1: World"
94        );
95    }
96}
swiftide_indexing/transformers/metadata_qa_text.rs

swiftide_indexing/transformers/
metadata_qa_text.rs
