如何解决Ollama Python库Asyncio任务无法提前取消的问题？

问题现象 #

使用Ollama Python API调用大语言模型生成回答，点击停止按钮仅能在响应开始打印后取消任务；若任务处于处理准备阶段（如模型加载、请求预处理），取消操作无效，最终仍会输出响应。例如点击按钮后，prompt_mistral("Testing")仍会执行并打印内容。

问题根源 #

1. 全局变量跟踪不足：仅用单个task_future跟踪最外层任务，内部嵌套调用的prompt_mistral等任务未被纳入跟踪，导致外层任务取消后，内部任务仍继续执行。
2. 准备阶段无取消检查：await client.chat()的准备阶段（模型初始化、请求构建）未触发异步取消点，此时调用cancel()无法触发CancelledError。
3. 线程与全局变量冲突：on_button_click函数中启动新的phi3任务会覆盖全局的worker_loop和task_future，导致原任务引用丢失，无法取消。

解决方案 #

1. 使用任务集合跟踪所有活跃异步任务，而非单个变量。
2. 在异步函数中，在关键节点（如await client.chat()之后）主动检查取消状态。
3. 修正停止按钮逻辑，专注取消现有任务，避免启动新任务干扰。
4. 优化线程管理，确保每个任务的事件循环和任务引用不被覆盖。

修改后的代码 #

import ollama
import asyncio
import threading
from typing import Optional, Set
import tkinter as tk

# 主窗口初始化
root = tk.Tk()
root.title("LLM 生成控制")

# 全局变量：跟踪所有活跃的事件循环和任务
active_loops: Set[asyncio.AbstractEventLoop] = set()
active_tasks: Set[asyncio.Future] = set()
lock = threading.Lock()

async def get_answer_from_phi3():
    print("Phi3 开始处理")
    messages = [{"role": "system", "content": "Hello"}]
    client = ollama.AsyncClient()

    try:
        # 发起请求后立即检查取消状态
        stream = await client.chat(
            model='phi3',
            messages=messages,
            stream=True,
            options={"top_k": 1}
        )
        # 检查是否已被取消
        if asyncio.current_task().cancelled():
            raise asyncio.CancelledError()
        
        async for chunk in stream:
            print(chunk['message']['content'], end='', flush=True)
            # 每处理一个chunk都检查取消状态
            if asyncio.current_task().cancelled():
                raise asyncio.CancelledError()

    except asyncio.exceptions.CancelledError:
        print("\nPhi3 任务已取消")
    except Exception as e:
        print(f"Phi3 错误: {e}")
        return "处理请求时发生错误"

async def prompt_mistral(query):
    print("Mistral 开始处理")
    messages = [{"role": "assistant", "content": "Write a song that celebrates the beauty, diversity, and importance of our planet, Earth. The song should evoke vivid imagery of the natural world, from lush forests and majestic mountains to serene oceans and vast deserts. It should capture the essence of Earth as a living, breathing entity that sustains all forms of life. Incorporate themes of harmony, unity, and interconnectedness, emphasizing how all elements of nature are intertwined and how humanity is an integral part of this complex web. The lyrics should reflect a sense of wonder and appreciation for the planet's resources and ecosystems, highlighting the delicate balance that sustains life. Include references to various landscapes, climates, and wildlife, painting a picture of Earth's diverse environments. The song should also touch on the responsibility we have to protect and preserve the planet for future generations, addressing issues like climate change, deforestation, pollution, and conservation efforts. Use poetic language and metaphors to convey the grandeur and fragility of Earth, and infuse the song with a hopeful and inspiring tone that encourages listeners to take action in safeguarding our shared home. The melody should be uplifting and emotionally resonant, complementing the powerful message of the lyrics"}]
    generated_answer = ''
    client = ollama.AsyncClient()

    try:
        stream = await client.chat(
            model='mistral',
            messages=messages,
            stream=True,
            options={"top_k": 1}
        )
        if asyncio.current_task().cancelled():
            raise asyncio.CancelledError()
        
        async for chunk in stream:
            generated_answer += chunk['message']['content']
            print(chunk['message']['content'])
            if asyncio.current_task().cancelled():
                raise asyncio.CancelledError()

    except asyncio.exceptions.CancelledError:
        print("\nMistral 任务已取消")
    except Exception as e:
        print(f"Mistral 错误: {e}")
        return "处理请求时发生错误"

async def prompt_llama_task(message):
    print("Llama2 开始处理")
    messages = [{"role": "assistant", "content": message}]
    client = ollama.AsyncClient()
    generated_answer = ''

    try:
        stream = await client.chat(
            model='llama2',
            messages=messages,
            stream=True,
            options={"top_k": 1}
        )
        if asyncio.current_task().cancelled():
            raise asyncio.CancelledError()
        
        async for chunk in stream:
            generated_answer += chunk['message']['content']
            print(chunk['message']['content'])
            if asyncio.current_task().cancelled():
                raise asyncio.CancelledError()
        
        if "help" in generated_answer:
            # 将嵌套任务加入全局跟踪集合
            mistral_task = asyncio.create_task(prompt_mistral("Testing"))
            with lock:
                active_tasks.add(mistral_task)
            await mistral_task
            with lock:
                active_tasks.discard(mistral_task)
        else:
            print(generated_answer)

    except asyncio.exceptions.CancelledError:
        print("\nLlama2 任务已取消")
    except Exception as e:
        print(f"Llama2 错误: {e}")
        return "处理请求时发生错误"

def prompt_llama(message):
    def worker_function():
        loop = asyncio.new_event_loop()
        with lock:
            active_loops.add(loop)
        task = loop.create_task(prompt_llama_task(message))
        with lock:
            active_tasks.add(task)
        try:
            loop.run_until_complete(task)
        finally:
            with lock:
                active_tasks.discard(task)
                active_loops.discard(loop)
            loop.close()

    print("启动Llama2线程")
    thread = threading.Thread(target=worker_function)
    thread.start()

def on_button_click():
    # 取消所有活跃任务
    with lock:
        for task in active_tasks:
            if not task.done():
                task.cancel()
        # 停止所有事件循环（可选，根据需求）
        for loop in active_loops:
            loop.call_soon_threadsafe(loop.stop)
    print("已触发所有任务取消")

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

# 启动Llama2任务
prompt_llama("Hi")

# 启动Tkinter主循环
root.mainloop()

关键修改点 #

• 新增active_loops和active_tasks集合，配合线程锁跟踪所有活跃的事件循环和任务，避免引用丢失。
• 在每个异步函数的await client.chat()后、async for循环内添加取消状态检查，确保准备阶段结束后能及时响应取消。
• 嵌套调用prompt_mistral时，将其任务加入全局跟踪集合，确保外层任务取消时，内部任务也能被取消。
• 修正on_button_click函数，仅负责取消所有活跃任务，不再启动新任务，避免干扰原任务的取消逻辑。

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