Python: Request cancellation sample (#459)

* request cancellation via tasks

* fix missing kwargs
This commit is contained in:
peterychang
2025-08-20 15:56:26 -04:00
committed by GitHub
Unverified
parent 011edfe420
commit feb4e908ae
3 changed files with 20 additions and 51 deletions
@@ -1,50 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import threading
from asyncio import Future
from collections.abc import Callable
from typing import Any
# from https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/src/autogen_core/_cancellation_token.py
class CancellationToken:
"""A token used to cancel pending async calls."""
def __init__(self) -> None:
self._cancelled: bool = False
self._lock: threading.Lock = threading.Lock()
self._callbacks: list[Callable[[], None]] = []
def cancel(self) -> None:
"""Cancel pending async calls linked to this cancellation token."""
with self._lock:
if not self._cancelled:
self._cancelled = True
for callback in self._callbacks:
callback()
def is_cancelled(self) -> bool:
"""Check if the CancellationToken has been used."""
with self._lock:
return self._cancelled
def add_callback(self, callback: Callable[[], None]) -> None:
"""Attach a callback that will be called when cancel is invoked."""
with self._lock:
if self._cancelled:
callback()
else:
self._callbacks.append(callback)
def link_future(self, future: Future[Any]) -> Future[Any]:
"""Link a pending async call to a token to allow its cancellation."""
with self._lock:
if self._cancelled:
future.cancel()
else:
def _cancel() -> None:
future.cancel()
self._callbacks.append(_cancel)
return future
@@ -90,7 +90,7 @@ def _tool_call_non_streaming(
response: ChatResponse | None = None
fcc_messages: list[ChatMessage] = []
for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
response = await func(self, messages=messages, chat_options=chat_options)
response = await func(self, messages=messages, chat_options=chat_options, **kwargs)
# if there are function calls, we will handle them first
function_results = {
it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
@@ -0,0 +1,19 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework.openai import OpenAIChatClient
async def main():
chat_client = OpenAIChatClient()
try:
task = asyncio.create_task(chat_client.get_response(messages=["Tell me a fantasy story."]))
await asyncio.sleep(1)
task.cancel()
await task
except asyncio.CancelledError:
print("Request was cancelled")
if __name__ == "__main__":
asyncio.run(main())