import gi, ollama, asyncio, threading gi.require_version("Gtk", "4.0") from gi.repository import GLib, Gtk, Gdk import util from chat_session import ChatSession @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") threads_spin = Gtk.Template.Child("threads_spin") temperature_spin = Gtk.Template.Child("temperature_spin") 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(), "num_thread": self.threads_spin.get_value_as_int(), "temperature": self.temperature_spin.get_value_as_int() } 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() @Gtk.Template.Callback() def gpu_layers_spin_changed(self, *args): pass @Gtk.Template.Callback() def low_vram_switch_set(self, *args): pass @Gtk.Template.Callback() def threads_spin_changed(self, *args): pass @Gtk.Template.Callback() def temperature_spin_changed(self, *args): pass 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()