vibefpmgui.window.ai_chat

  1"""Talking with the AI: Explain, the chat, and the edits it proposes.
  2
  3Findings and edit cards in the Review panel act on the editor from here.
  4"""
  5
  6from __future__ import annotations
  7
  8import re
  9from pathlib import Path
 10
 11from .. import ai_context as ctx
 12from .. import ai_edits as ed
 13from .. import ai_review as ai
 14from ..edit_cards import diff_rows
 15from ..editor import EditorTab
 16from ..result_tree import Problem
 17from ..review_monitor import Entry
 18
 19
 20class AIChatMixin:
 21    """Explain, the chat, and the AI's edits (a part of `vibefpmgui.main_window.MainWindow`)."""
 22
 23    EXPLAINABLE = ("build", "fortitude", "fortls")  # (Problems' sources: compiler, linter, language server)
 24
 25    def explain_problem(self, editor: EditorTab, problem):
 26        """Explain with AI: ask the model what a problem means and how to fix
 27        it; the answer appears in the Review panel as it's written."""
 28        if not self._can_ask("Explain with AI"):
 29            return
 30        self._explaining = (editor, problem)
 31        self._after_model = self._run_explain
 32        self._add_entry(Entry("explain", problem.message, editor.path, f"{editor.path.name}:{problem.line}",
 33                              self.ai_model))
 34        self._ensure_ollama(self._check_model)
 35
 36    def _run_explain(self):
 37        editor, problem = self._explaining
 38        if editor not in self.editors():
 39            self.review_monitor.finish(False, "The file was closed")
 40            return
 41        prompt = ai.build_explain_prompt(editor.path.name, editor.editor.toPlainText(), problem.line, problem.source,
 42                                         problem.message, problem.badge, self.effective_compiler()[0])
 43        self._explain_prompt = prompt
 44        self.review_monitor.set_status("Thinking")
 45        worker = ai.ReviewWorker(ai.DEFAULT_HOST, self.ai_model, prompt, system=ai.EXPLAIN_SYSTEM_PROMPT,
 46                                 **self._answer_options())
 47        # (Methods, not lambdas: a lambda would run in the worker's thread.)
 48        worker.chunk.connect(self.review_monitor.add_answer)
 49        worker.finished.connect(self._explain_finished)
 50        worker.failed.connect(self._explain_failed)
 51        self._start_answer(worker)
 52
 53    def _explain_finished(self, _text: str, stats: dict):
 54        entry = self._ai_entry
 55        self._remember(entry, ctx.Turn(entry.question, self._explain_prompt))
 56        detail = f"{stats.get('seconds', 0):.1f} s · {self.ai_model}" if stats.get("seconds") else self.ai_model
 57        self.review_monitor.finish(True, detail)
 58
 59    def _explain_failed(self, message: str):
 60        hint = f" (is '{self.ai_model}' downloaded? See AI → Model)" if "not found" in message.lower() else ""
 61        self.review_monitor.finish(False, message + hint)
 62
 63    def _end_explain_setup(self, why: str):
 64        if self.review_monitor.running:
 65            self.review_monitor.finish(False, why)
 66
 67    def send_chat(self, question: str, with_file: bool, kind: str = "ask"):
 68        """A question from the Review panel's box: "ask" (with ``with_file``,
 69        the open file goes along), or "edit" (a change to the open file: its
 70        edits are shown as cards to accept or reject). The model sees
 71        everything in the timeline since New Chat (see ai_context)."""
 72        editing = kind == "edit"
 73        editor = self.current_editor() if with_file or editing else None
 74        if editing and editor is None:
 75            self.statusBar().showMessage("Open the file to change in the editor first", 5000)
 76            return
 77        if not self._can_ask("AI chat"):
 78            return
 79        request, attachment = question, None
 80        if editor is not None:
 81            self._ai_editor = editor  # (the findings' lines are this file's)
 82            text = editor.editor.toPlainText()
 83            selection = None
 84            if editing and self._edit_selection is not None and self._edit_selection[0] is editor:
 85                selection = self._edit_selection[1:]
 86                first, last = selection
 87                part = "\n".join(ctx.file_lines(text)[first - 1:last])
 88                request += (f"\n\nChange only these lines (lines {first}–{last} of the file), unless the "
 89                            f"change needs more:\n```fortran\n{part}\n```")
 90            attachment = ctx.file_attachment(
 91                editor.path.name, text, numbers=not editing, context=self._chat_context(),
 92                ratio=self._token_ratio, selection=selection,
 93                cursor_line=editor.editor.textCursor().blockNumber() + 1)
 94        self._chat_turn = ctx.Turn(question, request, attachment)
 95        self._chat_kind, self._chat_editor = kind, editor
 96        self._edit_selection = None
 97        monitor = self.review_monitor
 98        monitor.set_selection_note("")
 99        monitor.chat_edit.clear()
100        self._after_model = self._run_chat
101        self._add_entry(Entry(kind, question, editor.path if editor is not None else None,
102                              editor.path.name if editor is not None else "", self.ai_model))
103        self._ensure_ollama(self._check_model)
104
105    def edit_with_ai(self, editor: EditorTab | None = None):
106        """AI → Edit with AI… (or the editor's Edit Selection with AI…): the
107        box, in Edit mode, for the selection if there is one."""
108        editor = editor or self.current_editor()
109        if editor is None:
110            self.statusBar().showMessage("Open the file to change in the editor first", 5000)
111            return
112        lines = editor.editor.selected_lines()
113        self._edit_selection = (editor, *lines) if lines else None
114        monitor = self.review_monitor
115        self.show_panel(monitor)
116        monitor.set_chat_kind("edit")
117        if lines:
118            monitor.set_selection_note(f"line {lines[0]}" if lines[0] == lines[1] else f"lines {lines[0]}–{lines[1]}")
119        else:
120            monitor.set_selection_note("")
121        monitor.chat_edit.setFocus()
122
123    def fix_with_ai(self, path, problem):
124        """Fix with AI… (a problem, or a finding in the timeline): a request
125        to address it, sent in Edit mode for its lines (shown in the box's
126        history, as if typed)."""
127        self.open_file(str(path), force_editor=True)
128        editor = next((e for e in self.editors() if e.path == Path(path).resolve()), None)
129        if editor is None:
130            return
131        first, last = problem.line, max(problem.line, problem.end_line)
132        self._edit_selection = (editor, first, last)
133        message = problem.message
134        what = {"build": "compiler message", "fortitude": f"fortitude issue {problem.badge}",
135                "fortls": "fortls message"}.get(problem.source, problem.source)
136        if problem.source == "review":  # ("AI (bug): …")
137            found = re.match(r"AI \((\w+)\):\s*(.*)", message, re.DOTALL)
138            what = f"AI review finding ({found.group(1)})" if found else "AI review finding"
139            message = found.group(2) if found else message
140        lines = f"line {first}" if first == last else f"lines {first}–{last}"
141        monitor = self.review_monitor
142        self.show_panel(monitor)
143        monitor.set_chat_kind("edit")
144        monitor.set_selection_note(lines)
145        monitor.chat_edit.setText(f"Address this {what} at {lines}: {message} "
146                                  "Change only what's needed for it.")
147        monitor.chat_edit.setFocus()
148        monitor.send_button.click()  # (if the AI is busy, the request stays in the box)
149
150    def new_chat(self):
151        """The model forgets everything so far (it stays in the timeline)."""
152        self.stop_review()
153        self._forget()
154        self.review_monitor.add_divider()
155        self.review_monitor.chat_edit.setFocus()
156
157    def clear_ai_history(self):
158        """Empty the timeline; the model forgets it too."""
159        self.stop_review()
160        self._forget()
161        self.review_monitor.clear()
162
163    def _forget(self):
164        """The model forgets everything so far (what's still running too)."""
165        self._turns = []
166        self._memory_generation += 1
167
168    def _chat_context(self) -> int:
169        """The most context a chat asks for: the model's own limit, at most
170        ai.MAX_CONTEXT."""
171        return min(self._model_contexts.get(self.ai_model) or ai.MAX_CONTEXT, ai.MAX_CONTEXT)
172
173    def _run_chat(self):
174        editing = self._chat_kind == "edit"
175        system = ai.EDIT_SYSTEM_PROMPT if editing else ai.CHAT_SYSTEM_PROMPT
176        cap = self._chat_context()
177        messages, tokens, dropped = ctx.compose(self._turns, self._chat_turn, system, cap, self._token_ratio)
178        options = self._answer_options(edit=editing)
179        room = options["num_predict"] + (ai.REASONING_TOKENS if options["think"] else 0) \
180            if options["num_predict"] else ai.REASONING_TOKENS
181        num_ctx = ai.context_for(tokens, cap, room)
182        self._chat_estimate = (tokens, len(self._turns) - dropped)
183        monitor = self.review_monitor
184        monitor.context_note = (f"context ~{tokens / 1000:.1f}k of {num_ctx // 1024}k"
185                                + (f" ({dropped} earliest left out)" if dropped else ""))
186        if tokens > cap * ctx.BUDGET * 0.9:
187            monitor.context_note += " · nearly full: New Chat starts over"
188        monitor.set_status("Thinking")
189        self._chat_keys: list = []
190        worker = ai.ReviewWorker(ai.DEFAULT_HOST, self.ai_model, messages[-1]["content"], system=system,
191                                 history=messages[1:-1], num_ctx=num_ctx,
192                                 temperature=0.1 if editing else 0.2, **options)
193        # (Methods, not lambdas: a lambda would run in the worker's thread.)
194        worker.chunk.connect(self._chat_chunk)
195        worker.finished.connect(self._chat_finished)
196        worker.failed.connect(self._chat_failed)
197        self._start_answer(worker)
198
199    def _chat_findings(self):
200        """Lines of the sent file the answer points to: listed in its entry,
201        and marked in the editor (as a review's)."""
202        entry, editor = self._ai_entry, self._chat_editor
203        if editor is None or editor not in self.editors() or self._chat_kind == "edit":
204            return
205        findings, _summary = ai.parse_findings(entry.answer, editor.editor.blockCount())
206        keys = [(f.line, f.end_line, f.category, f.comment) for f in findings]
207        if findings and keys != self._chat_keys:
208            self._chat_keys = keys
209            entry.findings = findings
210            self.review_monitor.refresh_soon(entry)
211            editor.show_review(findings)
212
213    def _chat_new_edits(self):
214        """Edit blocks the answer has finished: cards for them."""
215        entry = self._ai_entry
216        if self._chat_kind != "edit" or entry is None or entry.path is None:
217            return
218        found = ed.parse_edits(entry.answer)
219        if len(found) > len(entry.edits):
220            entry.edits += found[len(entry.edits):]
221            self._refresh_edits(entry)
222
223    def _chat_chunk(self, text: str):
224        self.review_monitor.add_answer(text)
225        self._chat_findings()
226        self._chat_new_edits()
227
228    def _chat_finished(self, _text: str, stats: dict):
229        self._chat_findings()
230        self._chat_new_edits()
231        tokens, earlier = self._chat_estimate
232        measured = stats.get("prompt_tokens") or 0
233        if measured and earlier == 0 and self._chat_turn.attachment is not None:
234            # (Only for a first question: on later ones Ollama reuses what it
235            # read before, and counts only the new part.)
236            raw = tokens / self._token_ratio
237            self._token_ratio = min(2.0, max(0.5, measured / raw))
238        self._remember(self._ai_entry, self._chat_turn)
239        detail = f"{stats.get('seconds', 0):.1f} s · {self.ai_model}" if stats.get("seconds") else self.ai_model
240        if stats.get("thinking"):
241            detail += " · the model spent its answer's length on reasoning (see AI → Answer Length)"
242        elif stats.get("truncated"):
243            detail += " · cut off at the length limit"
244        self.review_monitor.finish(True, detail)
245
246    def _chat_failed(self, message: str):
247        self._chat_new_edits()
248        if self._ai_entry is not None and self._ai_entry.answer:
249            self._remember(self._ai_entry, self._chat_turn)  # (what it had written: part of the conversation)
250        hint = f" (is '{self.ai_model}' downloaded? See AI → Model)" if "not found" in message.lower() else ""
251        self.review_monitor.finish(False, message + hint)
252
253    def _entry_editor(self, entry: Entry, open_it: bool = False) -> EditorTab | None:
254        """The tab of an entry's file (opened, with ``open_it``, if it isn't)."""
255        if entry.path is None:
256            return None
257        file = Path(entry.path).resolve()
258        editor = next((e for e in self.editors() if e.path == file), None)
259        if editor is None and open_it:
260            self.open_file(str(file), force_editor=True)
261            editor = next((e for e in self.editors() if e.path == file), None)
262        return editor
263
264    def _entry_text(self, entry: Entry) -> str | None:
265        """An entry's file as it is now: its tab's text, or the file's."""
266        editor = self._entry_editor(entry)
267        if editor is not None:
268            return editor.editor.toPlainText()
269        try:
270            return Path(entry.path).read_text(encoding="utf-8")
271        except (OSError, UnicodeDecodeError, TypeError):
272            return None
273
274    def go_to_finding(self, entry: Entry, finding):
275        editor = self._entry_editor(entry, open_it=True)
276        if editor is not None:
277            self.tabs.setCurrentWidget(editor)
278            editor.editor.go_to_line(finding.line, select=True)
279
280    def fix_finding(self, entry: Entry, finding):
281        """Fix with AI… on a finding in the timeline."""
282        if entry.path is not None:
283            self.fix_with_ai(entry.path, Problem(finding.line, 0, finding.end_line, 0, "review", "review",
284                                                 f"AI ({finding.category}): {finding.comment}", "", "review"))
285
286    def show_findings_in_editor(self, entry: Entry):
287        """An entry's findings, marked in the editor again."""
288        editor = self._entry_editor(entry, open_it=True)
289        if editor is not None:
290            self.tabs.setCurrentWidget(editor)
291            editor.show_review(entry.findings)
292
293    def _relocate(self, edit: ed.Edit, text: str | None):
294        """Find a pending edit again in the file as it is now: one that was
295        there and no longer is turns stale (and back, if the text returns)."""
296        if edit.status in (ed.APPLIED, ed.REJECTED):
297            return
298        was_there = edit.located or edit.status == ed.STALE
299        if text is None:
300            edit.start = edit.end = -1
301            edit.status = ed.STALE
302            return
303        ed.locate(text, edit)
304        if edit.located:
305            edit.status = ed.PENDING
306        elif was_there:
307            edit.status = ed.STALE
308
309    def _refresh_edits(self, entry: Entry):
310        """The edit cards of every entry about this entry's file, found again
311        in the file as it is now."""
312        for other in self.review_monitor.entries:
313            if other.edits and other.path == entry.path:
314                text = self._entry_text(other)
315                for edit in other.edits:
316                    self._relocate(edit, text)
317                self.review_monitor.refresh(other, text or "")
318
319    def accept_edit(self, entry: Entry, index: int):
320        self._make_edits(entry, [entry.edits[index]])
321
322    def accept_all_edits(self, entry: Entry):
323        self._make_edits(entry, [e for e in entry.edits if e.status == ed.PENDING])
324
325    def _make_edits(self, entry: Entry, edits: list[ed.Edit]):
326        """Make edits in the editor (one step to undo; the file is opened if
327        it isn't), each found again in the text as it is now; any no longer
328        there turn stale."""
329        editor = self._entry_editor(entry, open_it=True)
330        text = editor.editor.toPlainText() if editor is not None else None
331        for edit in edits:
332            self._relocate(edit, text)
333        if editor is None:
334            self._refresh_edits(entry)
335            return
336        located = [e for e in edits if e.located]
337        _new, made = ed.apply_all(text, located)
338        for edit in made:
339            edit.diff = diff_rows(text, edit)
340            edit.status = ed.APPLIED
341        for edit in located:
342            if edit not in made:
343                edit.status = ed.STALE  # (overlapped another one made just now)
344        if made:
345            editor.editor.replace_ranges([(e.start, e.end, e.new_text) for e in made])
346            self.statusBar().showMessage(f"Made {len(made)} edit{'s' * (len(made) != 1)} "
347                                         f"(Edit → Undo takes {'them' if len(made) > 1 else 'it'} back)", 5000)
348        self._refresh_edits(entry)
349
350    def reject_edit(self, entry: Entry, index: int):
351        self._reject(entry, [entry.edits[index]])
352
353    def reject_all_edits(self, entry: Entry):
354        self._reject(entry, [e for e in entry.edits if e.status == ed.PENDING])
355
356    def _reject(self, entry: Entry, edits: list[ed.Edit]):
357        text = self._entry_text(entry)
358        for edit in edits:
359            if edit.located and text is not None:
360                edit.diff = diff_rows(text, edit)
361            edit.status = ed.REJECTED
362        self._refresh_edits(entry)
363
364    def show_edit(self, entry: Entry, index: int):
365        """A card was clicked: show (and select) its place in the editor."""
366        edit, editor = entry.edits[index], self._entry_editor(entry, open_it=True)
367        if editor is None:
368            return
369        self._relocate(edit, editor.editor.toPlainText())
370        self.tabs.setCurrentWidget(editor)
371        if edit.located:
372            editor.editor.select_range(edit.start, edit.end)
373        elif edit.line:
374            editor.editor.go_to_line(edit.line)