I’m using ollama to stream a response from llama2 large language model. The functionality I need is, when I click the stop button, it should stop the thread immediately. The code below works but the problem is, it always waits for the first chunk to be available before it checks if stop is true. Sometimes the response can take some time before the first chunk is ready. It shouldn’t wait for the first chunk. It should immediately stop the thread when stop button is clicked.

My code:

import tkinter as tk
import ollama
import threading

stop = False


# Create the main window
root = tk.Tk()
root.title("Tkinter Button Example")

def get_answer():
    print("Start")
    stream = ollama.chat(
        model='llama2',
        messages=[{'role': 'user', 'content': 'write a story about earth'}],
        stream=True,
    )
    for chunk in stream:
        if stop is False:
            print(chunk['message']['content'], end='', flush=True)

        else:
            print("Stopped")   
            return
# Define the function to be called when the button is pressed
def on_button_click():
    global stop
    stop = True
    print("Button clicked!")

# Create the button
button = tk.Button(root, text="Stop", command=on_button_click)

# Place the button on the window
button.pack(pady=20)

thread = threading.Thread(target=get_answer)
thread.start()

# Start the Tkinter event loop
root.mainloop()