Skip to main content

lit/commands/
ai.rs

1use crate::errors::LitError;
2use crate::response::CommandResponse;
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Serialize, Deserialize)]
6pub struct AiResponse {
7    pub action: String,
8    pub generated: String,
9    pub message: String,
10    pub model: Option<String>,
11}
12
13impl CommandResponse for AiResponse {
14    fn command_name(&self) -> &'static str {
15        "ai"
16    }
17    fn human_readable(&self) -> String {
18        match self.action.as_str() {
19            "commit-message" => format!("Generated commit message:\n  {}\n", self.generated),
20            "branch-name" => format!("Suggested branch name: {}\n", self.generated),
21            "pr-description" => format!("Generated PR description:\n{}\n", self.generated),
22            _ => format!("{}: {}\n", self.action, self.generated),
23        }
24    }
25}
26
27/// AI configuration stored in lit config
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct AiConfig {
30    pub provider: String,
31    pub model: String,
32    pub api_key_env: String,
33    pub endpoint: Option<String>,
34}
35
36impl Default for AiConfig {
37    fn default() -> Self {
38        AiConfig {
39            provider: "openai".to_string(),
40            model: "gpt-4o-mini".to_string(),
41            api_key_env: "LIT_AI_API_KEY".to_string(),
42            endpoint: None,
43        }
44    }
45}
46
47/// Generate a commit message from the current staged diff
48pub fn execute_commit_message(
49    context: Option<String>,
50) -> Result<AiResponse, LitError> {
51    let repo_root = crate::core::find_repo_root()?;
52
53    // Get the current diff to use as context
54    let diff_result = crate::commands::diff::execute(true, false, false, None, None)?;
55    let diff_text = serde_json::to_string(&diff_result).unwrap_or_default();
56
57    if diff_text.is_empty() || diff_text == "{}" || diff_text.contains("\"files\":[]") {
58        return Err(LitError::general(
59            "No staged changes to generate commit message from",
60        ));
61    }
62
63    // Try to call AI API (requires configured API key)
64    let config = load_ai_config(&repo_root);
65    let api_key = std::env::var(&config.api_key_env).ok();
66
67    let generated = if let Some(key) = api_key {
68        call_ai_api(
69            &config,
70            &key,
71            &format!(
72                "Generate a concise, conventional commit message for the following diff. \
73                 Use imperative mood. Keep it under 72 characters for the subject line. \
74                 {} \n\nDiff:\n{}",
75                context.as_deref().unwrap_or(""),
76                &diff_text[..diff_text.len().min(4000)]
77            ),
78        )?
79    } else {
80        // Fallback: generate a basic message from file names
81        generate_fallback_commit_message(&diff_text)
82    };
83
84    Ok(AiResponse {
85        action: "commit-message".to_string(),
86        generated,
87        message: "Commit message generated".to_string(),
88        model: Some(config.model),
89    })
90}
91
92/// Generate a branch name from a description
93pub fn execute_branch_name(description: String) -> Result<AiResponse, LitError> {
94    let repo_root = crate::core::find_repo_root()?;
95    let config = load_ai_config(&repo_root);
96    let api_key = std::env::var(&config.api_key_env).ok();
97
98    let generated = if let Some(key) = api_key {
99        call_ai_api(
100            &config,
101            &key,
102            &format!(
103                "Generate a short, kebab-case git branch name (max 50 chars) for: {}",
104                description
105            ),
106        )?
107    } else {
108        // Fallback: simple kebab-case conversion
109        description
110            .to_lowercase()
111            .replace(|c: char| !c.is_alphanumeric() && c != '-', "-")
112            .trim_matches('-')
113            .to_string()
114    };
115
116    Ok(AiResponse {
117        action: "branch-name".to_string(),
118        generated,
119        message: "Branch name generated".to_string(),
120        model: Some(config.model),
121    })
122}
123
124/// Generate a PR description from branch diff
125pub fn execute_pr_description(
126    head: Option<String>,
127    base: Option<String>,
128) -> Result<AiResponse, LitError> {
129    let repo_root = crate::core::find_repo_root()?;
130    let config = load_ai_config(&repo_root);
131    let api_key = std::env::var(&config.api_key_env).ok();
132
133    let head_ref = head.unwrap_or_else(|| {
134        crate::core::get_current_branch(&repo_root).unwrap_or_else(|_| "HEAD".to_string())
135    });
136    let base_ref = base.unwrap_or_else(|| "main".to_string());
137
138    // Get diff between branches
139    let diff_result =
140        crate::commands::diff::execute(false, false, false, Some(base_ref.clone()), Some(head_ref.clone()))?;
141    let diff_text = serde_json::to_string(&diff_result).unwrap_or_default();
142
143    let generated = if let Some(key) = api_key {
144        call_ai_api(
145            &config,
146            &key,
147            &format!(
148                "Generate a pull request description for merging '{}' into '{}'. \
149                 Include: summary, changes made, testing notes. Use markdown formatting.\n\n\
150                 Diff:\n{}",
151                head_ref,
152                base_ref,
153                &diff_text[..diff_text.len().min(4000)]
154            ),
155        )?
156    } else {
157        format!(
158            "## Summary\n\nMerge `{}` into `{}`\n\n## Changes\n\n- See diff for details\n",
159            head_ref, base_ref
160        )
161    };
162
163    Ok(AiResponse {
164        action: "pr-description".to_string(),
165        generated,
166        message: "PR description generated".to_string(),
167        model: Some(config.model),
168    })
169}
170
171fn load_ai_config(repo_root: &std::path::Path) -> AiConfig {
172    let config_path = repo_root.join(".lit").join("ai.json");
173    if config_path.exists() {
174        match std::fs::read_to_string(&config_path) {
175            Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
176            Err(_) => AiConfig::default(),
177        }
178    } else {
179        AiConfig::default()
180    }
181}
182
183fn call_ai_api(config: &AiConfig, api_key: &str, prompt: &str) -> Result<String, LitError> {
184    let endpoint = config.endpoint.as_deref().unwrap_or(match config.provider.as_str() {
185        "openai" => "https://api.openai.com/v1/chat/completions",
186        "anthropic" => "https://api.anthropic.com/v1/messages",
187        _ => "https://api.openai.com/v1/chat/completions",
188    });
189
190    let body = match config.provider.as_str() {
191        "anthropic" => serde_json::json!({
192            "model": config.model,
193            "max_tokens": 1024,
194            "messages": [{"role": "user", "content": prompt}]
195        }),
196        _ => serde_json::json!({
197            "model": config.model,
198            "messages": [
199                {"role": "system", "content": "You are a helpful assistant for version control operations. Be concise."},
200                {"role": "user", "content": prompt}
201            ],
202            "max_tokens": 1024,
203            "temperature": 0.3
204        }),
205    };
206
207    let auth_header = match config.provider.as_str() {
208        "anthropic" => ("x-api-key", api_key.to_string()),
209        _ => ("Authorization", format!("Bearer {}", api_key)),
210    };
211
212    let response = ureq::post(endpoint)
213        .set(auth_header.0, &auth_header.1)
214        .set("Content-Type", "application/json")
215        .send_string(
216            &serde_json::to_string(&body)
217                .map_err(|e| LitError::general(format!("Failed to serialize request: {}", e)))?,
218        )
219        .map_err(|e| LitError::general(format!("AI API request failed: {}", e)))?;
220
221    let response_body: serde_json::Value = response
222        .into_json()
223        .map_err(|e| LitError::general(format!("Failed to parse AI response: {}", e)))?;
224
225    // Extract text from OpenAI-style or Anthropic-style response
226    let text = response_body["choices"][0]["message"]["content"]
227        .as_str()
228        .or_else(|| response_body["content"][0]["text"].as_str())
229        .unwrap_or("Failed to generate text")
230        .trim()
231        .to_string();
232
233    Ok(text)
234}
235
236fn generate_fallback_commit_message(diff_text: &str) -> String {
237    // Parse file names from diff output
238    let files: Vec<&str> = diff_text
239        .lines()
240        .filter(|l| l.contains("\"path\"") || l.contains("\"file\""))
241        .take(5)
242        .collect();
243
244    if files.is_empty() {
245        "Update files".to_string()
246    } else if files.len() == 1 {
247        format!("Update {}", files[0].trim().replace(['"', ','], ""))
248    } else {
249        format!("Update {} files", files.len())
250    }
251}