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 `Node`.
6use anyhow::Result;
7use async_trait::async_trait;
8use swiftide_core::{indexing::Node, Transformer};
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    /// Transforms an `Node` by extracting a keywords
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    /// a keywords from the provided prompt.
37    #[tracing::instrument(skip_all, name = "transformers.metadata_keywords")]
38    async fn transform_node(&self, mut node: Node) -> Result<Node> {
39        let prompt = self.prompt_template.to_prompt().with_node(&node);
40        let response = self.prompt(prompt).await?;
41
42        node.metadata.insert(NAME, response);
43
44        Ok(node)
45    }
46
47    fn concurrency(&self) -> Option<usize> {
48        self.concurrency
49    }
50}
51
52#[cfg(test)]
53mod test {
54    use swiftide_core::MockSimplePrompt;
55
56    use super::*;
57
58    #[test_log::test(tokio::test)]
59    async fn test_template() {
60        let template = default_prompt();
61
62        let prompt = template.to_prompt().with_node(&Node::new("test"));
63        insta::assert_snapshot!(prompt.render().await.unwrap());
64    }
65
66    #[tokio::test]
67    async fn test_metadata_keywords() {
68        let mut client = MockSimplePrompt::new();
69
70        client
71            .expect_prompt()
72            .returning(|_| Ok("important,keywords".to_string()));
73
74        let transformer = MetadataKeywords::builder().client(client).build().unwrap();
75        let node = Node::new("Some text");
76
77        let result = transformer.transform_node(node).await.unwrap();
78
79        assert_eq!(
80            result.metadata.get("Keywords").unwrap(),
81            "important,keywords"
82        );
83    }
84}
swiftide_indexing/transformers/metadata_keywords.rs

swiftide_indexing/transformers/
metadata_keywords.rs
