PyLlama/main.py

import gi, ollama, asyncio, threading
gi.require_version("Gtk", "4.0")
from gi.repository import GLib, Gtk, Gdk

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._current_response = None
        
    @property
    def prompt_buffer(self):
        return self.prompt.get_buffer()

    def prompt_clear(self):
        self.prompt_buffer.set_text("")

    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)
        
        row = Gtk.ListBoxRow()
        row.set_child(label)
        self.chat_list.append(row)        

        self.chat_scroll_to_bottom()
        return label
        
    def chat_add_response_part(self, part):
        self._current_response.set_label(self._current_response.get_label() + part)
        self.chat_scroll_to_bottom()

    @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_add_message(text, sender=True)
        self.chat_start_response()

        app = self.get_application()
        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, 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)

    @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
        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 _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, options):
        GLib.idle_add(self.window.prompt_hold, True)
        try:
            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)            
            
if __name__ == "__main__":
    app = PyLlama()
    app.run()