diff --git a/__pycache__/chat.cpython-313.pyc b/__pycache__/chat.cpython-313.pyc new file mode 100644 index 0000000..7fdb7b8 Binary files /dev/null and b/__pycache__/chat.cpython-313.pyc differ diff --git a/__pycache__/preferences.cpython-313.pyc b/__pycache__/preferences.cpython-313.pyc new file mode 100644 index 0000000..12bc8ea Binary files /dev/null and b/__pycache__/preferences.cpython-313.pyc differ diff --git a/__pycache__/prompt.cpython-313.pyc b/__pycache__/prompt.cpython-313.pyc new file mode 100644 index 0000000..d6c766f Binary files /dev/null and b/__pycache__/prompt.cpython-313.pyc differ diff --git a/__pycache__/util.cpython-313.pyc b/__pycache__/util.cpython-313.pyc new file mode 100644 index 0000000..b20692b Binary files /dev/null and b/__pycache__/util.cpython-313.pyc differ diff --git a/chat.py b/chat.py new file mode 100644 index 0000000..d92e4ce --- /dev/null +++ b/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) diff --git a/main.py b/main.py new file mode 100644 index 0000000..a8c242d --- /dev/null +++ b/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() diff --git a/preferences.py b/preferences.py new file mode 100644 index 0000000..fcb4fb4 --- /dev/null +++ b/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) diff --git a/prompt.py b/prompt.py new file mode 100644 index 0000000..2c6a8cd --- /dev/null +++ b/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 + diff --git a/pyllama.py b/pyllama.py deleted file mode 100644 index 833cbe9..0000000 --- a/pyllama.py +++ /dev/null @@ -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() diff --git a/util.py b/util.py new file mode 100644 index 0000000..fb6b602 --- /dev/null +++ b/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)