pythonmultithreadingtkinterllamaollama

how to instantly terminate a thread? Using ollama python API with tkinter to stream a response from llama2


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()

Solution

  • 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.