diff --git a/.github/workflows/graphite-ci.yml b/.github/workflows/graphite-ci.yml new file mode 100644 index 0000000..92aca47 --- /dev/null +++ b/.github/workflows/graphite-ci.yml @@ -0,0 +1,30 @@ +name: Graphite CI + +on: + push: + branches: + - main + - fix/** + pull_request: + +jobs: + regression-tests: + runs-on: ubuntu-latest + env: + QT_QPA_PLATFORM: offscreen + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + - name: Install Qt system libraries + run: | + sudo apt-get update + sudo apt-get install -y libegl1 libgl1 libopengl0 libxkbcommon-x11-0 libxcb-cursor0 + - name: Install dependencies + run: python -m pip install -r requirements.txt + - name: Compile modular application + run: python -m compileall -q graphite_app/graphite_app + - name: Run regression tests + run: python -m unittest discover -s tests -v diff --git a/README.md b/README.md index cb6ac0c..f60d551 100644 --- a/README.md +++ b/README.md @@ -132,58 +132,52 @@ Follow these steps to run Graphite locally. ### 1. Prerequisites -- **Python:** Python 3.8 or newer +- **Python:** Python 3.10 or newer - **Ollama:** Installed and running — https://ollama.com -### 2. Install an LLM Model - -Pull the default model used by the application: - -```bash -ollama pull qwen2.5:7b-instruct -```` - -Additional models can be configured inside the application using the model selection dialog. - -### 3. Clone and Install Dependencies - -Clone the repository: +### 2. Clone and Install Dependencies ```bash git clone https://github.com/dovvnloading/Graphite.git cd Graphite +python -m venv venv ``` -Create a virtual environment (recommended): +Activate the environment with `venv\\Scripts\\activate` on Windows or +`source venv/bin/activate` on macOS/Linux, then install the modular PySide6 app: ```bash -# Windows -python -m venv venv -.\venv\Scripts\activate - -# macOS/Linux -python3 -m venv venv -source venv/bin/activate +python -m pip install -r requirements.txt ``` -Install dependencies: +### 3. Install the Default Ollama Models + +Install and start [Ollama](https://ollama.com), then pull all three default models +configured in `graphite_app/graphite_app/graphite_config.py`: ```bash -pip install -r requirements.txt +ollama pull qwen2.5:3b +ollama pull qwen2.5:7b-instruct +ollama pull deepseek-coder:6.7b ``` -If no `requirements.txt` is available: +These are used for session titles, chat, and charts respectively. Graphite checks +for missing local models at startup and shows the required pull commands. +To use a remote API instead, select **API** mode and configure its provider and +models in Settings. -```bash -pip install PySide6 ollama matplotlib qtawesome -``` +### 4. Run the Modular Application -### 4. Run the Application +From the repository root: ```bash -python graphite_app.py +python graphite_app/graphite_app/graphite_app.py ``` +The top-level `Graphite.py` is an archived PyQt5 version, not the modular +application. To run that older version instead, install +`requirements-legacy.txt` and execute `python Graphite.py`. + ## Usage * **Sending Messages:** Enter a prompt in the input box and press Enter to create a new node followed by the AI response. diff --git a/graphite_app/graphite_app/api_provider.py b/graphite_app/graphite_app/api_provider.py index 04cf6c4..cc3b6de 100644 --- a/graphite_app/graphite_app/api_provider.py +++ b/graphite_app/graphite_app/api_provider.py @@ -175,6 +175,21 @@ def get_available_models(): except Exception as e: raise RuntimeError(f"Failed to fetch models from endpoint: {str(e)}") +def get_missing_ollama_models(): + """Check all configured local task models without triggering downloads.""" + response = ollama.Client(timeout=2.0).list() + installed = response.get('models', []) if isinstance(response, dict) else response.models + available = set() + for item in installed: + if isinstance(item, dict): + name = item.get('model') or item.get('name') + else: + name = getattr(item, 'model', None) or getattr(item, 'name', None) + if name: + available.add(name) + return sorted(set(config.OLLAMA_MODELS.values()) - available) + + def set_mode(use_api: bool): global USE_API_MODE USE_API_MODE = use_api diff --git a/graphite_app/graphite_app/graphite_agents.py b/graphite_app/graphite_app/graphite_agents.py index 551a595..f94481b 100644 --- a/graphite_app/graphite_app/graphite_agents.py +++ b/graphite_app/graphite_app/graphite_agents.py @@ -18,44 +18,41 @@ def __init__(self, agent, message, conversation_history): def run(self): try: # Use agent's get_response directly - response = self.agent.get_response(self.message) + response = self.agent.get_response(self.message, self.conversation_history) self.finished.emit(response) except Exception as e: self.error.emit(str(e)) class ChatWorker: - """Compose chat messages and call the configured provider synchronously.""" + """Compose chat messages for one branch and call the configured provider.""" def __init__(self, system_prompt, conversation_history): self.system_prompt = system_prompt - self.conversation_history = conversation_history - + self.conversation_history = list(conversation_history) + def run(self, user_message): - try: - messages = [ - {'role': 'system', 'content': self.system_prompt}, - *self.conversation_history, - {'role': 'user', 'content': user_message} - ] - response = api_provider.chat(task=config.TASK_CHAT, messages=messages) - ai_message = response['message']['content'] - return ai_message - except Exception as e: - return f"Error: {str(e)}" + # A node's history already ends with the new user message. Do not + # append it twice; also support callers that provide ancestor history. + messages = [ + {'role': 'system', 'content': self.system_prompt}, + *self.conversation_history + ] + user_turn = {'role': 'user', 'content': user_message} + if not self.conversation_history or self.conversation_history[-1] != user_turn: + messages.append(user_turn) + response = api_provider.chat(task=config.TASK_CHAT, messages=messages) + return response['message']['content'] + class ChatAgent: - """Stateful assistant that tracks conversation history per node/session.""" + """Stateless assistant: all context comes from the requested graph branch.""" def __init__(self, name, persona): self.name = name or "AI Assistant" self.persona = persona or "(default persona)" self.system_prompt = f"You are {self.name}. {self.persona}" - self.conversation_history = [] - - def get_response(self, user_message): - chat_worker = ChatWorker(self.system_prompt, self.conversation_history) - ai_response = chat_worker.run(user_message) - self.conversation_history.append({'role': 'user', 'content': user_message}) - self.conversation_history.append({'role': 'assistant', 'content': ai_response}) - return ai_response + + def get_response(self, user_message, conversation_history=None): + worker = ChatWorker(self.system_prompt, conversation_history or []) + return worker.run(user_message) class ExplainerAgent: """Generate plain-language explanations and normalize the output format.""" diff --git a/graphite_app/graphite_app/graphite_app.py b/graphite_app/graphite_app/graphite_app.py index 15fb165..14d1873 100644 --- a/graphite_app/graphite_app/graphite_app.py +++ b/graphite_app/graphite_app/graphite_app.py @@ -3,7 +3,7 @@ QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QToolBar, QToolButton, QLineEdit, QPushButton, QMessageBox, QSizePolicy, QLabel, QComboBox ) -from PySide6.QtCore import Qt, QSize, QPointF +from PySide6.QtCore import Qt, QSize, QPointF, QTimer from PySide6.QtGui import QKeySequence, QGuiApplication, QCursor, QShortcut import qtawesome as qta import json @@ -215,9 +215,46 @@ def keyPressEvent(self, event): super().keyPressEvent(event) def save_chat(self): - """Save the current chat session""" - self.session_manager.save_current_chat() - QMessageBox.information(self, "Success", "Chat saved successfully!") + """Save the session and report the actual result.""" + try: + saved = self.session_manager.save_current_chat() + except Exception as exc: + QMessageBox.critical(self, "Save Failed", str(exc)) + return + if saved: + QMessageBox.information(self, "Success", "Chat saved successfully!") + else: + QMessageBox.information(self, "Nothing to Save", "The canvas is empty.") + + def confirm_session_transition(self): + """Protect unsaved work before loading, clearing, or closing a session.""" + for worker_name in ('chat_thread', 'takeaway_thread', 'explainer_thread', 'chart_thread'): + worker = getattr(self, worker_name, None) + if worker is not None and worker.isRunning(): + QMessageBox.warning( + self, "Generation in Progress", + "Finish the current AI task before switching or closing chats." + ) + return False + if not self.session_manager.has_unsaved_changes(): + return True + choice = QMessageBox.question( + self, "Unsaved Changes", + "Save your changes before leaving this chat?", + QMessageBox.StandardButton.Save | + QMessageBox.StandardButton.Discard | + QMessageBox.StandardButton.Cancel, + QMessageBox.StandardButton.Save + ) + if choice == QMessageBox.StandardButton.Cancel: + return False + if choice == QMessageBox.StandardButton.Discard: + return True + try: + return self.session_manager.save_current_chat() + except Exception as exc: + QMessageBox.critical(self, "Save Failed", str(exc)) + return False def setup_toolbar(self, toolbar): """Setup toolbar with modern QToolButtons""" @@ -351,6 +388,27 @@ def show_help(self): center.y() - help_dialog.height() // 2) help_dialog.show() + def verify_default_models(self): + """Help first-time local users install every required Ollama task model.""" + if api_provider.get_mode() != "Ollama": + return + try: + missing = api_provider.get_missing_ollama_models() + except Exception as exc: + QMessageBox.warning( + self, "Ollama Unavailable", + "Graphite could not verify local models. Start Ollama and check " + f"your connection. Details: {exc}" + ) + return + if missing: + commands = "\n".join(f"ollama pull {model}" for model in missing) + QMessageBox.warning( + self, "Missing Ollama Models", + "Some local Graphite features need additional models. " + "Run these commands in a terminal:\n\n" + commands + ) + def show_settings(self): """Show the appropriate settings dialog based on the current mode.""" use_api = self.mode_combo.currentData() @@ -372,7 +430,18 @@ def on_mode_changed(self, index): def setCurrentNode(self, node): self.current_node = node self.message_input.setPlaceholderText(f"Responding to: {node.text[:30]}...") - + + def build_conversation_history(self, node): + """Read only the selected node's ancestor chain, oldest first.""" + turns = [] + while node is not None: + turns.append({ + 'role': 'user' if node.is_user else 'assistant', + 'content': node.text + }) + node = node.parent_node + return list(reversed(turns)) + def send_message(self): message = self.message_input.text().strip() if not message: @@ -385,15 +454,9 @@ def send_message(self): # Show loading overlay self.loading_overlay.show() - # Get conversation history up to current node - history = [] - if self.current_node: - temp_node = self.current_node - while temp_node: - if hasattr(temp_node, 'conversation_history'): - history = temp_node.conversation_history + history - temp_node = temp_node.parent_node - + # Ancestor nodes are the source of truth, not cumulative node histories. + history = self.build_conversation_history(self.current_node) + # Add user message node user_node = self.chat_view.scene().add_chat_node( message, @@ -424,7 +487,7 @@ def handle_response(self, response, user_node): response, is_user=False, parent_node=user_node, - conversation_history=self.agent.conversation_history + [ + conversation_history=user_node.conversation_history + [ {'role': 'assistant', 'content': response} ] ) @@ -639,17 +702,19 @@ def stop_all_workers(self): self.chart_thread.wait() def closeEvent(self, event): - """ - Overrides the default close event to ensure all background threads - are properly terminated before the application exits. - """ + """Save or explicitly discard canvas edits before closing.""" + if not self.confirm_session_transition(): + event.ignore() + return self.stop_all_workers() super().closeEvent(event) + def main(): app = QApplication(sys.argv) window = ChatWindow() window.show() + QTimer.singleShot(0, window.verify_default_models) sys.exit(app.exec()) if __name__ == "__main__": diff --git a/graphite_app/graphite_app/graphite_core.py b/graphite_app/graphite_app/graphite_core.py index 824cff8..5025eda 100644 --- a/graphite_app/graphite_app/graphite_core.py +++ b/graphite_app/graphite_app/graphite_core.py @@ -181,6 +181,50 @@ def load_notes(self, chat_id): }) return notes + def save_snapshot(self, title, chat_data, notes_data, pins_data, chat_id=None): + """Atomically save a canvas and its associated notes/navigation pins.""" + with self._connect() as conn: + payload = json.dumps(chat_data) + if chat_id is None: + cursor = conn.execute( + "INSERT INTO chats (title, data, updated_at) " + "VALUES (?, ?, CURRENT_TIMESTAMP)", (title, payload) + ) + chat_id = cursor.lastrowid + else: + cursor = conn.execute( + "UPDATE chats SET title = ?, data = ?, " + "updated_at = CURRENT_TIMESTAMP WHERE id = ?", + (title, payload, chat_id) + ) + if cursor.rowcount == 0: + raise ValueError(f"Chat {chat_id} no longer exists") + + conn.execute("DELETE FROM notes WHERE chat_id = ?", (chat_id,)) + for note in notes_data: + conn.execute( + "INSERT INTO notes (chat_id, content, position_x, position_y, " + "width, height, color, header_color) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + ( + chat_id, note['content'], note['position']['x'], + note['position']['y'], note['size']['width'], + note['size']['height'], note['color'], + note.get('header_color') + ) + ) + + conn.execute("DELETE FROM pins WHERE chat_id = ?", (chat_id,)) + for pin in pins_data: + conn.execute( + "INSERT INTO pins (chat_id, title, note, position_x, position_y) " + "VALUES (?, ?, ?, ?, ?)", + ( + chat_id, pin['title'], pin['note'], + pin['position']['x'], pin['position']['y'] + ) + ) + return chat_id + def save_chat(self, title, chat_data): """Persist a new chat snapshot and return the inserted row id.""" with self._connect() as conn: @@ -253,6 +297,7 @@ def __init__(self, window): self.db = ChatDatabase() self.title_generator = TitleGenerator() self.current_chat_id = None + self._saved_snapshot = self.capture_current_snapshot() def serialize_pin(self, pin): """Convert a navigation pin to a serializable dictionary""" @@ -319,50 +364,73 @@ def serialize_chart(self, chart): 'size': {'width': chart.width, 'height': chart.height} } - def serialize_current_chat(self): - """Serialize the current chat session with all elements""" + def capture_current_snapshot(self): + """Capture canvas, notes and pins without writing to the database.""" scene = self.window.chat_view.scene() - - # Get all notes, pins, and charts in the scene notes = [item for item in scene.items() if isinstance(item, Note)] pins = [item for item in scene.items() if isinstance(item, NavigationPin)] charts = [item for item in scene.items() if isinstance(item, ChartItem)] - - chat_data = { - 'nodes': [self.serialize_node(node) for node in scene.nodes], - 'connections': [self.serialize_connection(conn) for conn in scene.connections], - 'frames': [self.serialize_frame(frame) for frame in scene.frames], - 'charts': [self.serialize_chart(chart) for chart in charts], # Add charts - 'view_state': { - 'zoom_factor': self.window.chat_view._zoom_factor, - 'scroll_position': { - 'x': self.window.chat_view.horizontalScrollBar().value(), - 'y': self.window.chat_view.verticalScrollBar().value() + return { + 'chat_data': { + 'nodes': [self.serialize_node(node) for node in scene.nodes], + 'connections': [self.serialize_connection(conn) for conn in scene.connections], + 'frames': [self.serialize_frame(frame) for frame in scene.frames], + 'charts': [self.serialize_chart(chart) for chart in charts], + 'view_state': { + 'zoom_factor': self.window.chat_view._zoom_factor, + 'scroll_position': { + 'x': self.window.chat_view.horizontalScrollBar().value(), + 'y': self.window.chat_view.verticalScrollBar().value() + } } - } + }, + 'notes_data': [self.serialize_note(note) for note in notes], + 'pins_data': [self.serialize_pin(pin) for pin in pins] } - - # Save chat data first - if not self.current_chat_id: - last_message = scene.nodes[-1].text if scene.nodes else "New Chat" + + @staticmethod + def _has_content(snapshot): + scene_data = snapshot['chat_data'] + return bool( + any(scene_data[key] for key in ('nodes', 'connections', 'frames', 'charts')) + or snapshot['notes_data'] or snapshot['pins_data'] + ) + + def has_unsaved_changes(self): + current = self.capture_current_snapshot() + if self.current_chat_id is None and not self._has_content(current): + return False + return current != self._saved_snapshot + + def mark_current_state_saved(self): + """Record a clean baseline after a new session or successful load.""" + self._saved_snapshot = self.capture_current_snapshot() + + def serialize_current_chat(self): + """Persist a complete scene in one database transaction.""" + snapshot = self.capture_current_snapshot() + if self.current_chat_id is None and not self._has_content(snapshot): + return None + + chat_data = snapshot['chat_data'] + if self.current_chat_id is None: + last_message = chat_data['nodes'][-1]['text'] if chat_data['nodes'] else "New Chat" title = self.title_generator.generate_title(last_message) - self.current_chat_id = self.db.save_chat(title, chat_data) else: - chat = self.db.load_chat(self.current_chat_id) - if chat: - title = chat['title'] - self.db.update_chat(self.current_chat_id, title, chat_data) - - # Now save notes and pins separately - if self.current_chat_id: - notes_data = [self.serialize_note(note) for note in notes] - self.db.save_notes(self.current_chat_id, notes_data) - - pins_data = [self.serialize_pin(pin) for pin in pins] - self.db.save_pins(self.current_chat_id, pins_data) - + saved_chat = self.db.load_chat(self.current_chat_id) + if not saved_chat: + raise ValueError("The active chat has been deleted from the database") + title = saved_chat['title'] + + saved_id = self.db.save_snapshot( + title, chat_data, snapshot['notes_data'], snapshot['pins_data'], + self.current_chat_id + ) + self.current_chat_id = saved_id + self._saved_snapshot = snapshot return chat_data + def deserialize_chart(self, data, scene): """Recreate a chart from serialized data""" chart = scene.add_chart(data['data'], QPointF( @@ -593,6 +661,7 @@ def load_chat(self, chat_id): # Set current chat ID and update connections self.current_chat_id = chat_id scene.update_connections() + self.mark_current_state_saved() except Exception as e: print(f"Error loading chat: {str(e)}") @@ -604,18 +673,13 @@ def load_chat(self, chat_id): scene.pins.clear() if self.window and hasattr(self.window, 'pin_overlay'): self.window.pin_overlay.clear_pins() + self.current_chat_id = None + self.mark_current_state_saved() raise # Return loaded chat data return chat def save_current_chat(self): - """Save the current chat session""" - if not self.window.chat_view.scene().nodes: - return - - try: - self.serialize_current_chat() - except Exception as e: - print(f"Error saving chat: {str(e)}") - raise + """Save the current canvas, even if it contains only notes or pins.""" + return self.serialize_current_chat() is not None diff --git a/graphite_app/graphite_app/graphite_ui.py b/graphite_app/graphite_app/graphite_ui.py index a9c4180..225d3a5 100644 --- a/graphite_app/graphite_app/graphite_ui.py +++ b/graphite_app/graphite_app/graphite_ui.py @@ -500,19 +500,25 @@ def new_chat(self): 'Start a new chat? This will clear the current session.', QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No ) - - if reply == QMessageBox.StandardButton.Yes: - self.session_manager.current_chat_id = None - scene = self.session_manager.window.chat_view.scene() - scene.clear() - scene.nodes.clear() - scene.connections.clear() - scene.frames.clear() - # <<< FIX: Reset stale reference and UI placeholder - self.session_manager.window.current_node = None - self.session_manager.window.message_input.setPlaceholderText("Type your message...") - self.close() - + if reply != QMessageBox.StandardButton.Yes: + return + window = self.session_manager.window + if not window.confirm_session_transition(): + return + + scene = window.chat_view.scene() + scene.clear() + scene.nodes.clear() + scene.connections.clear() + scene.frames.clear() + scene.pins.clear() + window.pin_overlay.clear_pins() + window.current_node = None + window.message_input.setPlaceholderText("Type your message...") + self.session_manager.current_chat_id = None + self.session_manager.mark_current_state_saved() + self.close() + def delete_selected(self): current_item = self.chat_list.currentItem() if current_item: @@ -544,15 +550,16 @@ def rename_selected(self): def load_chat(self, item): chat_id = item.data(Qt.ItemDataRole.UserRole) + window = self.session_manager.window + if not window.confirm_session_transition(): + return try: self.session_manager.load_chat(chat_id) + window.message_input.setPlaceholderText("Type your message...") self.close() except Exception as e: - QMessageBox.critical( - self, - "Error", - f"Failed to load chat: {str(e)}" - ) + QMessageBox.critical(self, "Error", f"Failed to load chat: {str(e)}") + class CustomScrollBar(QWidget): valueChanged = Signal(float) @@ -1992,39 +1999,25 @@ def regenerate_response(self): from graphite_agents import ChatWorkerThread if not self.node.parent_node: return - - user_message = self.node.parent_node.text - - history = [] - temp_node = self.node.parent_node - while temp_node: - if hasattr(temp_node, 'conversation_history'): - parent_history = [msg for msg in temp_node.conversation_history - if msg['role'] != 'assistant' or - temp_node.parent_node is not None] - history = parent_history + history - temp_node = temp_node.parent_node - + main_window = self.node.scene().window if not main_window: return - + + user_message = self.node.parent_node.text + history = main_window.build_conversation_history(self.node.parent_node) main_window.message_input.setEnabled(False) main_window.send_button.setEnabled(False) main_window.loading_overlay.show() - + main_window.chat_thread = ChatWorkerThread( - main_window.agent, - user_message, - history - ) - - main_window.chat_thread.finished.connect( - lambda response: self.handle_regenerated_response(response) + main_window.agent, user_message, history ) + main_window.chat_thread.finished.connect(self.handle_regenerated_response) main_window.chat_thread.error.connect(main_window.handle_error) main_window.chat_thread.start() - + + def handle_regenerated_response(self, new_response): try: self.node.text = new_response diff --git a/requirements-legacy.txt b/requirements-legacy.txt new file mode 100644 index 0000000..93a1c90 --- /dev/null +++ b/requirements-legacy.txt @@ -0,0 +1,5 @@ +# Only for the archived Graphite.py (PyQt5) implementation. +PyQt5 +ollama +matplotlib +qtawesome diff --git a/requirements.txt b/requirements.txt index 7d4a6c3..221eff3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,6 @@ -PyQt5 -ollama -matplotlib -qtawesome +PySide6>=6.6 +ollama>=0.3 +matplotlib>=3.8 +qtawesome>=1.3 +openai>=1.0 +google-generativeai>=0.8 diff --git a/tests/test_graphite_regressions.py b/tests/test_graphite_regressions.py new file mode 100644 index 0000000..f9fec63 --- /dev/null +++ b/tests/test_graphite_regressions.py @@ -0,0 +1,159 @@ +"""Regression tests for branch isolation and durable chat snapshots. + +Run from the repository root: + python -m unittest discover -s tests -v +""" +import copy +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "graphite_app" / "graphite_app")) + +import api_provider +import graphite_config as config +from graphite_agents import ChatAgent +from graphite_core import ChatDatabase, ChatSessionManager +from graphite_app import ChatWindow + + +def empty_chat(): + return { + "nodes": [], + "connections": [], + "frames": [], + "charts": [], + "view_state": {"zoom_factor": 1.0, "scroll_position": {"x": 0, "y": 0}}, + } + + +def note(text): + return { + "content": text, + "position": {"x": 10, "y": 20}, + "size": {"width": 120, "height": 60}, + "color": "#444444", + "header_color": None, + } + + +def pin(title): + return {"title": title, "note": "", "position": {"x": 1, "y": 2}} + + +class ChatBranchTests(unittest.TestCase): + def test_branches_do_not_leak_context_or_duplicate_current_message(self): + agent = ChatAgent("Test", "You are testing.") + observed = [] + + def provider_chat(task, messages): + self.assertEqual(task, config.TASK_CHAT) + observed.append(copy.deepcopy(messages)) + return {"message": {"content": "ok"}} + + branch_a = [ + {"role": "user", "content": "Common question"}, + {"role": "assistant", "content": "Common answer"}, + {"role": "user", "content": "Private branch A"}, + ] + branch_b = [ + {"role": "user", "content": "Common question"}, + {"role": "assistant", "content": "Common answer"}, + {"role": "user", "content": "Branch B"}, + ] + + with patch("graphite_agents.api_provider.chat", side_effect=provider_chat): + agent.get_response("Private branch A", branch_a) + agent.get_response("Branch B", branch_b) + + self.assertEqual(observed[0][1:], branch_a) + self.assertEqual(observed[1][1:], branch_b) + self.assertNotIn("Private branch A", str(observed[1])) + self.assertEqual(observed[1][-1]["content"], "Branch B") + + def test_ancestry_has_each_turn_once(self): + root = SimpleNamespace(text="Hello", is_user=True, parent_node=None) + answer = SimpleNamespace(text="Hi", is_user=False, parent_node=root) + child = SimpleNamespace(text="Follow up", is_user=True, parent_node=answer) + self.assertEqual( + ChatWindow.build_conversation_history(None, child), + [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Follow up"}, + ], + ) + + +class SnapshotTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.db = ChatDatabase.__new__(ChatDatabase) + self.db.db_path = Path(self.directory.name) / "test.db" + self.db.init_database() + + def tearDown(self): + self.directory.cleanup() + + def test_notes_only_canvas_can_be_saved(self): + manager = ChatSessionManager.__new__(ChatSessionManager) + manager.db = self.db + manager.current_chat_id = None + manager.title_generator = SimpleNamespace(generate_title=lambda msg: "Notes") + snapshot = {"chat_data": empty_chat(), "notes_data": [note("Important")], "pins_data": []} + manager._saved_snapshot = {"chat_data": empty_chat(), "notes_data": [], "pins_data": []} + manager.capture_current_snapshot = lambda: copy.deepcopy(snapshot) + + self.assertTrue(manager.has_unsaved_changes()) + self.assertTrue(manager.save_current_chat()) + self.assertEqual(self.db.load_notes(manager.current_chat_id)[0]["content"], "Important") + self.assertFalse(manager.has_unsaved_changes()) + + snapshot["notes_data"][0]["content"] = "Edited" + self.assertTrue(manager.has_unsaved_changes()) + self.assertTrue(manager.save_current_chat()) + self.assertEqual(self.db.load_notes(manager.current_chat_id)[0]["content"], "Edited") + + def test_failed_save_rolls_back_chat_notes_and_pins(self): + original = empty_chat() + chat_id = self.db.save_snapshot("Test", original, [note("Original")], [pin("First")]) + updated = {**original, "nodes": [{"text": "Changed"}]} + + with self.assertRaises(KeyError): + self.db.save_snapshot("Test", updated, [{"content": "missing properties"}], [], chat_id) + + self.assertEqual(self.db.load_chat(chat_id)["data"], original) + self.assertEqual(self.db.load_notes(chat_id)[0]["content"], "Original") + self.assertEqual(self.db.load_pins(chat_id)[0]["title"], "First") + + def test_empty_new_canvas_is_not_written(self): + manager = ChatSessionManager.__new__(ChatSessionManager) + manager.current_chat_id = None + manager.capture_current_snapshot = lambda: { + "chat_data": empty_chat(), "notes_data": [], "pins_data": [] + } + self.assertFalse(manager.save_current_chat()) + + +class SetupTests(unittest.TestCase): + def test_default_local_models_are_checked(self): + installed = SimpleNamespace(models=[ + SimpleNamespace(model=config.OLLAMA_MODELS[config.TASK_CHAT]) + ]) + with patch("api_provider.ollama.Client") as client: + client.return_value.list.return_value = installed + missing = api_provider.get_missing_ollama_models() + self.assertEqual( + missing, + sorted({ + config.OLLAMA_MODELS[config.TASK_TITLE], + config.OLLAMA_MODELS[config.TASK_CHART], + }), + ) + + +if __name__ == "__main__": + unittest.main()