vibefpmgui.window.ai

  1"""The AI menu, the local Ollama server and model, AI reviews, and the AI's worker threads.
  2
  3The Review panel (`vibefpmgui.review_monitor`) shows what comes back.
  4"""
  5
  6from __future__ import annotations
  7
  8import time
  9
 10from qtpy.QtCore import QProcess, QTimer
 11from qtpy.QtGui import QActionGroup, QKeySequence, QTextCursor
 12from qtpy.QtWidgets import QInputDialog, QMessageBox, QSizePolicy, QWidget
 13
 14from .. import ai_context as ctx
 15from .. import ai_review as ai
 16from ..background import call_in_background
 17from ..editor import EditorTab
 18from ..processes import kill_process_tree
 19from ..review_monitor import Entry, ReviewMonitor
 20
 21
 22class AIMixin:
 23    """The AI menu, Ollama and AI reviews (a part of `vibefpmgui.main_window.MainWindow`)."""
 24
 25    def _build_ai_menu(self):
 26        self.ai_model = self.settings.value("ai_model", ai.DEFAULT_MODEL)
 27        self.ai_focus = self.settings.value("ai_focus", "general")
 28        self._ai_thread = None
 29        self._ai_worker = None
 30        self._ai_editor: EditorTab | None = None
 31        self._ollama_process: QProcess | None = None
 32        # Ollama is asked things on a pool thread (see background.call_in_background):
 33        self._ai_starting = False  # a review is getting ready (server, model): AI Review waits
 34        self._known_models: list[str] | None = None  # installed models, as last heard (None: not yet)
 35        self._after_model = lambda: None  # what runs once the model is there: a review, an explanation, a chat
 36        self._turns: list[ctx.Turn] = []  # what the model remembers: every entry since New Chat (see ai_context)
 37        self._token_ratio = 1.0  # measured tokens over estimated ones (see _chat_finished)
 38        self._model_contexts: dict[str, int | None] = {}  # each model's longest context, as Ollama says
 39        self._ai_entry: Entry | None = None  # the Review panel's entry for what's running (or last ran)
 40        # What the model remembers starts over with New Chat, Clear or another
 41        # project: a request started before (stopped, but still ending) mustn't
 42        # add to it.
 43        self._memory_generation = 0
 44        self._request_generation = 0
 45        self._edit_selection: tuple[EditorTab, int, int] | None = None  # Edit Selection with AI: the lines
 46        self._asking_models = False
 47
 48        menu = self.menuBar().addMenu("&AI")
 49        self.review_action = menu.addAction("Review Current File")
 50        self.review_action.setShortcut(QKeySequence("Ctrl+Alt+R"))
 51        self.review_action.setStatusTip("Ask a local AI model (via Ollama) to review the file in the editor")
 52        self.review_action.triggered.connect(self.start_review)
 53        # Also on the toolbar, at the far right.
 54        self.review_action.setIconText("AI Review")
 55        keys = self.review_action.shortcut().toString(QKeySequence.NativeText)
 56        self.review_action.setToolTip(f"AI Review ({keys}) — review the file in the editor "
 57                                      "with a local AI model")
 58        spacer = QWidget()
 59        spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)
 60        self.command_toolbar.addWidget(spacer)
 61        self.command_toolbar.addAction(self.format_action)
 62        self.command_icons[self.format_action] = "format"
 63        self.command_toolbar.addAction(self.review_action)
 64        self.command_icons[self.review_action] = "robot"
 65        self._update_command_icons()
 66        self.stop_review_action = menu.addAction("Stop")
 67        self.stop_review_action.setEnabled(False)
 68        self.stop_review_action.triggered.connect(self.stop_review)
 69        self.edit_with_ai_action = menu.addAction("Edit with AI…")
 70        self.edit_with_ai_action.setShortcut(QKeySequence("Ctrl+Alt+E"))
 71        self.edit_with_ai_action.setStatusTip("Ask the AI to change the file in the editor (the selection, if "
 72                                              "there is one); you accept or reject each edit")
 73        self.edit_with_ai_action.triggered.connect(self.edit_with_ai)
 74        self.stop_ollama_action = menu.addAction("Stop Ollama Server")
 75        self.stop_ollama_action.setStatusTip("Stop the Ollama server the app started (and the model it has "
 76                                             "loaded, freeing its memory); it starts again when the AI is next used")
 77        self.stop_ollama_action.triggered.connect(self.stop_ollama)
 78        menu.aboutToShow.connect(self._update_ollama_action)
 79        self._update_ollama_action()
 80        self.clear_review_action = menu.addAction("Clear Review Markers")
 81        self.clear_review_action.triggered.connect(
 82            lambda: self._with_editor(lambda e: e.clear_review()))
 83        menu.addSeparator()
 84        self.model_menu = menu.addMenu("Model")
 85        self.model_menu.aboutToShow.connect(self._fill_model_menu)
 86        self.ai_think = self.settings.value("ai_think", False, type=bool)
 87        self.ai_answer_tokens = self.settings.value("ai_answer_tokens", ai.MAX_ANSWER_TOKENS, type=int)
 88        self.think_action = menu.addAction("Think Before Answering")
 89        self.think_action.setCheckable(True)
 90        self.think_action.setChecked(self.ai_think)
 91        self.think_action.setStatusTip(
 92            "Let a thinking model (qwen3, deepseek-r1…) reason before it answers: slower, sometimes better. "
 93            f"The reasoning gets {ai.REASONING_TOKENS:,} tokens of its own, on top of the answer's length")
 94        self.think_action.toggled.connect(self._set_ai_think)
 95        length_menu = self.answer_length_menu = menu.addMenu("Answer Length")
 96        lengths = QActionGroup(self)
 97        for tokens in ai.ANSWER_LENGTHS:
 98            action = length_menu.addAction(f"{tokens:,} tokens" + (" (default)" if tokens == ai.MAX_ANSWER_TOKENS
 99                                                                     else "") if tokens else "No limit")
100            action.setCheckable(True)
101            action.setChecked(tokens == self.ai_answer_tokens)
102            action.setStatusTip("How much the model may write in an answer (edits get at least "
103                                f"{ai.EDIT_ANSWER_TOKENS:,}); a longer answer needs more of its context")
104            action.triggered.connect(lambda _=False, t=tokens: self._set_ai_answer_tokens(t))
105            lengths.addAction(action)
106        focus_menu = menu.addMenu("Focus")
107        group = QActionGroup(self)
108        for key, label in ai.FOCUS_LABELS.items():
109            action = focus_menu.addAction(label)
110            action.setCheckable(True)
111            action.setChecked(key == self.ai_focus)
112            action.setStatusTip(ai.FOCUS[key])
113            action.triggered.connect(lambda _=False, k=key: self._set_ai_focus(k))
114            group.addAction(action)
115
116    def _build_review_panel(self):
117        """The Review panel: the AI's timeline."""
118        self.review_monitor = ReviewMonitor()
119        monitor = self.review_monitor
120        monitor.stop_requested.connect(self.stop_review)
121        monitor.chat_requested.connect(self.send_chat)
122        monitor.new_chat_requested.connect(self.new_chat)
123        monitor.clear_requested.connect(self.clear_ai_history)
124        monitor.selection_cleared.connect(lambda: setattr(self, "_edit_selection", None))
125        monitor.finding_selected.connect(self.go_to_finding)
126        monitor.fix_requested.connect(self.fix_finding)
127        monitor.show_in_editor_requested.connect(self.show_findings_in_editor)
128        monitor.accept_requested.connect(self.accept_edit)
129        monitor.reject_requested.connect(self.reject_edit)
130        monitor.accept_all_requested.connect(self.accept_all_edits)
131        monitor.reject_all_requested.connect(self.reject_all_edits)
132        monitor.edit_clicked.connect(self.show_edit)
133
134    def _update_review_actions(self):
135        """AI review needs an editor tab in front (and no review running)."""
136        has_editor = self.current_editor() is not None
137        busy = self._ai_thread is not None or getattr(self, "_ai_starting", False)
138        self.review_action.setEnabled(has_editor and not busy)
139        self.review_monitor.set_sending(not busy)  # (a question while it's answering: wait for it)
140        self.clear_review_action.setEnabled(has_editor)
141
142    def _set_ai_think(self, on: bool):
143        self.ai_think = on
144        self.settings.setValue("ai_think", on)
145
146    def _set_ai_answer_tokens(self, tokens: int):
147        self.ai_answer_tokens = tokens
148        self.settings.setValue("ai_answer_tokens", tokens)
149
150    def _answer_options(self, edit: bool = False) -> dict:
151        """The worker's settings from AI → Think Before Answering and Answer
152        Length (an edit's answer, with its old and new text: at least
153        EDIT_ANSWER_TOKENS)."""
154        tokens = self.ai_answer_tokens
155        if edit and tokens:
156            tokens = max(tokens, ai.EDIT_ANSWER_TOKENS)
157        return {"num_predict": tokens, "think": self.ai_think}
158
159    def _start_answer(self, worker) -> None:
160        """Start a worker whose reasoning (if the model thinks) goes to the
161        running entry."""
162        worker.reasoning.connect(self.review_monitor.add_reasoning)
163        self._start_ai_worker(worker)
164
165    def _set_ai_focus(self, key: str):
166        self.ai_focus = key
167        self.settings.setValue("ai_focus", key)
168
169    def _set_ai_model(self, name: str):
170        self.ai_model = name
171        self.settings.setValue("ai_model", name)
172        self.statusBar().showMessage(f"AI model: {name}", 4000)
173
174    def _fill_model_menu(self):
175        """Installed models (if Ollama is running), plus the current choice:
176        the list as last heard, at once, while Ollama is asked again."""
177        self._show_model_menu()
178        if ai.ollama_available() and not self._asking_models:
179            self._asking_models = True
180            call_in_background(ai.installed_models, self._models_heard, self)
181
182    def _models_heard(self, models: list[str]):
183        self._asking_models = False
184        self._known_models = models
185        if self.model_menu.isVisible():
186            self._show_model_menu()
187
188    def _show_model_menu(self):
189        menu = self.model_menu
190        menu.clear()
191        models = self._known_models or []
192        names = sorted(set(models) | {self.ai_model})
193        group = QActionGroup(menu)
194        for name in names:
195            label = name if name in models else f"{name} (not downloaded yet)"
196            action = menu.addAction(label)
197            action.setCheckable(True)
198            action.setChecked(name == self.ai_model)
199            action.triggered.connect(lambda _=False, n=name: self._set_ai_model(n))
200            group.addAction(action)
201        if self._known_models is None and ai.ollama_available():
202            note = menu.addAction("(Checking for installed models…)")
203            note.setEnabled(False)
204        elif not models:
205            note = menu.addAction("(Ollama isn't running — installed models appear once it is)")
206            note.setEnabled(False)
207        menu.addSeparator()
208        other = menu.addAction("Other Model…")
209        other.triggered.connect(self._choose_other_model)
210
211    def _choose_other_model(self):
212        name, ok = QInputDialog.getText(
213            self, "AI model", "Ollama model name (e.g. qwen2.5-coder:14b, qwen3-coder:30b, "
214            "gpt-oss:20b). It's downloaded the first time it's used:", text=self.ai_model)
215        if ok and name.strip():
216            self._set_ai_model(name.strip())
217
218    def start_review(self):
219        editor = self.current_editor()
220        if editor is None:
221            QMessageBox.information(self, "AI review", "Open a file in the editor first "
222                                    "(double-click a file, or File → Open File in Editor).")
223            return
224        if not self._can_ask("AI review"):
225            return
226        self._ai_editor = editor
227        self._after_model = self._run_review
228        focus = ai.FOCUS_LABELS.get(self.ai_focus, self.ai_focus)
229        self._add_entry(Entry("review", focus, editor.path, editor.path.name, self.ai_model))
230        self._ensure_ollama(self._check_model)
231
232    def _can_ask(self, title: str) -> bool:
233        """The AI is free, and Ollama's Python package is there (if not: why)."""
234        if self._ai_thread is not None or self._ai_starting:
235            self.statusBar().showMessage("The AI is busy: try again when it's done", 5000)
236            return False
237        if not ai.ollama_available():
238            QMessageBox.warning(self, title, "The 'ollama' Python package isn't installed.\n"
239                                "Run 'pixi install' (or 'pip install ollama').")
240            return False
241        return True
242
243    def _add_entry(self, entry: Entry) -> Entry:
244        """A new entry in the Review panel's timeline, for what's starting."""
245        self._ai_entry = entry
246        self._request_generation = self._memory_generation
247        self.show_panel(self.review_monitor)
248        if entry.kind not in ("ask", "edit"):
249            self.review_monitor.context_note = ""
250        return self.review_monitor.add_entry(entry)
251
252    def _remember(self, entry: Entry, turn: ctx.Turn):
253        """What the model was asked and answered: it's part of what later
254        questions send (see ai_context)."""
255        if self._request_generation != self._memory_generation:
256            return  # (asked before New Chat, Clear or another project: forgotten)
257        turn.answer, turn.edits = entry.answer, entry.edits
258        entry.turn = turn
259        self._turns.append(turn)
260
261    def _set_ai_starting(self, on: bool):
262        """A review is getting ready (asking Ollama, starting it): AI Review
263        waits, so a second click doesn't start a second server."""
264        self._ai_starting = on
265        self._update_review_actions()
266
267    def _ensure_ollama(self, then):
268        """Start "ollama serve" if no server answers, then call ``then``
269        (Ollama is asked on a pool thread: the window never waits for it)."""
270        self._set_ai_starting(True)
271        call_in_background(ai.server_running, lambda up: self._server_checked(up, then), self)
272
273    def _server_checked(self, up: bool, then):
274        if up:
275            then()
276            return
277        exe = ai.ollama_executable()
278        if exe is None:
279            self._set_ai_starting(False)
280            self._end_explain_setup("Ollama isn't running")
281            QMessageBox.warning(self, "AI review", "Ollama isn't running and the 'ollama' program "
282                                "wasn't found. Run 'pixi install', or install Ollama from ollama.com.")
283            return
284        self._ai_status("starting the local Ollama server", "Starting Ollama")
285        self._ollama_process = QProcess(self)
286        self._ollama_process.setProcessChannelMode(QProcess.MergedChannels)
287        self._ollama_process.start(exe, ["serve"])
288        deadline = time.monotonic() + 20
289
290        def poll():
291            call_in_background(ai.server_running, polled, self)
292
293        def polled(up: bool):
294            if up:
295                then()
296            elif time.monotonic() > deadline or self._ollama_process.state() == QProcess.NotRunning:
297                self._set_ai_starting(False)
298                self.review_monitor.finish(False, "Could not start the Ollama server.")
299            else:
300                QTimer.singleShot(300, self, poll)
301
302        QTimer.singleShot(300, self, poll)
303
304    def _check_model(self):
305        """Is the model downloaded? (Asked on a pool thread.)"""
306        call_in_background(ai.installed_models, self._model_checked, self)
307
308    def _model_checked(self, installed: list[str]):
309        self._known_models = installed
310        self._set_ai_starting(False)  # (from here a worker thread, or nothing, follows)
311        model = self.ai_model if ":" in self.ai_model else self.ai_model + ":latest"
312        if model in installed or self.ai_model in installed:
313            if self.ai_model not in self._model_contexts:  # (for the next chat question: asked once)
314                self._model_contexts[self.ai_model] = None
315                call_in_background(lambda m=self.ai_model: ai.model_context_length(m),
316                                   lambda n, m=self.ai_model: self._model_contexts.__setitem__(m, n), self)
317            self._after_model()
318            return
319        answer = QMessageBox.question(
320            self, "Download AI model",
321            f"The model '{self.ai_model}' isn't downloaded yet. Download it now?\n\n"
322            "Models are a few GB; this happens once, and everything runs locally afterwards.")
323        if answer != QMessageBox.Yes:
324            self._end_explain_setup(f"{self.ai_model} isn't downloaded")
325            return
326        self._ai_status(f"downloading {self.ai_model}", "Downloading")
327        worker = ai.ReviewWorker(ai.DEFAULT_HOST, self.ai_model, pull_only=True)
328        worker.progress.connect(self._pull_progress)
329        # (A method of the window, not a lambda: a lambda would run in the
330        # worker's thread, which is ending, and the review would never start.)
331        worker.finished.connect(self._pull_finished)
332        worker.failed.connect(self._review_failed)
333        self._start_ai_worker(worker)
334
335    def _ai_status(self, description: str, status: str):
336        """Getting ready (starting Ollama, downloading the model): the status
337        of the entry that's waiting for it."""
338        self.show_panel(self.review_monitor)
339        if self.review_monitor.running:
340            self.review_monitor.set_status(status)
341
342    def _pull_finished(self, *_):
343        self._ai_done(then=self._after_model)
344
345    def _pull_progress(self, status: str, fraction: float):
346        if fraction >= 0:
347            self.review_monitor.set_status(f"Downloading {self.ai_model} {fraction * 100:.0f}%", fraction)
348        else:
349            self.review_monitor.set_status(status.capitalize() or "Downloading")
350
351    def _run_review(self):
352        editor = self._ai_editor
353        if editor is None or editor not in self.editors():
354            self.review_monitor.finish(False, "The file was closed")
355            return
356        code = editor.editor.toPlainText()
357        diagnostics = editor.compiler_diagnostics()
358        lines = ctx.file_lines(code)
359        options = self._answer_options()
360        cap = self._chat_context()
361        room = options["num_predict"] + (ai.REASONING_TOKENS if options["think"] else 0) \
362            if options["num_predict"] else ai.REASONING_TOKENS
363
364        def prompt_for(first: int, last: int) -> str:
365            return ai.build_prompt(editor.path.name, editor.highlighter.lexer.name, "\n".join(lines[first - 1:last]),
366                                   self.ai_focus, diagnostics, first, len(lines))
367
368        # The selection (two lines or more), or the whole file if it fits the
369        # model's context with room for the answer; else the procedure at the
370        # cursor. (Sent anyway, a file too long loses its start, instructions
371        # included: Ollama doesn't refuse, it drops what doesn't fit.)
372        first, last, why = 1, len(lines), ""
373        selection = editor.editor.selected_lines()
374        if selection is not None and selection[1] > selection[0]:
375            (first, last), why = selection, "the selection"
376        elif ctx.estimate_tokens(prompt_for(1, len(lines)) + ai.SYSTEM_PROMPT, self._token_ratio) + room \
377                > cap * 0.9:
378            first, last = ctx.window(lines, None, editor.editor.textCursor().blockNumber() + 1)
379            why = "the procedure at the cursor: the file is too long to review at once"
380        prompt = prompt_for(first, last)
381        entry = self._ai_entry
382        if why:
383            entry.where = f"{editor.path.name}:{first}–{last}"
384            entry.question = f"{entry.question} · {why}"
385            self.review_monitor.refresh(entry)
386        num_ctx = ai.context_for(ctx.estimate_tokens(prompt + ai.SYSTEM_PROMPT, self._token_ratio), cap, room)
387        self.review_monitor.set_status("Thinking")
388        self._show_single_output(f"AI review of {editor.path.name} ({self.ai_model})", "")
389        self._review_text = ""
390        self._review_keys = []
391        self._review_code = code
392        editor.clear_review()  # markers from an earlier review would be misleading
393        worker = ai.ReviewWorker(ai.DEFAULT_HOST, self.ai_model, prompt, num_ctx=num_ctx, **options)
394        worker.chunk.connect(self._review_chunk)
395        worker.finished.connect(self._review_finished)
396        worker.failed.connect(self._review_failed)
397        self._start_answer(worker)
398
399    def _start_ai_worker(self, worker):
400        self._ai_worker = worker
401        self._ai_thread = ai.start_worker(worker, self)
402        self._ai_thread.finished.connect(self._ai_thread_done)
403        self._update_review_actions()
404        self.stop_review_action.setEnabled(True)
405        self.review_monitor.set_stoppable(True)
406
407    def _ai_thread_done(self):
408        self._ai_thread = None
409        self._ai_worker = None
410        self._update_review_actions()
411        self.stop_review_action.setEnabled(False)
412        self.review_monitor.set_stoppable(False)
413
414    def _ai_done(self, then=None):
415        # Let the download thread finish before starting the review thread.
416        if then is not None:
417            QTimer.singleShot(100, self, lambda: self._after_ai_thread(then))
418
419    def _after_ai_thread(self, then):
420        if self._ai_thread is not None:
421            QTimer.singleShot(100, self, lambda: self._after_ai_thread(then))
422        else:
423            then()
424
425    def _review_chunk(self, text: str):
426        cursor = self.test_output.textCursor()
427        cursor.movePosition(QTextCursor.End)
428        cursor.insertText(text)
429        self.test_output.setTextCursor(cursor)
430        self.test_output.ensureCursorVisible()
431        self.review_monitor.add_answer(text)
432        # Text arrives a line at a time: re-read the answer so far and show
433        # the findings found so far (the entry's list and the editor's markers).
434        self._review_text += text
435        editor, entry = self._ai_editor, self._ai_entry
436        max_line = editor.editor.blockCount() if editor in self.editors() else None
437        findings, _summary = ai.parse_findings(self._review_text, max_line)
438        keys = [(f.line, f.end_line, f.category, f.comment) for f in findings]
439        if keys != self._review_keys:
440            self._review_keys = keys
441            entry.findings = findings
442            self.review_monitor.refresh_soon(entry)  # (with the answer's next showing)
443            if editor in self.editors():
444                editor.show_review(findings)
445
446    def _review_finished(self, text: str, stats: dict):
447        editor, entry = self._ai_editor, self._ai_entry
448        max_line = editor.editor.blockCount() if editor in self.editors() else None
449        findings, summary = ai.parse_findings(text, max_line)
450        entry.findings, entry.summary = findings, summary
451        name = editor.path.name if editor is not None else entry.where.split(":")[0]
452        self._remember(entry, ctx.Turn(
453            f"Review {entry.where}", f"Review my file {entry.where} ({entry.question.lower()}): list the issues "
454            "you find.",
455            ctx.file_attachment(name, self._review_code, numbers=True, context=self._chat_context(),
456                                ratio=self._token_ratio)))
457        detail = f"{stats.get('seconds', 0):.1f} s · {self.ai_model}" if stats else self.ai_model
458        if stats.get("repeated"):
459            detail += " · stopped a repetition loop"
460        elif stats.get("thinking"):
461            detail += " · the model spent its answer's length on reasoning (see AI → Answer Length)"
462        elif stats.get("truncated"):
463            detail += " · answer cut off at the length limit"
464        self.review_monitor.finish(True, detail)
465        if editor in self.editors():
466            editor.show_review(findings)
467        self.statusBar().showMessage(f"AI review: {len(findings)} findings", 5000)
468
469    def _review_failed(self, message: str):
470        hint = ""
471        if "not found" in message.lower():
472            hint = f" (is '{self.ai_model}' downloaded? See AI → Model)"
473        self.review_monitor.finish(False, message + hint)
474        self.statusBar().showMessage(f"AI: {message}{hint}", 6000)
475
476    def stop_review(self):
477        if self._ai_worker is not None:
478            self._ai_worker.cancel()
479
480    def _ollama_running(self) -> bool:
481        """A server the app started is running (one started otherwise isn't ours to stop)."""
482        return self._ollama_process is not None and self._ollama_process.state() != QProcess.NotRunning
483
484    def _update_ollama_action(self):
485        self.stop_ollama_action.setEnabled(self._ollama_running())
486
487    def stop_ollama(self, quitting: bool = False):
488        """AI → Stop Ollama Server: the server the app started, with the
489        model runner it started (killing only the server would leave that
490        running, holding the model's memory)."""
491        if not self._ollama_running():
492            return
493        if not quitting:
494            self.stop_review()
495        kill_process_tree(self._ollama_process)
496        self._ollama_process.waitForFinished(3000)
497        self._ollama_process = None
498        self._update_ollama_action()
499        if not quitting:
500            self.statusBar().showMessage("Stopped the Ollama server (it starts again when the AI is next used)", 5000)
501
502    def _shutdown_ai(self):
503        if self._ai_worker is not None:
504            self._ai_worker.cancel()
505        if self._ai_thread is not None:
506            self._ai_thread.quit()
507            self._ai_thread.wait(3000)
508        self.stop_ollama(quitting=True)