Major speed improvement, split the project into several files

chat.py

@@ -0,0 +1,43 @@
+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)

main.py

@@ -0,0 +1,127 @@
+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
+        
+class PyLlamaWindow(Gtk.ApplicationWindow):
+    def __init__(self, application=None, title=None):
+        super().__init__(application=application, title=title)	
+        
+        self.set_default_size(800, 600)
+        
+        self.model_store = Gtk.StringList()
+        self.model_dropdown = Gtk.DropDown(tooltip_text="Select model")
+        self.model_dropdown.set_model(self.model_store)
+
+        model_preferences = Gtk.Button(icon_name="view-more-symbolic",
+                                       tooltip_text="Model preferences")
+        model_preferences.connect("clicked", self.on_preferences_clicked)
+
+        self.model_prefs = ModelPreferenceView(position=Gtk.PositionType.BOTTOM,
+                                               on_gpu_change=self.on_gpu_changed)
+        self.model_prefs.set_parent(model_preferences)
+        
+        header = Gtk.HeaderBar(show_title_buttons=True)
+        header.pack_start(self.model_dropdown)
+        header.pack_start(model_preferences)
+        
+        self.set_titlebar(header)
+
+        vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
+        util.set_margin_all(vbox, 8)
+
+        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()
+
+        app = self.get_application()
+        model = self.model_dropdown.get_selected_item().get_string()
+
+        async def task():
+            await app.send_prompt(model, text)
+
+        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):
+        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.client = ollama.AsyncClient()        
+        self.async_loop = asyncio.new_event_loop()
+        self.current_async_task = None
+        threading.Thread(target=self._run_asyncio_loop, daemon=True).start()
+        
+    def do_activate(self):
+        self.window = PyLlamaWindow(application=self, title="PyLlama")
+        self.window.present()
+
+    def _run_asyncio_loop(self):
+        asyncio.set_event_loop(self.async_loop)
+        self.async_loop.run_forever()
+        return True
+
+    def stop_generation(self):
+        if self.current_async_task and not self.current_async_task.done():
+            self.current_async_task.cancel()
+            
+    async def send_prompt(self, model, text=""):
+        messages = [{ "role": "user", "content": text }]
+        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)
+        except asyncio.CancelledError:
+            raise
+        finally:
+            GLib.idle_add(self.window.prompt.unhold, True)            
+            
+if __name__ == "__main__":
+    app = PyLlama()
+    app.run()

preferences.py

@@ -0,0 +1,24 @@
+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)

prompt.py

@@ -0,0 +1,65 @@
+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
+

pyllama.py

@@ -1,174 +0,0 @@
-import gi, ollama, asyncio
-gi.require_version("Gtk", "4.0")
-from gi.repository import GLib, Gtk
-
-def set_margin_all(widget, m):
-    widget.set_margin_start(m)
-    widget.set_margin_top(m)
-    widget.set_margin_end(m)
-    widget.set_margin_bottom(m)
-
-class PyLlamaWindow(Gtk.ApplicationWindow):
-    def __init__(self, application=None, title=None):
-        super().__init__(application=application, title=title)	
-        self.set_default_size(800, 600)
-        self.cur_ai_response = None
-        self.gpu_spin = None
-        
-        self.model_store = Gtk.StringList()
-        self.model_dropdown = Gtk.DropDown()
-        self.model_dropdown.set_model(self.model_store)
-        self.model_dropdown.set_tooltip_text("Select model")
-
-        model_preferences = Gtk.Button()
-        icon = Gtk.Image.new_from_icon_name("view-more-symbolic")
-        model_preferences.set_child(icon)
-        model_preferences.set_tooltip_text("Preferences")
-        model_preferences.connect("clicked", self.on_preferences_clicked)
-
-        header = Gtk.HeaderBar()
-        header.set_show_title_buttons(True)
-        header.pack_start(self.model_dropdown)
-        header.pack_start(model_preferences)
-        
-        self.set_titlebar(header)
-        
-        vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
-        set_margin_all(vbox, 8)
-        
-        self.listbox = Gtk.ListBox()
-        self.listbox.set_selection_mode(Gtk.SelectionMode.NONE)
-        self.listbox.set_focusable(True)
-        
-        scrolled = Gtk.ScrolledWindow()
-        scrolled.set_vexpand(True)
-        scrolled.set_child(self.listbox)
-        vbox.append(scrolled)
-
-        # self.prompt = Gtk.TextView()
-        # self.prompt.set_hexpand(True)
-        
-        self.entry = Gtk.Entry()
-        self.entry.set_placeholder_text("Type message and press Enter...")
-        self.entry.set_hexpand(True)
-        self.entry.set_icon_from_icon_name(Gtk.EntryIconPosition.SECONDARY, "go-next-symbolic")
-        self.entry.connect("activate", self.send_message)
-        self.entry.connect("icon_press", self.send_message)
-
-        vbox.append(self.entry)
-        self.set_child(vbox)
-
-        self.current_model = None
-        self.populate_models()
-
-    def hold_prompt(self, *_):
-        self.entry.set_placeholder_text("Wait for response...")
-        self.listbox.grab_focus()
-        self.entry.set_sensitive(False)
-
-    def unhold_prompt(self, *_):
-        self.entry.set_placeholder_text("Type message and press Enter...")
-        self.entry.set_sensitive(True)
-        self.entry.grab_focus()
-        
-    def on_preferences_clicked(self, button):
-        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
-        set_margin_all(box, 8)
-
-        label = Gtk.Label(label="GPU Layers")
-        label.set_xalign(0.0)
-        
-        self.gpu_spin = Gtk.SpinButton(adjustment=Gtk.Adjustment(
-            value=10, lower=0, upper=100, step_increment=1
-        ), digits=0, hexpand=True)
-
-        box.append(label)
-        box.append(self.gpu_spin)
-
-        preferences_popover = Gtk.Popover(position=Gtk.PositionType.BOTTOM)
-        preferences_popover.set_child(box)
-        preferences_popover.set_parent(button)
-        preferences_popover.popup()
-        
-    def populate_models(self):
-        try:
-            models = ollama.list()
-        except Exception as e:
-            models = []
-            print ("Failed to fetch models: ", e)
-
-        names = [m.model for m in models["models"]]
-        self.model_store.splice(0, len(list(self.model_store)), names)
-
-    def add_message(self, text="", sender=True):
-        label = Gtk.Label()
-        label.set_selectable(True)
-        label.set_wrap(True)
-        label.set_xalign(1.0 if sender else 0.0)
-        label.set_text(text)
-        set_margin_all(label, 8)
-        
-        row = Gtk.ListBoxRow()
-        row.set_child(label)
-
-        self.listbox.append(row)
-        
-        adj = self.listbox.get_adjustment()
-        adj.set_value(adj.get_upper())
-
-        return label
-
-    def add_response_part(self, part):
-        if self.cur_ai_response == None:
-            return
-
-        self.cur_ai_response.set_label(self.cur_ai_response.get_label() + part)
-        
-    def send_message(self, entry, _=None):
-        buffer = entry.get_buffer()
-        text = buffer.get_text()
-        if text == None or text == "":
-            return
-
-        buffer.delete_text(0, -1) # Delete all text
-        self.add_message(text=text)
-        self.cur_ai_response = self.add_message(sender=False)
-
-        self.listbox.grab_focus()
-        self.entry.set_sensitive(False)
-        
-        app = self.get_application()
-        model = self.model_dropdown.get_selected_item().get_string()
-        num_gpu = self.gpu_spin.get_value_as_int() if self.gpu_spin else 10
-        app.loop.call_soon_threadsafe(lambda: asyncio.create_task(
-            app.send_prompt(model, text, num_gpu)
-        ))
-        
-class PyLlama(Gtk.Application):
-    def __init__(self):
-        super().__init__(application_id="eur.cfpi-fpsi.pyllama")
-        GLib.set_application_name("PyLlama")
-        
-        self.client = ollama.AsyncClient()
-        self.loop = asyncio.new_event_loop()
-        GLib.timeout_add(100, self.run_asyncio_loop)
-        
-    def do_activate(self):
-        self.window = PyLlamaWindow(application=self, title="PyLlama")
-        self.window.present()
-
-    def run_asyncio_loop(self):
-        self.loop.stop()
-        self.loop.run_forever()
-        return True
-    
-    async def send_prompt(self, model, text="", num_gpu=10):
-        messages = [{ "role": "user", "content": text }]
-        GLib.idle_add(self.window.hold_prompt, True)
-        async for part in await self.client.chat(model=model, messages=messages, stream=True, options={ "num_gpu": num_gpu }):
-            GLib.idle_add(self.window.add_response_part, part.message.content)
-        GLib.idle_add(self.window.unhold_prompt, True)
-        
-if __name__ == "__main__":
-    app = PyLlama()
-    app.run()

util.py

@@ -0,0 +1,5 @@
+def set_margin_all(widget, m):
+    widget.set_margin_start(m)
+    widget.set_margin_top(m)
+    widget.set_margin_end(m)
+    widget.set_margin_bottom(m)