MX Linux Forum

Code: Select all

import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib
import requests
import threading
import subprocess
import json

# Ollama API endpoint
OLLAMA_API_URL = "http://localhost:11434/api/generate"

class ChatApp(Gtk.Window):
    def __init__(self):
        Gtk.Window.__init__(self, title="Ollama Chat")
        self.set_default_size(900, 600)

        # Main vertical box
        self.vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        self.add(self.vbox)

        # Dropdown for model selection
        self.model_store = Gtk.ListStore(str)  # Store model names
        self.model_combo = Gtk.ComboBox.new_with_model(self.model_store)
        renderer = Gtk.CellRendererText()
        self.model_combo.pack_start(renderer, True)
        self.model_combo.add_attribute(renderer, "text", 0)
        self.model_combo.set_active(0)  # Default to the first model
        self.model_combo.connect("changed", self.on_model_changed)
        self.vbox.pack_start(self.model_combo, False, True, 0)

        # Chat history (TextView inside ScrolledWindow)
        self.scrolled_window = Gtk.ScrolledWindow()
        self.scrolled_window.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)  # Vertical scrollbar only
        self.vbox.pack_start(self.scrolled_window, True, True, 0)

        self.chat_history = Gtk.TextView()
        self.chat_history.set_editable(False)
        self.chat_history.set_wrap_mode(Gtk.WrapMode.WORD)
        self.scrolled_window.add(self.chat_history)  # Add TextView to ScrolledWindow

        # Input box (Entry)
        self.input_entry = Gtk.Entry()
        self.input_entry.set_placeholder_text("Type your message here...")
        self.input_entry.connect("activate", self.on_input_activate)
        self.vbox.pack_start(self.input_entry, False, True, 0)

        # Send button
        self.send_button = Gtk.Button(label="Send")
        self.send_button.connect("clicked", self.on_send_clicked)
        self.vbox.pack_start(self.send_button, False, True, 0)

        # Initialize model list and selected model
        self.models = self.get_installed_models()
        self.selected_model = self.models[0] if self.models else None
        self.populate_model_dropdown()

    def get_installed_models(self):
        """Run `ollama list` to get installed models."""
        try:
            result = subprocess.run(["ollama", "list"], capture_output=True, text=True, check=True)
            lines = result.stdout.splitlines()
            models = [line.split()[0] for line in lines[1:]]  # Skip header line
            return models
        except Exception as e:
            print(f"Error fetching models: {e}")
            return []

    def populate_model_dropdown(self):
        """Populate the dropdown with installed models."""
        for model in self.models:
            self.model_store.append([model])
        if self.models:
            self.model_combo.set_active(0)  # Select the first model by default

    def on_model_changed(self, combo):
        """Handle model selection change."""
        tree_iter = combo.get_active_iter()
        if tree_iter is not None:
            model = self.model_store[tree_iter][0]
            self.selected_model = model
            self.append_to_chat_history(f"Switched to model: {model}\n")

    def on_input_activate(self, entry):
        self.send_message()

    def on_send_clicked(self, button):
        self.send_message()

    def send_message(self):
        user_input = self.input_entry.get_text().strip()
        if not user_input:
            return

        # Append user input to chat history
        self.append_to_chat_history(f"You: {user_input}\n")
        self.input_entry.set_text("")  # Clear input box

        # Send request to Ollama in a separate thread
        threading.Thread(target=self.query_ollama, args=(user_input,)).start()

    def query_ollama(self, prompt):
        payload = {
            "model": self.selected_model,
            "prompt": prompt,
            "stream": False
        }

        try:
            response = requests.post(OLLAMA_API_URL, json=payload)
            if response.status_code == 200:
                response_data = response.json()
                response_text = response_data.get("response", "").strip()
                GLib.idle_add(self.append_to_chat_history, f"{self.selected_model}: {response_text}\n\n")
            else:
                GLib.idle_add(self.append_to_chat_history, f"Error: {response.status_code} - {response.text}\n\n")
        except Exception as e:
            GLib.idle_add(self.append_to_chat_history, f"Error: {str(e)}\n\n")

    def append_to_chat_history(self, text):
        buffer = self.chat_history.get_buffer()
        buffer.insert(buffer.get_end_iter(), text)
        # Scroll to the end of the chat history
        self.chat_history.scroll_to_iter(buffer.get_end_iter(), 0, False, 0, 0)

# Run the application
if __name__ == "__main__":
    app = ChatApp()
    app.connect("destroy", Gtk.main_quit)
    app.show_all()
    Gtk.main()