swiftide_indexing/transformers/
metadata_title.rs

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

swiftide_indexing/transformers/
metadata_title.rs
