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()
You cannot instantly terminate a thread in python.
This ollama API currently offers an async client, you can use the async client and cancel the Task, this should close the async connection almost instantly.
import tkinter as tk
import ollama
import threading
import asyncio
from typing import Optional
# Create the main window
root = tk.Tk()
root.title("Tkinter Button Example")
client = ollama.AsyncClient()
async def get_answer():
print("Start")
stream = await client.chat(
model='llama2',
messages=[{'role': 'user', 'content': 'write a story about earth'}],
stream=True,
)
async for chunk in stream:
print(chunk['message']['content'], end='', flush=True)
worker_loop: Optional[asyncio.AbstractEventLoop] = None
task_future: Optional[asyncio.Future] = None
def worker_function():
global worker_loop, task_future
worker_loop = asyncio.new_event_loop()
task_future = worker_loop.create_task(get_answer())
worker_loop.run_until_complete(task_future)
# Define the function to be called when the button is pressed
def on_button_click():
# the loop and the future are not threadsafe
worker_loop.call_soon_threadsafe(
lambda: task_future.cancel()
)
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=worker_function)
thread.start()
# Start the Tkinter event loop
root.mainloop()
this does raise an asyncio.CancelledError in the worker thread, so you may want to catch it if you don't want errors on your screen, and you may want to wrap the whole thing in a class rather than relying on globals.