122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
import gi
|
|
gi.require_version("Gtk", "4.0")
|
|
from gi.repository import GLib, Gtk
|
|
|
|
import ollama
|
|
|
|
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.set_title_widget(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)
|
|
|
|
scrolled = Gtk.ScrolledWindow()
|
|
scrolled.set_vexpand(True)
|
|
scrolled.set_child(self.listbox)
|
|
vbox.append(scrolled)
|
|
|
|
entry = Gtk.Entry()
|
|
entry_buffer = entry.get_buffer()
|
|
entry.set_placeholder_text("Type message and press Enter...")
|
|
entry.set_hexpand(True)
|
|
entry.set_icon_from_icon_name(Gtk.EntryIconPosition.SECONDARY, "go-next-symbolic")
|
|
entry.connect("activate", self.send_message)
|
|
entry.connect("icon_press", self.send_message)
|
|
|
|
vbox.append(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.get_application().send_prompt(self.model_dropdown.get_selected_item().get_string(), text)
|
|
|
|
class PyLlama(Gtk.Application):
|
|
def __init__(self):
|
|
super().__init__(application_id="eur.cfpi-fpsi.pyllama")
|
|
GLib.set_application_name("PyLlama")
|
|
|
|
def do_activate(self):
|
|
self.client = ollama.Client()
|
|
self.window = PyLlamaWindow(application=self, title="PyLlama")
|
|
self.window.present()
|
|
|
|
def send_prompt(self, model, text):
|
|
messages = [{
|
|
"role": "user",
|
|
"content": text
|
|
}]
|
|
|
|
for part in self.client.chat(model, messages=messages, stream=True):
|
|
self.window.add_response_part(part.message.content)
|
|
|
|
if __name__ == "__main__":
|
|
app = PyLlama()
|
|
app.run()
|