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#[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
47pub fn execute_commit_message(
49 context: Option<String>,
50) -> Result<AiResponse, LitError> {
51 let repo_root = crate::core::find_repo_root()?;
52
53 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 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 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
92pub 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 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
124pub 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 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 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 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}