如何立即终止线程？Ollama Python API+Tkinter实现Llama2流式响应中断

问题分析 #

当前代码使用全局stop变量控制线程终止，但ollama.chat()在生成第一个响应块前会阻塞线程，这段时间内无法检查stop变量，导致点击停止按钮后无法立即终止操作。

解决方案1：使用多进程替代线程 #

Python线程无法被强制终止，而进程可以通过terminate()方法立即停止，完美解决阻塞阶段无法终止的问题。

修改后的代码 #

import tkinter as tk
import ollama
import multiprocessing

# 主窗口创建
root = tk.Tk()
root.title("Ollama Llama2 流式响应控制")

def get_answer():
    print("开始生成响应")
    stream = ollama.chat(
        model='llama2',
        messages=[{'role': 'user', 'content': 'write a story about earth'}],
        stream=True,
    )
    for chunk in stream:
        print(chunk['message']['content'], end='', flush=True)

# 停止按钮回调
def on_button_click():
    if process.is_alive():
        process.terminate()
        print("\n已立即终止响应生成")

# 创建停止按钮
button = tk.Button(root, text="停止", command=on_button_click)
button.pack(pady=20)

# 使用多进程执行响应生成任务
process = multiprocessing.Process(target=get_answer)
process.start()

# 启动Tkinter事件循环
root.mainloop()

原理说明 #

multiprocessing.Process的terminate()方法会直接终止子进程，无论它当前处于阻塞还是运行状态，点击停止按钮后就能立即中断Ollama的请求，无需等待第一个响应块生成。

解决方案2：直接调用Ollama HTTP API #

通过直接调用Ollama的HTTP接口，利用requests库的流式请求特性，在停止时关闭请求连接，实现立即中断。

修改后的代码 #

import tkinter as tk
import requests
import threading
import json

# 主窗口创建
root = tk.Tk()
root.title("Ollama Llama2 流式响应控制")

current_response = None
current_session = None

def get_answer():
    global current_response, current_session
    print("开始生成响应")
    url = "http://localhost:11434/api/chat"
    payload = {
        "model": "llama2",
        "messages": [{"role": "user", "content": "write a story about earth"}],
        "stream": True
    }
    current_session = requests.Session()
    try:
        current_response = current_session.post(url, json=payload, stream=True)
        for chunk in current_response.iter_lines(decode_unicode=True):
            if chunk:
                data = json.loads(chunk)
                print(data['message']['content'], end='', flush=True)
    except requests.exceptions.RequestException:
        print("\n响应生成已停止")
    finally:
        if current_response:
            current_response.close()
        if current_session:
            current_session.close()

# 停止按钮回调
def on_button_click():
    global current_response, current_session
    if current_response:
        current_response.close()
    if current_session:
        current_session.close()
    print("\n已立即终止响应生成")

# 创建停止按钮
button = tk.Button(root, text="停止", command=on_button_click)
button.pack(pady=20)

# 使用线程执行响应生成任务
thread = threading.Thread(target=get_answer)
thread.start()

# 启动Tkinter事件循环
root.mainloop()

原理说明 #

直接调用Ollama的HTTP API后，点击停止按钮时关闭当前请求的响应对象和会话，强制中断流式请求，无需等待第一个响应块生成。

内容的提问来源于stack exchange，提问作者noocoder777