136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
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.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")
|
|
self.model_dropdown.set_enable_search(True)
|
|
|
|
header = Gtk.HeaderBar()
|
|
header.set_show_title_buttons(True)
|
|
header.pack_start(self.model_dropdown)
|
|
|
|
self.set_titlebar(header)
|
|
|
|
vbox = Gtk.Box.new(Gtk.Orientation.VERTICAL, 6)
|
|
set_margin_all(vbox, 8)
|
|
|
|
self.listbox = Gtk.ListBox.new()
|
|
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 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=None, 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("" if not text else 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(text=None, sender=False)
|
|
|
|
self.listbox.grab_focus()
|
|
self.entry.set_sensitive(False)
|
|
|
|
app = self.get_application()
|
|
model = self.model_dropdown.get_selected_item().get_string()
|
|
app.loop.call_soon_threadsafe(lambda: asyncio.create_task(
|
|
app.send_prompt(model, text)
|
|
))
|
|
|
|
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):
|
|
messages = [{ "role": "user", "content": text, "options": { "num_gpu": 100, "main_gpu": 0 } }]
|
|
async for part in await self.client.chat(model=model, messages=messages, stream=True):
|
|
GLib.idle_add(self.window.add_response_part, part.message.content)
|
|
GLib.idle_add(self.window.entry.set_sensitive, True)
|
|
|
|
if __name__ == "__main__":
|
|
app = PyLlama()
|
|
app.run()
|