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