Как немедленно отменить задачу Asyncio, которая использует библиотеку Ollama Python для генерации ответа?

Нам нужно обновить функции on_button_click() и prompt_llama().

def on_button_click():
    global worker_loop, task_future
    if worker_loop and task_future:
        worker_loop.call_soon_threadsafe(task_future.cancel)

    def phi3_worker_function():
        global worker_loop, task_future
        worker_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(worker_loop)
        task_future = worker_loop.create_task(get_answer_from_phi3())
        worker_loop.run_until_complete(task_future)

    print("Starting thread")
    thread = threading.Thread(target=phi3_worker_function)
    thread.start()

def prompt_llama(message):
    async def prompt():
        messages = [{"role": "assistant", "content": message}]
        try:
            client = ollama.AsyncClient()
            stream = await client.chat(
                model='llama2',
                messages=messages,
                stream=True,
                options = {"top_k": 1}
            )

            generated_answer = ''
            async for chunk in stream:
                generated_answer += chunk['message']['content']
                print(chunk['message']['content'])

            if "help" in generated_answer:
                await prompt_mistral("Testing")
            else:
                print(generated_answer)

        except asyncio.exceptions.CancelledError:
            print("Cancelled")
        except Exception as e:
            print(f"Error: {e}")

    def mistral_worker_function():
        global worker_loop, task_future
        worker_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(worker_loop)
        task_future = worker_loop.create_task(prompt())
        worker_loop.run_until_complete(task_future)

    print("Starting thread")
    thread = threading.Thread(target=mistral_worker_function)
    thread.start()

Чтобы немедленно отменить задачу asyncio, которая использует библиотеку Ollama Python для генерации ответа, вам необходимо убедиться, что запрос на отмену может быстро прервать выполнение сопрограммы. Это предполагает более частую проверку отмены в сопрограмме и соответствующую обработку CancelledError.

Ключевые изменения:

1. Проверка отмены: добавлен if Task_future.cancelled(): прерывание внутри асинхронных циклов for для частой проверки отмены.
2. Обработка исключений: asyncio.Exceptions.CancelledError изменен на asyncio.CancelledError, чтобы правильно перехватывать исключение отмены.
3. Отмена Threadsafe: гарантировалось, что отмена безопасно вызывается в цикле событий с помощью worker_loop.call_soon_threadsafe(task_future.cancel).

Обновить код:

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

root = tk.Tk()
root.title("Tkinter Button Example")

worker_loop: Optional[asyncio.AbstractEventLoop] = None
task_future: Optional[asyncio.Future] = None

async def get_answer_from_phi3():
    print("Trying")

    messages = [
        {"role": "system", "content": "Hello"}
    ]

    client = ollama.AsyncClient()

    try:
        stream = await client.chat(
            model='phi3',
            messages=messages,
            stream=True,
            options = {"top_k": 1}
        )

        async for chunk in stream:
            if task_future.cancelled():
                break
            print(chunk['message']['content'], end='', flush=True)

    except asyncio.CancelledError:
        print("Cancelled")
    except Exception as e:
        print(e)
        return "Sorry, an error occurred while processing your request."

async def prompt_mistral(query):
    messages = []
    messages.append({"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"})
    
    try:
        client = ollama.AsyncClient()
        stream = await client.chat(
            model='mistral',
            messages=messages,
            stream=True,
            options = {"top_k": 1}
        )

        async for chunk in stream:
            if task_future.cancelled():
                break
            print(chunk['message']['content'])

    except asyncio.CancelledError:
        print("Cancelled response")
        return
    except Exception as e:
        print(e)
        return "Sorry, an error occurred while processing your request."

def prompt_llama(message):
    async def prompt():
        messages = []
        messages.append({"role": "assistant", "content": message})
        try:
            client = ollama.AsyncClient()
            stream = await client.chat(
                model='llama2',
                messages=messages,
                stream=True,
                options = {"top_k": 1}
            )

            generated_answer = ''

            async for chunk in stream:
                if task_future.cancelled():
                    break
                generated_answer += chunk['message']['content']
                print(chunk['message']['content'])

            if "help" in generated_answer:
                await prompt_mistral("Testing")
            else:
                print(generated_answer)

        except asyncio.CancelledError:
            print("Cancelled")
            return
        except Exception as e:
            print(e)
            return "Sorry, an error occurred while processing your request."

    def mistral_worker_function():
        global worker_loop, task_future
        worker_loop = asyncio.new_event_loop()
        task_future = worker_loop.create_task(prompt())
        worker_loop.run_until_complete(task_future)

    print("Starting thread")
    thread = threading.Thread(target=mistral_worker_function)
    thread.start()
    client = ollama.AsyncClient()

def on_button_click():
    global worker_loop, task_future
    if task_future and not task_future.done():
        worker_loop.call_soon_threadsafe(task_future.cancel)
    
    def phi3_worker_function():
        global worker_loop, task_future
        worker_loop = asyncio.new_event_loop()
        task_future = worker_loop.create_task(get_answer_from_phi3())
        worker_loop.run_until_complete(task_future)

    print("Starting thread")
    thread = threading.Thread(target=phi3_worker_function)
    thread.start()

button = tk.Button(root, text = "Stop", command=on_button_click)
button.pack(pady=20)

prompt_llama("Hi")
root.mainloop()

Обновите Ollama до последней версии с помощью curl https://ollama.ai/install.sh | sh в Linux.

Это автоматически исправит эту проблему. Код работает именно так, как есть. Просто нужно обновить Олламу.