rich-prompt 0.2.0

A Rust CLI tool for rich prompts with file and directory selection.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use crate::domain::models::{ContextOutput, FileContext};
use log::{debug, info};

fn count_tokens(content: &str) -> usize {
    content.split_whitespace().count()
}

pub fn build_context_output(
    files: Vec<FileContext>,
    file_map: String,
    user_prompt: Option<String>,
) -> ContextOutput {
    debug!("Building context output from {} files", files.len());
    let mut file_contents = String::new();
    let mut total_tokens = 0;

    for file in &files {
        let tokens = count_tokens(&file.content);
        total_tokens += tokens;

        debug!("Adding file {} with {} tokens", file.path.display(), tokens);
        file_contents.push_str(&format!(
            "\nFile: {}\n```{}\n{}\n```\n",
            file.path.display(),
            file.path.extension().and_then(|e| e.to_str()).unwrap_or(""),
            file.content
        ));
    }

    let user_instructions = match user_prompt {
        Some(prompt) => {
            info!("Including user prompt in context");
            prompt
        }
        None => {
            debug!("No user prompt provided");
            String::new()
        }
    };

    ContextOutput {
        file_map,
        file_contents,
        user_instructions,
        token_count: total_tokens,
    }
}

pub fn format_output(output: &ContextOutput) -> String {
    debug!(
        "Formatting context output with {} tokens",
        output.token_count
    );
    let mut result = String::new();

    result.push_str("<file_map>\n");
    result.push_str(&output.file_map);
    result.push_str("</file_map>\n\n\n");

    result.push_str("<file_contents>");
    result.push_str(&output.file_contents);
    result.push_str("</file_contents>");

    if !output.user_instructions.is_empty() {
        result.push_str("\n\n<user_instructions>\n");
        result.push_str(&output.user_instructions);
        result.push_str("\n</user_instructions>");
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_count_tokens() {
        assert_eq!(count_tokens("hello world"), 2);
        assert_eq!(count_tokens(""), 0);
        assert_eq!(count_tokens("one\ntwo\nthree"), 3);
    }

    #[test]
    fn test_build_context_output() {
        let files = vec![
            FileContext {
                path: PathBuf::from("test/file1.rs"),
                content: "fn test() {}".to_string(),
            },
            FileContext {
                path: PathBuf::from("test/file2.rs"),
                content: "struct Test {}".to_string(),
            },
        ];

        let file_map = "test\n├── test/file1.rs\n├── test/file2.rs\n".to_string();
        let user_prompt = Some("Refactor this code".to_string());

        let output = build_context_output(files, file_map, user_prompt);

        assert_eq!(output.token_count, 6);
        assert_eq!(output.user_instructions, "Refactor this code");
        assert!(output.file_contents.contains("fn test() {}"));
        assert!(output.file_contents.contains("struct Test {}"));
    }

    #[test]
    fn test_format_output() {
        let output = ContextOutput {
            file_map: "dir1\n".to_string(),
            file_contents: "content1\n".to_string(),
            user_instructions: "prompt1".to_string(),
            token_count: 3,
        };

        let formatted = format_output(&output);

        assert!(formatted.contains("<file_map>\ndir1\n</file_map>"));
        assert!(formatted.contains("<file_contents>content1\n</file_contents>"));
        assert!(formatted.contains("<user_instructions>\nprompt1\n</user_instructions>"));
    }
}