luckyshot 0.0.5

A CLI tool for one-shot code generations using RAG and file watching
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
125
126
127
128
129

use clap::{Parser, Subcommand};
use dotenvy::dotenv;
use std::env;
use std::io::Read;

mod files;
mod openai;
mod scan;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Scan files matching a glob pattern
    Scan {
        /// The glob pattern to match files
        #[arg(required = true)]
        pattern: String,

        /// Size of chunks to split files into (0 for no chunking)
        #[arg(long, default_value = "0")]
        chunk_size: usize,

        /// Size of overlap between chunks (0 for no overlap)
        #[arg(long, default_value = "0")]
        overlap_size: usize,

        /// Include file metadata in embeddings
        #[arg(long, default_value = "false")]
        embed_metadata: bool,
    },

    /// Ask a question about the codebase
    #[command(trailing_var_arg = true)]
    Ask {
        /// The question to ask (optional if using stdin)
        #[arg(trailing_var_arg = true)]
        prompt: Vec<String>,
    },

    /// Expand a prompt using a system prompt
    Expand {
        /// System prompt for expanding the question
        #[arg(required = true)]
        system_prompt: String,

        /// The prompt to expand (optional if using stdin)
        #[arg(trailing_var_arg = true)]
        prompt: Vec<String>,
    },
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not found in environment");

    let cli = Cli::parse();

    match cli.command {
        Commands::Scan {
            pattern,
            chunk_size,
            overlap_size,
            embed_metadata,
        } => {
            if chunk_size > 0 && overlap_size >= chunk_size {
                eprintln!("Error: overlap_size must be less than chunk_size");
                std::process::exit(1);
            }
            scan::scan_files(&pattern, &api_key, chunk_size, overlap_size, embed_metadata).await?;
        }
        Commands::Ask { prompt } => {
            let prompt_text = if prompt.is_empty() {
                let mut buffer = String::new();
                std::io::stdin().read_to_string(&mut buffer)?;
                buffer.trim().to_string()
            } else {
                prompt.join(" ")
            };

            if prompt_text.is_empty() {
                eprintln!("Error: No input provided via arguments or stdin");
                std::process::exit(1);
            }

            match openai::get_embedding(&prompt_text, &api_key).await {
                Ok(embedding) => {
                    let _related_files = openai::find_related_files(embedding).await;
                }
                Err(e) => {
                    eprintln!("Error getting embedding: {}", e);
                }
            }
        }
        Commands::Expand {
            prompt,
            system_prompt,
        } => {
            let prompt_text = if prompt.is_empty() {
                // Read from stdin if no prompt arguments provided
                let mut buffer = String::new();
                std::io::stdin().read_line(&mut buffer)?;
                buffer
            } else {
                prompt.join(" ")
            };

            if !prompt_text.trim().is_empty() {
                match openai::get_openai_chat_completion(&prompt_text, &system_prompt, &api_key)
                    .await
                {
                    Ok(expanded) => println!("{}", expanded),
                    Err(e) => eprintln!("Error expanding prompt: {}", e),
                }
            } else {
                eprintln!("Error: No prompt provided via arguments or stdin");
                std::process::exit(1);
            }
        }
    }
    Ok(())
}