Handling Cancelled UserProxyAgent input requests

[jspv]

I'm having difficultly figuring out the right way to cleanly catch the cancellation from the token when leveraged in the input_func parameter to UserProxyAgent. I want to catch the expected CancelledError, but I'm seeing a pedantic exception which appears to be coming from _user_proxy_agent.py which I can't seem to catch or handle more gracefully. Below is a minimal example:

import asyncio

from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient


async def response(prompt: str, cancellation_token: CancellationToken) -> str:
    await asyncio.sleep(5)
    cancellation_token.cancel()


# Create the agents.
model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")
assistant = AssistantAgent("assistant", model_client=model_client)
user_proxy = UserProxyAgent(
    "user_proxy", input_func=response
)

# Create the termination condition which will end the conversation when the user says "APPROVE".
termination = TextMentionTermination("APPROVE")

# Create the team.
team = RoundRobinGroupChat([assistant, user_proxy], termination_condition=termination)

# Run the conversation and stream to the console.
cancellation_token = CancellationToken()
stream = team.run_stream(
    task="Greet User and answer any question they ask:",
    cancellation_token=cancellation_token,
)

async def main():
    try:
        await Console(stream)
    except asyncio.exceptions.CancelledError:
        print("Conversation Cancelled.")
    except BaseException as e:
        print(f"A different error occurred:{type(e)}")
        # Clean up the model client.

    print("Done.")


if __name__ == "__main__":
    asyncio.run(main())

---------- user ----------
Greet User and answer any question they ask:
---------- assistant ----------
Hello! How can I assist you today? If you have any questions, feel free to ask!
Error processing publish message for user_proxy_0ef32129-ee05-4536-a001-a82e39d73e36/0ef32129-ee05-4536-a001-a82e39d73e36
Traceback (most recent call last):
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 230, in on_messages_stream
    yield Response(chat_message=TextMessage(content=user_input, source=self.name))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/pydantic/main.py", line 214, in __init__
    validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 1 validation error for TextMessage
content
  Input should be a valid string [type=string_type, input_value=None, input_type=NoneType]
    For further information visit https://errors.pydantic.dev/2.10/v/string_type

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_core/_single_threaded_agent_runtime.py", line 510, in _on_message
    return await agent.on_message(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_core/_base_agent.py", line 113, in on_message
    return await self.on_message_impl(message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_agentchat/teams/_group_chat/_sequential_routed_agent.py", line 67, in on_message_impl
    return await super().on_message_impl(message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_core/_routed_agent.py", line 485, in on_message_impl
    return await h(self, message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_core/_routed_agent.py", line 268, in wrapper
    return_value = await func(self, message, ctx)  # type: ignore
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_agentchat/teams/_group_chat/_chat_agent_container.py", line 69, in handle_request
    async for msg in self._agent.on_messages_stream(self._message_buffer, ctx.cancellation_token):
  File "/Users/jspv/src/mchat/.venv/lib/python3.11/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 235, in on_messages_stream
    raise RuntimeError(f"Failed to get user input: {str(e)}") from e
RuntimeError: Failed to get user input: 1 validation error for TextMessage
content
  Input should be a valid string [type=string_type, input_value=None, input_type=NoneType]
    For further information visit https://errors.pydantic.dev/2.10/v/string_type
A different error occurred:<class 'RuntimeError'>
Done.

I also tried reducing it to just down to UserProxyAgent, no RoundRobin and while here I can catch the exception, it exhibits an odd behavior that I have to press "Enter" to complete the process, even though it is cancelled.

import asyncio
from typing import Any

from autogen_agentchat.agents import UserProxyAgent
from autogen_agentchat.messages import TextMessage
from autogen_core import CancellationToken

token = CancellationToken()
agent = UserProxyAgent("user_proxy")


async def timeout(delay: float):
    await asyncio.sleep(delay)


def cancellation_callback(task: asyncio.Task[Any]):
    token.cancel()


async def cancellable_user_agent():
    response = None
    try:
        timeout_task = asyncio.create_task(timeout(3))
        timeout_task.add_done_callback(cancellation_callback)
        agent_task = asyncio.create_task(
            agent.on_messages(
                [TextMessage(content="What is your name? ", source="user")],
                cancellation_token=token,
            )
        )
        response = await agent_task
        print(f"Your name is {response.chat_message.content}")
        print(f"Agent status is {agent_task._state}")
    except asyncio.exceptions.CancelledError:
        print(f"Cancelled, response = {response}")
        print(f"Agent status is {agent_task._state}")
    except BaseException as e:
        print(f"BaseException: {e} {type(e)}")


if __name__ == "__main__":
    asyncio.run(cancellable_user_agent())

Letting the timer cancel the task:

$python test.py
Enter your response: Cancelled, response = None
Agent status is CANCELLED

$  <-- Enter needed to be pressed to exit

Completing the task via input:

$python test.py
Enter your response: hello
Your name is hello
Agent status is FINISHED
$

No enter needed. Not sure if this second issue is related, it seems to have a process sitting out there waiting. I can hit ctrl-c instead of enter and I see that something is keeping cancellable_user_agent active even after the task has been cancelled.

Any help appreciated.

Enter your response: Cancelled, response = None
Agent status is CANCELLED
^CTraceback (most recent call last):
  File "/Users/jspv/src/mchat/foo2.py", line 43, in <module>
    asyncio.run(cancellable_user_agent())
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/runners.py", line 189, in run
    with Runner(debug=debug) as runner:
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/runners.py", line 63, in __exit__
    self.close()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/runners.py", line 73, in close
    loop.run_until_complete(loop.shutdown_default_executor())
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/base_events.py", line 641, in run_until_complete
    self.run_forever()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/base_events.py", line 608, in run_forever
    self._run_once()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/asyncio/base_events.py", line 1898, in _run_once
    event_list = self._selector.select(timeout)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/selectors.py", line 566, in select
    kev_list = self._selector.control(None, max_ev, timeout)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt
^CException ignored in: <module 'threading' from '/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py'>
Traceback (most recent call last):
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py", line 1560, in _shutdown
    atexit_call()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/concurrent/futures/thread.py", line 31, in _python_exit
    t.join()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py", line 1119, in join
    self._wait_for_tstate_lock()
  File "/opt/homebrew/Cellar/[email protected]/3.11.11/Frameworks/Python.framework/Versions/3.11/lib/python3.11/threading.py", line 1139, in _wait_for_tstate_lock
    if lock.acquire(block, timeout):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt: 

[ekzhu]

@jackgerrits thoughts?

[jackgerrits]

@jspv if you wish to cancel a user input function you should raise a CancelledError from within the callback. The cancellation token is to allow you to respond to being cancelled.