swiftide_indexing/transformers/
metadata_keywords.rs

1//! Extract keywords from a node and add them as metadata
2//! This module defines the `MetadataKeywords` struct and its associated methods,
3//! which are used for generating metadata in the form of keywords
4//! for a given text. It interacts with a client (e.g., `OpenAI`) to generate
5//! the keywords based on the text chunk in a `TextNode`.
6use anyhow::Result;
7use async_trait::async_trait;
8use swiftide_core::{Transformer, indexing::TextNode};
9
10/// `MetadataKeywords` is responsible for generating keywords
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    default_prompt_file = "prompts/metadata_keywords.prompt.md",
15    metadata_field_name = "Keywords"
16)]
17pub struct MetadataKeywords {}
18
19#[async_trait]
20impl Transformer for MetadataKeywords {
21    type Input = String;
22    type Output = String;
23
24    /// Transforms an `TextNode` by extracting a keywords
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 keywords from the provided prompt.
40    #[tracing::instrument(skip_all, name = "transformers.metadata_keywords")]
41    async fn transform_node(&self, mut node: TextNode) -> Result<TextNode> {
42        let prompt = self.prompt_template.clone().with_node(&node);
43        let response = self.prompt(prompt).await?;
44
45        node.metadata.insert(NAME, response);
46
47        Ok(node)
48    }
49
50    fn concurrency(&self) -> Option<usize> {
51        self.concurrency
52    }
53}
54
55#[cfg(test)]
56mod test {
57    use swiftide_core::MockSimplePrompt;
58
59    use super::*;
60
61    #[test_log::test(tokio::test)]
62    async fn test_template() {
63        let template = default_prompt();
64
65        let prompt = template.clone().with_node(&TextNode::new("test"));
66        insta::assert_snapshot!(prompt.render().unwrap());
67    }
68
69    #[tokio::test]
70    async fn test_metadata_keywords() {
71        let mut client = MockSimplePrompt::new();
72
73        client
74            .expect_prompt()
75            .returning(|_| Ok("important,keywords".to_string()));
76
77        let transformer = MetadataKeywords::builder().client(client).build().unwrap();
78        let node = TextNode::new("Some text");
79
80        let result = transformer.transform_node(node).await.unwrap();
81
82        assert_eq!(
83            result.metadata.get("Keywords").unwrap(),
84            "important,keywords"
85        );
86    }
87}
swiftide_indexing/transformers/metadata_keywords.rs

swiftide_indexing/transformers/
metadata_keywords.rs
