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(context: Option<String>) -> Result<AiResponse, LitError> {
49 let repo_root = crate::core::find_repo_root()?;
50
51 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 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 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
90pub 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 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
122pub 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 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 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 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}