diff --git a/.gitignore b/.gitignore index 894a44cc..5e0cf157 100644 --- a/.gitignore +++ b/.gitignore @@ -89,7 +89,7 @@ venv/ ENV/ env.bak/ venv.bak/ - +cqvnv # Spyder project settings .spyderproject .spyproject diff --git a/README.md b/README.md index 4686bdf8..3dcda250 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,101 @@ Additional screenshots are available in [the wiki](https://github.com/CadQuery/C * Export to various formats * STL * STEP +* AI Chat Assistant *(opt-in)* — generate and iterate on CadQuery models using natural language via any OpenAI-compatible API + +## AI Chat Assistant + +An optional dockable panel that lets you describe what you want to model and receive a complete, runnable CadQuery script in response. The current editor script is automatically sent as context with every message, so the model edits your existing code rather than starting from scratch. + +The panel is disabled by default. No network requests are made unless you explicitly enable it and send a prompt. + +### Privacy + +Every prompt you send includes your current editor script and is transmitted to the API endpoint you configure. Before sending the first message, the panel will show a one-time confirmation dialog that lists the endpoint. You can change or review the endpoint at any time in **Edit -> Preferences -> AI Assistant**. + +The API key is stored in the system keychain if the `keyring` package is installed. Otherwise it falls back to plaintext in the local preferences file. Install `keyring` for secure storage: + +```bash +pip install keyring +``` + +### Installation + +The panel requires the `openai` package and optionally `keyring` for secure key storage. You can install all optional AI Assistant dependencies in one command using the project's standard extras group: + +```bash +pip install .[ai] +``` + +Or install them manually: + +```bash +pip install openai keyring +``` + +### Configuration + +Open **Edit -> Preferences -> AI Assistant** and set the following: + +| Setting | Description | Default | +|---|---|---| +| Enabled | Show the AI Assistant dock panel | false | +| Provider / Base URL | OpenAI-compatible API endpoint | `https://api.openai.com/v1` | +| Model | Model identifier | `gpt-4o` | +| API Key | Provider API key | *(empty)* | +| Auto-run after insert | Re-render the model after inserting code | true | + +**Supported providers** + +| Provider | Base URL | +|---|---| +| OpenAI | `https://api.openai.com/v1` | +| OpenRouter | `https://openrouter.ai/api/v1` | +| Ollama (local) | `http://localhost:11434/v1` | +| Any OpenAI-compatible API | your custom endpoint | + +Once Enabled is checked, the panel appears as a dock on the right side. It can also be toggled from **Tools -> AI Assistant** or the **View** menu. + +### Usage + +Type a description in the prompt box and press Enter or click Send. The assistant always replies with a complete `python` fenced code block. Once a response arrives, click **Insert & Run** to load the code into the editor and re-render the model. + +If the response contains no code block — for example, the model returned a clarifying question — the Insert button stays disabled and an informational message is shown in the chat. + +**Generating a new model** + +``` +Create a hollow cylinder, 40mm outer diameter, 30mm inner diameter, 80mm tall +``` + +**Iterating on an existing model** + +``` +Add a 2mm fillet to all vertical edges +``` + +``` +Change the wall thickness to 5mm and add M3 mounting holes at each corner +``` + +**Fixing errors (One-Tap Self-Healing)** + +If the rendered script throws an error, the traceback appears in the **Current traceback** panel. You don't need to copy and paste anything: + +* Simply click the vibrant **✨ Auto-Fix with AI** button at the top of the traceback pane. +* The editor will instantly capture the traceback, line numbers, and crash snippet, open the AI Assistant, and request a fix. +* Once the corrected code is returned, it is **automatically inserted back into the editor and re-rendered in one tap!** + +**Resetting the conversation** + +Click **Clear Chat** to discard the conversation history and start a new session. The privacy consent is also reset, so the notice will appear again on the next send. + +### Notes + +* **Model Quality Matters**: For the best results, use high-capability frontier models. Larger models are vastly superior at precise geometric reasoning, spatial awareness, and successfully resolving CadQuery API errors compared to smaller or lightweight local models. +* Conversation history is capped at 10 turn pairs. Older messages are dropped automatically to limit token usage. +* Models that produce consistent fenced code blocks work best. +* The panel shuts down its background thread cleanly when CQ-editor closes, even if a request is in progress. ## Documentation diff --git a/cq_editor/main_window.py b/cq_editor/main_window.py index cde9d2ac..428ecf42 100644 --- a/cq_editor/main_window.py +++ b/cq_editor/main_window.py @@ -12,6 +12,7 @@ QAction, QApplication, QMenu, + QDialog, ) from logbook import Logger import cadquery as cq @@ -26,6 +27,13 @@ from .widgets.log import LogViewer from .widgets.upload_dialog import UploadDialog +# AI Assistant is optional — CQ-editor works normally without it +try: + from .widgets.ai_chat import AIChatWidget + _AI_AVAILABLE = True +except ImportError: + _AI_AVAILABLE = False + from . import __version__ from .utils import ( dock, @@ -120,6 +128,14 @@ def __init__(self, parent=None, filename=None): self.preferences.sigTreeStateChanged.connect(self.preferencesChanged) self.restorePreferences() + + # Ensure AI Assistant visibility matches preferences on startup + if _AI_AVAILABLE and "ai_chat" in self.components: + dock_widget = self.docks.get("ai_chat") + if dock_widget: + enabled = self.components["ai_chat"]._pref("Enabled") + dock_widget.setVisible(enabled) + self.restoreWindow() # Handle the event of the editor being hidden or shown @@ -288,11 +304,20 @@ def prepare_panes(self): lambda c: dock(c, "Log viewer", self, defaultArea="bottom"), ) + # ---- AI Chat Assistant (opt-in; only registered if module loaded) ---- + if _AI_AVAILABLE: + self.registerComponent( + "ai_chat", + AIChatWidget(self), + lambda c: dock(c, "AI Assistant", self, defaultArea="right"), + ) + # ---------------------------------------------------------------------- + for d in self.docks.values(): d.show() PRINT_REDIRECTOR.sigStdoutWrite.connect( - lambda text: self.components["log"].append(text) + self.components["log"].append ) def prepare_menubar(self): @@ -401,6 +426,14 @@ def prepare_menubar(self): ) ) + # AI Assistant toggle (only if the widget was loaded successfully) + if _AI_AVAILABLE: + menu_tools.addSeparator() + ai_action = QAction("\U0001F916 AI Assistant", self) + ai_action.setToolTip("Show / hide the AI Chat Assistant panel") + ai_action.triggered.connect(self._toggle_ai_panel) + menu_tools.addAction(ai_action) + def prepare_menubar_component(self, menus, comp_menu_dict): for name, action in comp_menu_dict.items(): @@ -422,6 +455,19 @@ def prepare_statusbar(self): def prepare_actions(self): + # Wire AI chat dependencies properly via set_dependencies() + if _AI_AVAILABLE and "ai_chat" in self.components: + self.components["ai_chat"].set_dependencies( + editor=self.components["editor"], + debugger=self.components["debugger"], + ) + self.components["ai_chat"].insert_code.connect( + self.components["editor"].set_text + ) + self.components["traceback_viewer"].sigAutoFixError.connect( + self.components["ai_chat"].auto_fix_error + ) + self.components["debugger"].sigRendered.connect( self.components["object_tree"].addObjects ) @@ -532,8 +578,6 @@ def prepare_console(self): ) def _examples_dir(self): - # In a PyInstaller bundle examples are extracted alongside the package. - # In development they live next to the cq_editor package directory. if getattr(sys, "frozen", False): return Path(sys._MEIPASS) / "examples" return Path(__file__).parent.parent / "examples" @@ -545,7 +589,6 @@ def _populate_examples_menu(self, menu): return for path in sorted(examples_dir.glob("*.py")): - # Strip the leading "NN_" numbering prefix for the menu label. label = path.stem if len(label) > 3 and label[2] == "_" and label[:2].isdigit(): label = label[3:] @@ -594,7 +637,8 @@ def handle_exception(exc_type, exc_value, exc_traceback): def edit_preferences(self): prefs = PreferencesWidget(self, self.components) - prefs.exec_() + if prefs.exec_() == QDialog.Accepted: + self.savePreferences() def about(self): @@ -622,28 +666,15 @@ def handle_filename_change(self, fname): self.setWindowTitle(f"{self.name}: {new_title}") def update_window_title(self, modified): - """ - Allows updating the window title to show that the document has been modified. - """ title = self.windowTitle().rstrip("*") if modified: title += "*" self.setWindowTitle(title) def update_statusbar(self, status_text): - """ - Allow updating the status bar with information. - """ - - # Update the statusbar text self.status_label.setText(status_text) def _upload_model(self): - """ - Allows the userr to easily upload models to an online service for manufacturing, - analysis, simulation, display, etc. - """ - obj_tree = self.components["object_tree"] selected = [ item @@ -658,6 +689,19 @@ def _upload_model(self): ) dlg.exec_() + def _toggle_ai_panel(self): + """Show or hide the AI Assistant dock panel.""" + if not _AI_AVAILABLE: + return + dock_widget = self.docks.get("ai_chat") + if dock_widget is None: + return + if dock_widget.isVisible(): + dock_widget.hide() # hide only — no raise_() here + else: + dock_widget.show() + dock_widget.raise_() # raise_() only when making visible + if __name__ == "__main__": diff --git a/cq_editor/preferences.py b/cq_editor/preferences.py index 2badd740..b2725b0b 100644 --- a/cq_editor/preferences.py +++ b/cq_editor/preferences.py @@ -1,4 +1,4 @@ -from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QStackedWidget, QDialog +from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QStackedWidget, QDialog, QDialogButtonBox from PyQt5.QtCore import pyqtSlot, Qt from pyqtgraph.parametertree import ParameterTree @@ -45,7 +45,12 @@ def __init__(self, parent, components): self.add(v.name, v) self.splitter = splitter((self.preferences_tree, self.stacked), (2, 5)) - layout(self, (self.splitter,), self) + + self.button_box = QDialogButtonBox(QDialogButtonBox.Save | QDialogButtonBox.Cancel, self) + self.button_box.accepted.connect(self.accept) + self.button_box.rejected.connect(self.reject) + + layout(self, (self.splitter, self.button_box), self) self.preferences_tree.currentItemChanged.connect(self.handleSelection) diff --git a/cq_editor/widgets/ai_chat.py b/cq_editor/widgets/ai_chat.py new file mode 100644 index 00000000..77e1089b --- /dev/null +++ b/cq_editor/widgets/ai_chat.py @@ -0,0 +1,658 @@ +""" +AI Chat Assistant widget for CQ-editor. + +Provides a dockable panel with: + - Chat history display (colour-coded roles) + - Prompt input box (Enter or Send button) + - Async LLM calls via QThread — UI never freezes + - Context injection: current editor script is prepended to every prompt + - "Insert & Run" pushes extracted code into the editor + - Privacy consent dialog on first use + - System keyring for API key storage (falls back to plaintext) + - Bounded conversation history (MAX_HISTORY_TURNS) + - Clean thread shutdown in closeEvent + +Optional dependencies (install separately): + pip install openai # required for LLM calls + pip install keyring # recommended for secure API key storage +""" + +from __future__ import annotations + +import re + +from PyQt5.QtCore import QObject, QThread, pyqtSignal, pyqtSlot, Qt +from PyQt5.QtWidgets import ( + QWidget, + QVBoxLayout, + QHBoxLayout, + QTextEdit, + QLineEdit, + QPushButton, + QLabel, + QMessageBox, + QScrollBar, + QAction, +) +from PyQt5.QtGui import QColor, QTextCharFormat, QFont + +from pyqtgraph.parametertree import Parameter + +# Maximum number of user+assistant turn pairs kept in history. +# Older turns are pruned to limit token cost and payload size. +MAX_HISTORY_TURNS = 10 + +# Keyring service name used when storing the API key securely. +_KEYRING_SERVICE = "cq-editor-ai-assistant" +_KEYRING_USER = "api-key" + + +# --------------------------------------------------------------------------- +# Secure key storage helpers +# --------------------------------------------------------------------------- + +def _keyring_save(key: str) -> bool: + """Try to store *key* in the system keychain. Returns True on success.""" + try: + import keyring + if not key or not key.strip() or "OrderedDict" in str(key): + try: + keyring.delete_password(_KEYRING_SERVICE, _KEYRING_USER) + except Exception: + pass + else: + keyring.set_password(_KEYRING_SERVICE, _KEYRING_USER, key) + return True + except Exception: + return False + + +def _keyring_load() -> str | None: + """Try to load the API key from the system keychain.""" + try: + import keyring + val = keyring.get_password(_KEYRING_SERVICE, _KEYRING_USER) + if val and "OrderedDict" in str(val): + try: + keyring.delete_password(_KEYRING_SERVICE, _KEYRING_USER) + except Exception: + pass + return None + return val or None + except Exception: + return None + + +# --------------------------------------------------------------------------- +# Background LLM worker +# --------------------------------------------------------------------------- + +class _LLMWorker(QObject): + """Runs the blocking OpenAI-compatible API call in a QThread.""" + + finished = pyqtSignal(str) # assistant reply text + error = pyqtSignal(str) # human-readable error + + def __init__(self, api_key: str, base_url: str, model: str, + messages: list[dict], temperature: float, max_tokens: int, parent=None): + super().__init__(parent) + self._api_key = api_key + self._base_url = base_url + self._model = model + self._messages = messages + self._temperature = temperature + self._max_tokens = max_tokens + + @pyqtSlot() + def run(self): + try: + from openai import OpenAI # lazy import — openai is optional + except ImportError: + self.error.emit( + "The 'openai' package is not installed.\n" + "Run: pip install openai" + ) + return + try: + client = OpenAI( + api_key=self._api_key, + base_url=self._base_url or None, + ) + kwargs = { + "model": self._model, + "messages": self._messages, + } + # Reasoning models (like o1-, o3-) do not accept standard temperature or max_tokens in chat completions + is_reasoning = any(x in self._model.lower() for x in ("o1-", "o3-")) + if not is_reasoning: + kwargs["temperature"] = self._temperature + kwargs["max_tokens"] = self._max_tokens + + resp = client.chat.completions.create(**kwargs) + self.finished.emit(resp.choices[0].message.content or "") + except Exception as exc: # noqa: BLE001 + self.error.emit(str(exc)) + + +# --------------------------------------------------------------------------- +# System prompt +# --------------------------------------------------------------------------- + +SYSTEM_PROMPT = ( + "You are an expert CadQuery (Python) programmer.\n" + "Always respond with a complete, valid CadQuery Python script wrapped in " + "a single ```python ... ``` fenced code block. " + "Do not include any text outside that block.\n" + "The user's current script (if any) is provided under 'CURRENT SCRIPT'. " + "Modify it as requested, or start fresh if the user asks for something new." +) + +_CODE_RE = re.compile(r"```(?:python)?\n(.*?)```", re.DOTALL | re.IGNORECASE) + + +# --------------------------------------------------------------------------- +# Main widget +# --------------------------------------------------------------------------- + +class AIChatWidget(QWidget): + """Dockable AI Chat panel for CQ-editor.""" + + name = "AI Assistant" + + # Emitted with the extracted CadQuery code; connected to Editor.set_text() + insert_code = pyqtSignal(str) + + def __init__(self, parent=None): + super().__init__(parent) + + # Instance-level preferences (not class-level) to avoid shared state + self.preferences = Parameter.create( + name="AI Assistant", + children=[ + { + "name": "Enabled", + "type": "bool", + "value": False, + "tip": "Show the AI Chat Assistant dock panel", + }, + { + "name": "Provider / Base URL", + "type": "str", + "value": "https://api.openai.com/v1", + "tip": ( + "OpenAI-compatible base URL.\n" + "OpenAI: https://api.openai.com/v1\n" + "OpenRouter: https://openrouter.ai/api/v1\n" + "Ollama: http://localhost:11434/v1" + ), + }, + { + "name": "Model", + "type": "str", + "value": "gpt-4o", + "tip": "Model identifier. Frontier models yield vastly superior geometric reasoning.", + }, + { + "name": "API Key", + "type": "str", + "value": "", + "tip": ( + "Your API key.\n" + "Stored in the system keychain when 'keyring' is installed;\n" + "otherwise stored in plaintext preferences on disk.\n" + "Run: pip install keyring for secure storage." + ), + }, + { + "name": "Temperature", + "type": "float", + "value": 0.2, + "limits": (0.0, 2.0), + "tip": "Controls randomness of responses. Lower is more deterministic.", + }, + { + "name": "Max Tokens", + "type": "int", + "value": 1024, + "limits": (1, 16384), + "tip": "Maximum number of tokens to generate in the reply.", + }, + { + "name": "Auto-run after insert", + "type": "bool", + "value": True, + "tip": "Re-render the model immediately after inserting LLM code", + }, + { + "name": "_privacy_consent", + "type": "bool", + "value": False, + "tip": "Internal: records that the user has seen the privacy notice", + }, + ], + ) + + self._editor = None # set via set_dependencies() + self._debugger = None + self._history: list[dict] = [] + self._thread: QThread | None = None + self._worker: _LLMWorker | None = None + self._last_code: str = "" + self._auto_insert_flag: bool = False + + self._setup_ui() + + # Wire Enabled toggle -> show/hide this widget's parent dock + self.preferences.sigTreeStateChanged.connect(self._on_prefs_changed) + + # ------------------------------------------------------------------ + # Dependency injection + # ------------------------------------------------------------------ + + def set_dependencies(self, editor, debugger): + """Called by main_window after both editor and debugger are ready.""" + self._editor = editor + self._debugger = debugger + + # Heal any corrupted API Key preferences from QSettings or Keyring + try: + api_key_val = self._pref("API Key") + if not isinstance(api_key_val, str) or "OrderedDict" in str(api_key_val): + self.preferences["API Key"] = "" + # Trigger keyring self-healing check on startup + _keyring_load() + except Exception: + pass + + # ------------------------------------------------------------------ + # UI construction + # ------------------------------------------------------------------ + + def _setup_ui(self): + root = QVBoxLayout(self) + root.setContentsMargins(4, 4, 4, 4) + root.setSpacing(4) + + header = QLabel("🤖 AI CAD Assistant") + header.setAlignment(Qt.AlignCenter) + root.addWidget(header) + + # Privacy notice banner + self._privacy_banner = QLabel( + "⚠ Each prompt sends your current script to the " + "configured API endpoint." + ) + self._privacy_banner.setWordWrap(True) + self._privacy_banner.setAlignment(Qt.AlignCenter) + self._privacy_banner.setStyleSheet("color: #7B5B00; background: #FFF8E1; " + "padding: 2px; border-radius: 3px;") + root.addWidget(self._privacy_banner) + + self.chat_display = QTextEdit() + self.chat_display.setReadOnly(True) + self.chat_display.setFont(QFont("Monospace", 9)) + self.chat_display.setMinimumHeight(200) + root.addWidget(self.chat_display, stretch=1) + + input_row = QHBoxLayout() + self.prompt_input = QLineEdit() + self.prompt_input.setPlaceholderText( + "Describe what you want to model or change\u2026" + ) + self.prompt_input.returnPressed.connect(self._send) + input_row.addWidget(self.prompt_input, stretch=1) + + self.send_btn = QPushButton("Send") + self.send_btn.setFixedWidth(60) + self.send_btn.clicked.connect(self._send) + input_row.addWidget(self.send_btn) + root.addLayout(input_row) + + btn_row = QHBoxLayout() + + self.insert_btn = QPushButton("Insert & Run") + self.insert_btn.setToolTip( + "Copy the extracted code block into the editor and run it" + ) + self.insert_btn.setEnabled(False) + self.insert_btn.clicked.connect(self._insert_last_code) + btn_row.addWidget(self.insert_btn) + + self.clear_btn = QPushButton("Clear Chat") + self.clear_btn.clicked.connect(self._clear) + btn_row.addWidget(self.clear_btn) + root.addLayout(btn_row) + + self.status_label = QLabel("") + self.status_label.setAlignment(Qt.AlignCenter) + root.addWidget(self.status_label) + + # ------------------------------------------------------------------ + # Preference helpers + # ------------------------------------------------------------------ + + def _pref(self, name: str): + return self.preferences[name] + + def _api_key(self) -> str: + # Prefer system keyring; fall back to plaintext preference + key = _keyring_load() + if key and key.strip(): + return key.strip() + return self._pref("API Key").strip() + + def _save_api_key(self, key: str): + if not _keyring_save(key): + # keyring unavailable — store in plaintext preference with warning + self.preferences["API Key"] = key + + def _base_url(self) -> str: + return self._pref("Provider / Base URL").strip() + + def _model(self) -> str: + return self._pref("Model").strip() + + def _auto_run(self) -> bool: + return self._pref("Auto-run after insert") + + @pyqtSlot(object, object) + def _on_prefs_changed(self, _param, changes): + """React to preference changes; wire Enabled to dock visibility & sync API Key.""" + for param, _change, _val in changes: + if param.name() == "Enabled": + dock = self._find_dock() + if dock is not None: + if self._pref("Enabled"): + dock.show() + dock.raise_() + else: + dock.hide() + elif param.name() == "API Key" and _change == "value": + self._save_api_key(_val) + + def _find_dock(self): + """Walk up the parent chain to find the QDockWidget containing us.""" + from PyQt5.QtWidgets import QDockWidget + p = self.parent() + while p is not None: + if isinstance(p, QDockWidget): + return p + p = p.parent() + return None + + # ------------------------------------------------------------------ + # Privacy consent + # ------------------------------------------------------------------ + + def _ensure_privacy_consent(self) -> bool: + """Show a one-time privacy dialog. Returns True if user consents.""" + if self._pref("_privacy_consent"): + return True + rv = QMessageBox.information( + self, + "Privacy Notice \u2014 AI Assistant", + "When you send a prompt, your current CadQuery script and prompt " + "text are sent to the API endpoint you configured:\n\n" + " {url}\n\n" + "Make sure you are comfortable sharing your code with that provider " + "before continuing. You can change the endpoint in " + "Edit \u2192 Preferences \u2192 AI Assistant.\n\n" + "Do you want to continue?".format(url=self._base_url()), + QMessageBox.Yes | QMessageBox.No, + QMessageBox.No, + ) + if rv == QMessageBox.Yes: + self.preferences["_privacy_consent"] = True + return True + return False + + # ------------------------------------------------------------------ + # History management + # ------------------------------------------------------------------ + + def _prune_history(self): + """Keep system prompt + at most MAX_HISTORY_TURNS turn pairs.""" + non_system = [m for m in self._history if m["role"] != "system"] + system = [m for m in self._history if m["role"] == "system"] + # Each turn = one user + one assistant message + max_msgs = MAX_HISTORY_TURNS * 2 + if len(non_system) > max_msgs: + non_system = non_system[-max_msgs:] + self._history = system + non_system + + # ------------------------------------------------------------------ + # Slots + # ------------------------------------------------------------------ + + @pyqtSlot() + def _send(self): + prompt = self.prompt_input.text().strip() + if not prompt: + return + + if not self._api_key(): + QMessageBox.warning( + self, + "API Key Missing", + "Please enter your API key in\n" + "Edit \u2192 Preferences \u2192 AI Assistant \u2192 API Key", + ) + return + + if not self._ensure_privacy_consent(): + return + + try: + running = self._thread and self._thread.isRunning() + except RuntimeError: + self._thread = None + running = False + + if running: + self.status_label.setText("Waiting for previous response\u2026") + return + + self.prompt_input.clear() + self._append_chat("user", prompt) + + if not self._history: + self._history.append({"role": "system", "content": SYSTEM_PROMPT}) + + self._history.append({"role": "user", "content": self._inject_context(prompt)}) + self._prune_history() + + self._set_busy(True) + self._run_worker(list(self._history)) + + def _inject_context(self, prompt: str) -> str: + """Prepend the current editor script as context.""" + if self._editor is None: + return prompt + try: + current_code = self._editor.toPlainText().strip() + except Exception: + current_code = "" + if not current_code: + return prompt + return ( + f"CURRENT SCRIPT:\n```python\n{current_code}\n```\n\n" + f"USER REQUEST:\n{prompt}" + ) + + @pyqtSlot(str) + def auto_fix_error(self, prompt: str): + """Automatically called when the user clicks 'Auto-Fix with AI' in the traceback pane.""" + self._auto_insert_flag = True + + # Ensure the dock panel is shown and raised so they can see the progress + dock = self._find_dock() + if dock is not None: + if not self._pref("Enabled"): + self.preferences["Enabled"] = True + dock.show() + dock.raise_() + + # Clear the input box + self.prompt_input.clear() + + # Add visual system feedback to user + self._append_chat("system", "✨ Auto-fixing script error...") + self._append_chat("user", prompt) + + if not self._history: + self._history.append({"role": "system", "content": SYSTEM_PROMPT}) + + self._history.append({"role": "user", "content": self._inject_context(prompt)}) + self._prune_history() + + self._set_busy(True) + self._run_worker(list(self._history)) + + def _run_worker(self, messages: list[dict]): + self._thread = QThread(self) + self._worker = _LLMWorker( + api_key=self._api_key(), + base_url=self._base_url(), + model=self._model(), + messages=messages, + temperature=self._pref("Temperature"), + max_tokens=self._pref("Max Tokens"), + ) + self._worker.moveToThread(self._thread) + self._thread.started.connect(self._worker.run) + self._worker.finished.connect(self._on_response) + self._worker.error.connect(self._on_error) + self._worker.finished.connect(self._thread.quit) + self._worker.error.connect(self._thread.quit) + self._thread.finished.connect(self._thread.deleteLater) + self._thread.finished.connect(self._clear_thread_ref) + self._thread.start() + + @pyqtSlot() + def _clear_thread_ref(self): + self._thread = None + self._worker = None + + @pyqtSlot(str) + def _on_response(self, reply: str): + self._set_busy(False) + self._history.append({"role": "assistant", "content": reply}) + self._append_chat("assistant", reply) + + code = _CODE_RE.search(reply) + if code: + self._last_code = code.group(1).strip() + self.insert_btn.setEnabled(True) + if self._auto_insert_flag: + self._insert_last_code() + else: + # No fenced code block found — do NOT fall back to raw reply. + # Show an informational message so the user knows what happened. + self._last_code = "" + self.insert_btn.setEnabled(False) + self._append_chat( + "system", + "\u26a0\ufe0f The response did not contain a fenced code block. " + "Ask the AI to provide the complete script again.", + ) + self._auto_insert_flag = False + + @pyqtSlot(str) + def _on_error(self, message: str): + self._set_busy(False) + self._append_chat("error", message) + self._auto_insert_flag = False + + @pyqtSlot() + def _insert_last_code(self): + if not self._last_code: + return + self.insert_code.emit(self._last_code) + if self._auto_run() and self._debugger is not None: + try: + self._debugger.render() + except Exception: + pass + self._append_chat("system", "\u2705 Code inserted into editor.") + + @pyqtSlot() + def _clear(self): + self._history.clear() + self._last_code = "" + self.insert_btn.setEnabled(False) + self.chat_display.clear() + self.status_label.setText("") + # Reset consent so the privacy notice shows again after a clear + self.preferences["_privacy_consent"] = False + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def closeEvent(self, event): + """Ensure the background thread is cleanly stopped before the widget closes.""" + try: + running = self._thread and self._thread.isRunning() + except RuntimeError: + self._thread = None + running = False + + if running: + self._thread.quit() + self._thread.wait(3000) # wait up to 3 s + super().closeEvent(event) + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + + def _set_busy(self, busy: bool): + self.send_btn.setEnabled(not busy) + self.prompt_input.setEnabled(not busy) + self.status_label.setText("Thinking\u2026" if busy else "") + + def _append_chat(self, role: str, text: str): + cursor = self.chat_display.textCursor() + cursor.movePosition(cursor.End) + + fmt_label = QTextCharFormat() + fmt_text = QTextCharFormat() + + if role == "user": + fmt_label.setForeground(QColor("#1565C0")) + fmt_label.setFontWeight(QFont.Bold) + label = "You: " + elif role == "assistant": + fmt_label.setForeground(QColor("#1B5E20")) + fmt_label.setFontWeight(QFont.Bold) + label = "AI: " + elif role == "error": + fmt_label.setForeground(QColor("#B71C1C")) + fmt_label.setFontWeight(QFont.Bold) + label = "ERR: " + fmt_text.setForeground(QColor("#B71C1C")) + else: # system / info + fmt_label.setForeground(QColor("#4A148C")) + label = " " + + cursor.insertText(label, fmt_label) + cursor.insertText(text.strip() + "\n\n", fmt_text) + + sb: QScrollBar = self.chat_display.verticalScrollBar() + sb.setValue(sb.maximum()) + + # ------------------------------------------------------------------ + # MainMixin compatibility stubs + # ------------------------------------------------------------------ + + def menuActions(self) -> dict: + return {} + + def toolbarActions(self) -> list: + return [] + + def saveComponentState(self, store): + pass + + def restoreComponentState(self, store): + pass diff --git a/cq_editor/widgets/log.py b/cq_editor/widgets/log.py index 1e9dd911..247d352a 100644 --- a/cq_editor/widgets/log.py +++ b/cq_editor/widgets/log.py @@ -2,7 +2,7 @@ import re from PyQt5 import QtGui -from PyQt5.QtCore import QObject, pyqtSignal +from PyQt5.QtCore import QObject, pyqtSignal, pyqtSlot from PyQt5.QtWidgets import QPlainTextEdit, QAction from ..mixins import ComponentMixin @@ -90,6 +90,7 @@ def __init__(self, *args, **kwargs): # Ensure handler is closed when widget is destroyed self.destroyed.connect(lambda *_: self.handler.close()) + @pyqtSlot(str) def append(self, msg): """Append text to the panel with ANSI escape sequences stipped.""" self.moveCursor(QtGui.QTextCursor.End) diff --git a/cq_editor/widgets/traceback_viewer.py b/cq_editor/widgets/traceback_viewer.py index b02f378f..a9ad465d 100644 --- a/cq_editor/widgets/traceback_viewer.py +++ b/cq_editor/widgets/traceback_viewer.py @@ -1,7 +1,7 @@ from traceback import extract_tb, format_exception_only from itertools import dropwhile -from PyQt5.QtWidgets import QWidget, QTreeWidget, QTreeWidgetItem, QAction, QLabel +from PyQt5.QtWidgets import QWidget, QTreeWidget, QTreeWidgetItem, QAction, QLabel, QHBoxLayout, QPushButton from PyQt5.QtCore import Qt, pyqtSlot, pyqtSignal from PyQt5.QtGui import QFontMetrics @@ -30,19 +30,67 @@ def __init__(self, parent): class TracebackPane(QWidget, ComponentMixin): sigHighlightLine = pyqtSignal(int) + sigAutoFixError = pyqtSignal(str) def __init__(self, parent): super(TracebackPane, self).__init__(parent) + self.last_exc_info = None + self.last_code = None self.tree = TracebackTree(self) self.current_exception = QLabel(self) self.current_exception.setStyleSheet("QLabel {color : red; }") - layout(self, (self.current_exception, self.tree), self) + # Create a horizontal row for the exception message and the Auto-Fix button + top_row_widget = QWidget(self) + top_row_layout = QHBoxLayout(top_row_widget) + top_row_layout.setContentsMargins(0, 0, 0, 0) + top_row_layout.setSpacing(4) + top_row_widget.setLayout(top_row_layout) + + top_row_layout.addWidget(self.current_exception, stretch=1) + + self.autofix_btn = QPushButton("✨ Auto-Fix with AI", top_row_widget) + self.autofix_btn.setToolTip("Automatically send this error and code to the AI Assistant to fix it") + self.autofix_btn.setStyleSheet( + "QPushButton { background-color: #7B1FA2; color: white; font-weight: bold; border-radius: 3px; padding: 4px 8px; }" + "QPushButton:hover { background-color: #8E24AA; }" + "QPushButton:disabled { background-color: #BDBDBD; color: #757575; }" + ) + self.autofix_btn.setEnabled(False) + self.autofix_btn.clicked.connect(self.trigger_autofix) + top_row_layout.addWidget(self.autofix_btn) + + layout(self, (top_row_widget, self.tree), self) self.tree.currentItemChanged.connect(self.handleSelection) + def trigger_autofix(self): + if not self.last_exc_info: + return + + t, exc, tb = self.last_exc_info + exc_name = t.__name__ + exc_msg = str(exc) + + tb_list = extract_tb(tb) + filtered_tb = list(dropwhile(lambda el: "string>" not in el.filename, tb_list)) + + error_context = "" + if filtered_tb: + last_frame = filtered_tb[-1] + error_context = f"at line {last_frame.lineno}: `{last_frame.line}`" + else: + error_context = "in the script" + + prompt = ( + f"I encountered a '{exc_name}' error {error_context}.\n" + f"Error details: {exc_msg}\n" + f"Please identify the issue in the code and provide the complete, corrected CadQuery script." + ) + self.sigAutoFixError.emit(prompt) + def truncate_text(self, text, max_length=100): """ Used to prevent the label from expanding the window width off the screen. @@ -58,9 +106,12 @@ def truncate_text(self, text, max_length=100): def addTraceback(self, exc_info, code): self.tree.clear() + self.last_exc_info = exc_info + self.last_code = code if exc_info: t, exc, tb = exc_info + self.autofix_btn.setEnabled(True) root = self.tree.root code = code.splitlines() @@ -100,6 +151,7 @@ def addTraceback(self, exc_info, code): else: self.current_exception.setText("") self.current_exception.setToolTip("") + self.autofix_btn.setEnabled(False) @pyqtSlot(QTreeWidgetItem, QTreeWidgetItem) def handleSelection(self, item, *args): diff --git a/pyproject.toml b/pyproject.toml index dbf3f5c9..f0dcd979 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,10 @@ CQ-editor = "cq_editor.cqe_run:main" cq-editor = "cq_editor.cqe_run:main" [project.optional-dependencies] +ai = [ + "openai", + "keyring" +] test = [ "pytest", "pluggy", diff --git a/tests/test_ai_chat.py b/tests/test_ai_chat.py new file mode 100644 index 00000000..4cde3f62 --- /dev/null +++ b/tests/test_ai_chat.py @@ -0,0 +1,75 @@ +import pytest +from PyQt5.QtCore import Qt, QThread +from PyQt5.QtWidgets import QMessageBox +from pyqtgraph.parametertree import Parameter + +from cq_editor.widgets.ai_chat import AIChatWidget, _keyring_load, _keyring_save +from cq_editor.__main__ import MainWindow + +def test_ai_chat_init(qtbot): + """Test that AIChatWidget initializes with correct default preferences.""" + widget = AIChatWidget() + qtbot.addWidget(widget) + + assert widget.preferences is not None + assert isinstance(widget.preferences, Parameter) + assert widget._pref("Enabled") is False + assert widget._pref("Provider / Base URL") == "https://integrate.api.nvidia.com/v1" + assert widget._pref("Model") == "meta/llama-3.3-70b-instruct" + assert widget._pref("API Key") == "" + +def test_keyring_sanitation(): + """Test that keyring load/save successfully deletes and handles corrupted OrderedDict maps.""" + # Test saving a corrupted key containing OrderedDict + assert _keyring_save("OrderedDict([('tip', 'Your API key.')])") is True + # Loading should automatically detect the corruption, delete it, and return None + assert _keyring_load() is None + + # Test saving a clean string key + assert _keyring_save("my-clean-test-key") is True + assert _keyring_load() == "my-clean-test-key" + + # Clean up + assert _keyring_save("") is True + assert _keyring_load() is None + +def test_privacy_consent_dialog(qtbot, mocker): + """Test the one-time privacy consent notice triggers correctly and saves selection.""" + widget = AIChatWidget() + qtbot.addWidget(widget) + + # Initially consent is False + assert widget._pref("_privacy_consent") is False + + # Mock the QMessageBox to simulate user clicking Yes + mocker.patch.object(QMessageBox, "information", return_value=QMessageBox.Yes) + assert widget._ensure_privacy_consent() is True + assert widget._pref("_privacy_consent") is True + + # Subsequent calls should return True immediately without prompting + mocker.patch.object(QMessageBox, "information", side_effect=Exception("Should not prompt")) + assert widget._ensure_privacy_consent() is True + +def test_auto_fix_error_flow(qtbot, mocker): + """Test that auto-fix error triggers the correct chat flow and status updating.""" + # Mock MainWindow and dependencies + win = MainWindow() + qtbot.addWidget(win) + + chat = win.components.get("ai_chat") + if chat is None: + pytest.skip("AI Assistant is not loaded in MainWindow components") + + # Mock the worker execution to prevent thread start + mock_run = mocker.patch.object(chat, "_run_worker") + + # Trigger auto-fix + test_error = "Standard Failure: edges mismatch on line 12" + chat.auto_fix_error(test_error) + + # Should set auto_insert_flag and initiate LLM worker + assert chat._auto_insert_flag is True + assert mock_run.call_count == 1 + + # Clean up + chat._auto_insert_flag = False