diff --git a/.gitignore b/.gitignore index ba0430d..352a86b 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -__pycache__/ \ No newline at end of file +__pycache__/ +*~ \ No newline at end of file diff --git a/chat.py b/chat.py deleted file mode 100644 index d92e4ce..0000000 --- a/chat.py +++ /dev/null @@ -1,43 +0,0 @@ -import gi -from gi.repository import GLib, Gtk, Gdk - -import util - -class ChatView(Gtk.ScrolledWindow): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.cur_ai_response = None - - self.listbox = Gtk.ListBox() - self.listbox.set_selection_mode(Gtk.SelectionMode.NONE) - self.listbox.set_focusable(True) - - self.set_child(self.listbox) - - def scroll_to_bottom(self, *args): - adj = self.get_vadjustment() - adj.set_value(adj.get_upper() - adj.get_page_size()) - - def start_response(self): - self.cur_ai_response = self.add_message(sender=False) - return self.cur_ai_response - - def add_message(self, text="", sender=True): - label = Gtk.Label(selectable=True, wrap=True, xalign=1.0 if sender else 0.0) - label.set_text(text) - util.set_margin_all(label, 8) - - row = Gtk.ListBoxRow() - row.set_child(label) - self.listbox.append(row) - - GLib.idle_add(self.scroll_to_bottom) - return label - - def add_response_part(self, part): - if self.cur_ai_response is None: - return - - self.cur_ai_response.set_label(self.cur_ai_response.get_label() + part) - GLib.idle_add(self.scroll_to_bottom) diff --git a/main.py b/main.py index a8c242d..77c41d0 100644 --- a/main.py +++ b/main.py @@ -2,96 +2,128 @@ import gi, ollama, asyncio, threading gi.require_version("Gtk", "4.0") from gi.repository import GLib, Gtk, Gdk -from prompt import PromptView -from chat import ChatView -from preferences import ModelPreferenceView import util - + +@Gtk.Template(filename="pyllama.ui") class PyLlamaWindow(Gtk.ApplicationWindow): + __gtype_name__ = "main_window" + + model_list = Gtk.Template.Child("model_list") + model_parameters = Gtk.Template.Child("model_parameters") + chat_scroll = Gtk.Template.Child("chat_scroll") + chat_list = Gtk.Template.Child("chat_list") + prompt = Gtk.Template.Child("prompt") + prompt_controller = Gtk.Template.Child("prompt_controller") + send_button = Gtk.Template.Child("send_button") + cancel_button = Gtk.Template.Child("cancel_button") + gpu_layers_spin = Gtk.Template.Child("gpu_layers_spin") + low_vram_switch = Gtk.Template.Child("low_vram_switch") + def __init__(self, application=None, title=None): super().__init__(application=application, title=title) + self.model_list.set_model(self.get_application().model_store) + self.prompt.add_controller(self.prompt_controller) - self.set_default_size(800, 600) + self._current_response = None - self.model_store = Gtk.StringList() - self.model_dropdown = Gtk.DropDown(tooltip_text="Select model") - self.model_dropdown.set_model(self.model_store) + @property + def prompt_buffer(self): + return self.prompt.get_buffer() - model_preferences = Gtk.Button(icon_name="view-more-symbolic", - tooltip_text="Model preferences") - model_preferences.connect("clicked", self.on_preferences_clicked) + def prompt_clear(self): + self.prompt_buffer.set_text("") - self.model_prefs = ModelPreferenceView(position=Gtk.PositionType.BOTTOM, - on_gpu_change=self.on_gpu_changed) - self.model_prefs.set_parent(model_preferences) + def prompt_get_text(self): + start = self.prompt_buffer.get_start_iter() + end = self.prompt_buffer.get_end_iter() + return self.prompt_buffer.get_text(start, end, False).strip() + + def prompt_hold(self, *args): + self.grab_focus() + self.prompt.set_sensitive(False) + self.send_button.set_visible(False) + self.cancel_button.set_visible(True) + + def prompt_unhold(self, *args): + self.prompt.set_sensitive(True) + self.prompt.grab_focus() + self.send_button.set_visible(True) + self.cancel_button.set_visible(False) + + def chat_scroll_to_bottom(self): + adj = self.chat_scroll.get_vadjustment() + adj.set_value(adj.get_upper() - adj.get_page_size()) + + def chat_start_response(self): + self._current_response = self.chat_add_message(sender=False) + + def chat_add_message(self, text="", sender=True): + label = Gtk.Label(selectable=True, wrap=True, xalign=1.0 if sender else 0.0) + label.set_text(text) + util.set_margin_all(label, 8) - header = Gtk.HeaderBar(show_title_buttons=True) - header.pack_start(self.model_dropdown) - header.pack_start(model_preferences) + row = Gtk.ListBoxRow() + row.set_child(label) + self.chat_list.append(row) + + self.chat_scroll_to_bottom() + return label - self.set_titlebar(header) + def chat_add_response_part(self, part): + self._current_response.set_label(self._current_response.get_label() + part) + self.chat_scroll_to_bottom() - vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6) - util.set_margin_all(vbox, 8) + @Gtk.Template.Callback() + def prompt_key_pressed(self, _a, keyval, _b, state): + match keyval: + case Gdk.KEY_Return if not state & Gdk.ModifierType.SHIFT_MASK: + self.send_button_clicked(self.send_button) + return True + case _: + return False + + @Gtk.Template.Callback() + def send_button_clicked(self, button): + text = self.prompt_get_text() + self.prompt_clear() - self.chat = ChatView(vexpand=True) - vbox.append(self.chat) - - self.prompt = PromptView(spacing=6) - self.prompt.on_send = self.on_send - self.prompt.on_cancel = self.on_cancel - vbox.append(self.prompt) - - self.set_child(vbox) - self.populate_models() - - def populate_models(self): - try: - models = ollama.list() - names = [m.model for m in models["models"]] - except Exception as e: - print ("Failed to fetch models: ", e) - names = [] - - self.model_store.splice(0, len(self.model_store), names) - - def on_gpu_changed(self, value): - self.get_application().num_gpu = value - - def on_preferences_clicked(self, button): - self.model_prefs.popup() - - def on_send(self, *args): - text = self.prompt.get_text() - if not text: - return - - self.prompt.clear() - self.chat.add_message(text, sender=True) - self.chat.start_response() + self.chat_add_message(text, sender=True) + self.chat_start_response() app = self.get_application() - model = self.model_dropdown.get_selected_item().get_string() + model = self.model_list.get_selected_item().get_string() + app.history_push(role="user", text=text) + + options = { + "num_gpu": self.gpu_layers_spin.get_value_as_int(), + "low_vram": self.low_vram_switch.get_state() + } + async def task(): - await app.send_prompt(model, text) + await app.send_prompt(model, options) if app.current_async_task and not app.current_async_task.done(): app.current_async_task.cancel() app.current_async_task = asyncio.run_coroutine_threadsafe(task(), app.async_loop) - def on_cancel(self, *args): + @Gtk.Template.Callback() + def cancel_button_clicked(self, *args): app = self.get_application() app.stop_generation() - + class PyLlama(Gtk.Application): def __init__(self): super().__init__(application_id="eur.cfpi-fpsi.pyllama") GLib.set_application_name("PyLlama") self.num_gpu = 15 - + self.history_size = 4 + self.history = [] + self.model_store = Gtk.StringList() + self._populate_models() + self.client = ollama.AsyncClient() self.async_loop = asyncio.new_event_loop() self.current_async_task = None @@ -106,21 +138,42 @@ class PyLlama(Gtk.Application): self.async_loop.run_forever() return True + def _populate_models(self): + try: + models = ollama.list() + names = [m.model for m in models["models"]] + except Exception as e: + print ("Failed to fetch models: ", e) + names = [] + + self.model_store.splice(0, len(self.model_store), names) + def stop_generation(self): if self.current_async_task and not self.current_async_task.done(): self.current_async_task.cancel() + + def history_push(self, role, text): + self.history.append({"role": role, "content": text}) + while len(self.history) > self.history_size: + self.history.pop(0) - async def send_prompt(self, model, text=""): - messages = [{ "role": "user", "content": text }] - GLib.idle_add(self.window.prompt.hold, True) - + async def send_prompt(self, model, options): + GLib.idle_add(self.window.prompt_hold, True) try: - async for part in await self.client.chat(model=model, messages=messages, stream=True, options={ "num_gpu": self.num_gpu }): - GLib.idle_add(self.window.chat.add_response_part, part.message.content) + messages = list(self.history)[-self.history_size:] + async for part in await self.client.chat(model=model, messages=messages, stream=True, options=options): + chunk = part.message.content + + if not self.history or self.history[-1]["role"] != "assistant": + self.history.append({"role": "assistant", "content": chunk}) + else: + self.history[-1]["content"] += chunk + + GLib.idle_add(self.window.chat_add_response_part, part.message.content) except asyncio.CancelledError: raise finally: - GLib.idle_add(self.window.prompt.unhold, True) + GLib.idle_add(self.window.prompt_unhold, True) if __name__ == "__main__": app = PyLlama() diff --git a/preferences.py b/preferences.py deleted file mode 100644 index fcb4fb4..0000000 --- a/preferences.py +++ /dev/null @@ -1,24 +0,0 @@ -import gi -from gi.repository import Gtk - -import util - -class ModelPreferenceView(Gtk.Popover): - def __init__(self, on_gpu_change=None, *args, **kwargs): - super().__init__(*args, **kwargs) - self.on_gpu_change = on_gpu_change - - box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6) - util.set_margin_all(box, 8) - - label = Gtk.Label(label="GPU Layers") - box.append(label) - - gpu_spin = Gtk.SpinButton(adjustment=Gtk.Adjustment( - value=15, lower=0, upper=50, step_increment=1, - ), digits=0, hexpand=True) - - gpu_spin.connect("value-changed", lambda spin: self.on_gpu_change(spin.get_value_as_int())) - box.append(gpu_spin) - - self.set_child(box) diff --git a/prompt.py b/prompt.py deleted file mode 100644 index 2c6a8cd..0000000 --- a/prompt.py +++ /dev/null @@ -1,65 +0,0 @@ -import gi -from gi.repository import Gtk, Gdk - -class PromptView(Gtk.Box): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.prompt = Gtk.TextView(left_margin=12, top_margin=14, - right_margin=12, bottom_margin=14) - - prompt_scroll = Gtk.ScrolledWindow(hexpand=True) - prompt_scroll.set_propagate_natural_height(True) - prompt_scroll.set_max_content_height(150) - prompt_scroll.set_child(self.prompt) - - key_controller = Gtk.EventControllerKey() - key_controller.connect("key-pressed", self.on_key_pressed) - self.prompt.add_controller(key_controller) - - self.prompt_button = Gtk.Button(icon_name="go-next-symbolic") - self.btn_handler = self.prompt_button.connect("clicked", self.on_send) - - self.append (prompt_scroll) - self.append (self.prompt_button) - - @property - def buffer(self): - return self.prompt.get_buffer() - - def clear(self): - self.buffer.set_text("") - - def get_text(self): - return self.buffer.get_text(self.buffer.get_start_iter(), - self.buffer.get_end_iter(), False).strip() - - def hold(self, *args): - self.grab_focus() - self.prompt.set_sensitive(False) - self.prompt_button.set_icon_name("process-stop-symbolic") - self.prompt_button.disconnect(self.btn_handler) - self.btn_handler = self.prompt_button.connect("clicked", self.on_cancel) - - def unhold(self, *args): - self.prompt.set_sensitive(True) - self.prompt.grab_focus() - self.prompt_button.set_icon_name("go-next-symbolic") - self.prompt_button.disconnect(self.btn_handler) - self.btn_handler = self.prompt_button.connect("clicked", self.on_send) - - - def on_key_pressed(self, _a, keyval, _b, state): - if keyval == Gdk.KEY_Return: - if state & Gdk.ModifierType.SHIFT_MASK: - return False - self.on_send() - return True - return False - - def on_send(self, *args): - pass - - def on_cancel(self, *args): - pass - diff --git a/pyllama.cmb b/pyllama.cmb new file mode 100644 index 0000000..f53c030 --- /dev/null +++ b/pyllama.cmb @@ -0,0 +1,6 @@ + + + + + + diff --git a/pyllama.ui b/pyllama.ui new file mode 100644 index 0000000..48f35da --- /dev/null +++ b/pyllama.ui @@ -0,0 +1,176 @@ + + + + + + + + + + + + + diff --git a/shell.nix b/shell.nix index 8a597df..a0db324 100644 --- a/shell.nix +++ b/shell.nix @@ -2,6 +2,7 @@ let pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/tarball/nixos-25.11") {}; in pkgs.mkShell { packages = with pkgs; [ + cambalache gtk4 gobject-introspection