Initial commit

src/controller/image/prompt.rs

@@ -0,0 +1,111 @@
+use crate::conversation::llm::{Author, Message};
+
+/// Builds a prompt from the original prompt and other messages in the conversation.
+///
+/// Only messages authored by the user are considered.
+///
+/// Messages that say "Again" (regardless of casing) are ignored. They are considered special messages
+/// which trigger re-generation, but do not need to be included in the prompt criteria.
+pub fn build(original_prompt: &str, other_messages: Vec<Message>) -> String {
+    let mut prompt = original_prompt.to_owned();
+
+    // Make a new messages vector that only contains messages we care about
+    let other_messages: Vec<Message> = other_messages
+        .into_iter()
+        .filter(|message| {
+            if let Author::User = message.author {
+                message.message_text.to_lowercase() != "again"
+            } else {
+                false
+            }
+        })
+        .collect();
+
+    if !other_messages.is_empty() {
+        prompt.push_str("\nOther criteria:");
+        for message in other_messages {
+            prompt.push_str(
+                format!("\n- {}", message.message_text.replace("\n", ". ").as_str()).as_str(),
+            );
+        }
+    }
+
+    prompt
+}
+
+#[cfg(test)]
+mod tests {
+    use super::build;
+    use super::{Author, Message};
+
+    struct TestCase {
+        original_prompt: &'static str,
+        messages: Vec<Message>,
+        expected_prompt: &'static str,
+    }
+
+    #[test]
+    fn test_build_prompt() {
+        let test_cases = vec![
+            // Simple case
+            TestCase {
+                original_prompt: "Generate a picture of a cat",
+                messages: vec![],
+                expected_prompt: "Generate a picture of a cat",
+            },
+            // Only a single user message
+            TestCase {
+                original_prompt: "Generate a picture of a dog",
+                messages: vec![Message {
+                    author: Author::User,
+                    message_text: "Must be blue".to_owned(),
+                }],
+                expected_prompt: "Generate a picture of a dog\nOther criteria:\n- Must be blue",
+            },
+            // Multiple complex user messages dispersed with assistant messages
+            TestCase {
+                original_prompt: "Generate a picture of an elephant",
+                messages: vec![Message {
+                    author: Author::User,
+                    message_text: "Must be blue".to_owned(),
+                },
+                Message {
+                    author: Author::Assistant,
+                    message_text: "Whatever".to_owned(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "Must be 3-legged.\nMust be flying.".to_owned(),
+                }],
+                expected_prompt: "Generate a picture of an elephant\nOther criteria:\n- Must be blue\n- Must be 3-legged.. Must be flying.",
+            },
+            // "Again" is ignored.
+            TestCase {
+                original_prompt: "Generate a picture of a grizzly bear",
+                messages: vec![Message {
+                    author: Author::User,
+                    message_text: "Must be blue".to_owned(),
+                },
+                Message {
+                    author: Author::Assistant,
+                    message_text: "Whatever".to_owned(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "Again".to_owned(),
+                },
+                Message {
+                    author: Author::User,
+                    message_text: "again".to_owned(),
+                }],
+                expected_prompt: "Generate a picture of a grizzly bear\nOther criteria:\n- Must be blue",
+            },
+        ];
+
+        for test_case in test_cases {
+            let actual_prompt = build(test_case.original_prompt, test_case.messages);
+
+            assert_eq!(actual_prompt, test_case.expected_prompt);
+        }
+    }
+}