fix_engine/
llm_client.rs

1//! LLM client for AI-assisted fix generation.
2//!
3//! Sends code snippets + violation messages to an LLM endpoint and parses
4//! the response into text edits.
5
6use anyhow::Result;
7use fix_engine_core::{FixConfidence, FixSource, LlmFixRequest, PlannedFix, TextEdit};
8use serde::{Deserialize, Serialize};
9
10use crate::context::FixContext;
11
12/// An OpenAI-compatible chat completion request.
13#[derive(Serialize)]
14struct ChatRequest {
15    model: String,
16    messages: Vec<ChatMessage>,
17    temperature: f32,
18}
19
20#[derive(Serialize)]
21struct ChatMessage {
22    role: String,
23    content: String,
24}
25
26/// An OpenAI-compatible chat completion response.
27#[derive(Deserialize)]
28struct ChatResponse {
29    choices: Vec<ChatChoice>,
30}
31
32#[derive(Deserialize)]
33struct ChatChoice {
34    message: ChatChoiceMessage,
35}
36
37#[derive(Deserialize)]
38struct ChatChoiceMessage {
39    content: String,
40}
41
42/// Send an LLM fix request and return planned fixes.
43pub async fn request_llm_fix(
44    endpoint: &str,
45    request: &LlmFixRequest,
46    ctx: &dyn FixContext,
47) -> Result<Vec<PlannedFix>> {
48    // Read the source file for full context
49    let source = std::fs::read_to_string(&request.file_path)?;
50
51    let system_prompt = ctx.llm_system_prompt();
52
53    let user_prompt = format!(
54        "File: {}\nLine: {}\n\nMigration rule: {}\n\nMessage: {}\n\nFull file source:\n```\n{}\n```",
55        request.file_path.display(),
56        request.line,
57        request.rule_id,
58        request.message,
59        source,
60    );
61
62    let chat_request = ChatRequest {
63        model: "gpt-4".to_string(),
64        messages: vec![
65            ChatMessage {
66                role: "system".to_string(),
67                content: system_prompt,
68            },
69            ChatMessage {
70                role: "user".to_string(),
71                content: user_prompt,
72            },
73        ],
74        temperature: 0.0,
75    };
76
77    let client = reqwest::Client::new();
78    let response = client
79        .post(endpoint)
80        .json(&chat_request)
81        .send()
82        .await?
83        .json::<ChatResponse>()
84        .await?;
85
86    let content = response
87        .choices
88        .first()
89        .map(|c| c.message.content.as_str())
90        .unwrap_or("");
91
92    let edits = parse_llm_fix_response(content, &request.rule_id);
93
94    if edits.is_empty() {
95        return Ok(Vec::new());
96    }
97
98    Ok(vec![PlannedFix {
99        edits,
100        confidence: FixConfidence::Medium,
101        source: FixSource::Llm,
102        rule_id: request.rule_id.clone(),
103        file_uri: request.file_uri.clone(),
104        line: request.line,
105        description: format!("LLM-generated fix for {}", request.rule_id),
106    }])
107}
108
109/// Parse the LLM response format into text edits.
110fn parse_llm_fix_response(content: &str, rule_id: &str) -> Vec<TextEdit> {
111    let mut edits = Vec::new();
112    let mut in_fix_block = false;
113    let mut current_line: Option<u32> = None;
114    let mut current_old: Option<String> = None;
115
116    for line in content.lines() {
117        let trimmed = line.trim();
118
119        if trimmed == "```fix" {
120            in_fix_block = true;
121            continue;
122        }
123        if trimmed == "```" && in_fix_block {
124            in_fix_block = false;
125            continue;
126        }
127
128        if !in_fix_block {
129            continue;
130        }
131
132        if let Some(rest) = trimmed.strip_prefix("LINE:") {
133            current_line = rest.trim().parse().ok();
134        } else if let Some(rest) = trimmed.strip_prefix("OLD:") {
135            current_old = Some(rest.to_string());
136        } else if let Some(rest) = trimmed.strip_prefix("NEW:") {
137            if let (Some(line_num), Some(old_text)) = (current_line, current_old.take()) {
138                edits.push(TextEdit {
139                    line: line_num,
140                    old_text,
141                    new_text: rest.to_string(),
142                    rule_id: rule_id.to_string(),
143                    description: "LLM-generated fix".to_string(),
144                    replace_all: false,
145                });
146            }
147            current_line = None;
148        }
149    }
150
151    edits
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_parse_llm_response() {
160        let response = r#"
161```fix
162LINE:56
163OLD:<BarsIcon />
164NEW:<PageToggleButton isHamburgerButton />
165```
166
167```fix
168LINE:10
169OLD:import { Button, BarsIcon } from '@patternfly/react-core';
170NEW:import { PageToggleButton } from '@patternfly/react-core';
171```
172"#;
173        let edits = parse_llm_fix_response(response, "test-rule");
174        assert_eq!(edits.len(), 2);
175        assert_eq!(edits[0].line, 56);
176        assert_eq!(edits[0].old_text, "<BarsIcon />");
177        assert_eq!(edits[0].new_text, "<PageToggleButton isHamburgerButton />");
178        assert_eq!(edits[1].line, 10);
179    }
180
181    #[test]
182    fn test_parse_llm_response_empty_input() {
183        let edits = parse_llm_fix_response("", "rule-1");
184        assert!(edits.is_empty());
185    }
186
187    #[test]
188    fn test_parse_llm_response_no_fix_blocks() {
189        let response = "Here is some explanation text without any fix blocks.";
190        let edits = parse_llm_fix_response(response, "rule-1");
191        assert!(edits.is_empty());
192    }
193
194    #[test]
195    fn test_parse_llm_response_non_fix_code_blocks_ignored() {
196        let response = r#"
197```typescript
198const x = 1;
199```
200
201```javascript
202console.log("hello");
203```
204"#;
205        let edits = parse_llm_fix_response(response, "rule-1");
206        assert!(edits.is_empty());
207    }
208
209    #[test]
210    fn test_parse_llm_response_single_fix() {
211        let response = r#"
212```fix
213LINE:1
214OLD:import { Chip } from '@patternfly/react-core';
215NEW:import { Label } from '@patternfly/react-core';
216```
217"#;
218        let edits = parse_llm_fix_response(response, "rename-rule");
219        assert_eq!(edits.len(), 1);
220        assert_eq!(edits[0].line, 1);
221        assert_eq!(
222            edits[0].old_text,
223            "import { Chip } from '@patternfly/react-core';"
224        );
225        assert_eq!(
226            edits[0].new_text,
227            "import { Label } from '@patternfly/react-core';"
228        );
229        assert_eq!(edits[0].rule_id, "rename-rule");
230    }
231
232    #[test]
233    fn test_parse_llm_response_incomplete_fix_block_skipped() {
234        // Missing NEW: line — should not produce an edit
235        let response = r#"
236```fix
237LINE:5
238OLD:something
239```
240"#;
241        let edits = parse_llm_fix_response(response, "rule-1");
242        assert!(edits.is_empty());
243    }
244
245    #[test]
246    fn test_parse_llm_response_missing_line_skipped() {
247        // Has OLD and NEW but no LINE: — should not produce an edit
248        let response = r#"
249```fix
250OLD:old text
251NEW:new text
252```
253"#;
254        let edits = parse_llm_fix_response(response, "rule-1");
255        assert!(edits.is_empty());
256    }
257
258    #[test]
259    fn test_parse_llm_response_whitespace_tolerance() {
260        let response = r#"
261```fix
262LINE:  42
263OLD:  <Chip />
264NEW:  <Label />
265```
266"#;
267        let edits = parse_llm_fix_response(response, "rule-1");
268        assert_eq!(edits.len(), 1);
269        assert_eq!(edits[0].line, 42);
270        // OLD/NEW preserve the text after the prefix
271        assert_eq!(edits[0].old_text, "  <Chip />");
272        assert_eq!(edits[0].new_text, "  <Label />");
273    }
274
275    #[test]
276    fn test_parse_llm_response_new_can_be_empty() {
277        // Removing a line entirely — NEW is empty
278        let response = r#"
279```fix
280LINE:10
281OLD:  isHidden={true}
282NEW:
283```
284"#;
285        let edits = parse_llm_fix_response(response, "rule-1");
286        assert_eq!(edits.len(), 1);
287        assert_eq!(edits[0].old_text, "  isHidden={true}");
288        assert_eq!(edits[0].new_text, "");
289    }
290}
fix_engine/llm_client.rs

fix_engine/
llm_client.rs
