PyLlama/pyllama.py

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()