Skip to main content

lib/utils/
input.rs

1use anyhow::{Context, Result};
2use std::io::{self, Write};
3use crate::core::prompts::generate_prompts;
4use crate::core::assimilator::Assimilator;
5use crate::adapters::llm::LLMProvider;
6
7pub fn input_use_case_option() -> Result<usize> {
8    loop {
9        println!("Please select a use case from the following options:");
10        println!("1. Creative Writing");
11        println!("2. Problem Solving");
12        println!("3. Code Problem Solving");
13        println!("4. Explanation");
14        println!("5. Custom");
15        println!("Enter the number of your choice:");
16
17        let mut user_input = String::new();
18        io::stdin().read_line(&mut user_input)
19            .context("Failed to read line")?;
20
21        match user_input.trim().parse() {
22            Ok(num) if (1..=5).contains(&num) => return Ok(num),
23            _ => {
24                println!("Invalid option. Please enter a number between 1 and 5.");
25            }
26        }
27    }
28}
29
30pub fn input_use_case_custom() -> Result<String> {
31    loop {
32        print!("Please enter the use case: ");
33        io::stdout().flush()
34            .context("Failed to flush stdout")?;
35
36        let mut user_input = String::new();
37        io::stdin().read_line(&mut user_input)
38            .context("Failed to read line")?;
39
40        let trimmed = user_input.trim();
41        if !trimmed.is_empty() {
42            return Ok(trimmed.to_string());
43        } else {
44            println!("Use case cannot be empty. Please try again.");
45        }
46    }
47}
48
49pub async fn fetch_prompts<T: LLMProvider>(assimilator: &Assimilator<T>) -> Result<Vec<String>> {
50    let use_case_option = input_use_case_option()
51        .context("Failed to get use case option")?;
52
53    let prompts = match use_case_option {
54        1..=4 => generate_prompts(use_case_option),
55        5 => {
56            let use_case = input_use_case_custom()
57                .context("Failed to get custom use case")?;
58            assimilator.tune_prompt(&use_case).await
59                .context("Failed to tune prompt")?
60        },
61        _ => unreachable!("input_use_case_option() should only return 1-5"),
62    };
63
64    Ok(prompts)
65}