Update send_message to be a single async operation. Add start helper to runtime to manage this (#165)

python/examples/core/inner_outer_direct.py

@@ -47,12 +47,12 @@ async def main() -> None:
     runtime = SingleThreadedAgentRuntime()
     inner = runtime.register_and_get("inner", Inner)
     outer = runtime.register_and_get("outer", lambda: Outer(inner))
+
+    run_context = runtime.start()
+
     response = await runtime.send_message(MessageType(body="Hello", sender="external"), outer)
-
-    while not response.done():
-        await runtime.process_next()
-
-    print(await response)
+    print(response)
+    await run_context.stop()
 
 
 if __name__ == "__main__":

python/examples/core/one_agent_direct.py

@@ -50,19 +50,19 @@ async def main() -> None:
         lambda: ChatCompletionAgent("Chat agent", get_chat_completion_client_from_envs(model="gpt-3.5-turbo")),
     )
 
+    run_context = runtime.start()
+
     # Send a message to the agent.
     message = Message(content="Can you tell me something fun about SF?")
     result = await runtime.send_message(message, agent)
 
-    # Process messages until the agent responds.
-    while result.done() is False:
-        await runtime.process_next()
-
     # Get the response from the agent.
     response = await result
     assert isinstance(response, Message)
     print(response.content)
 
+    await run_context.stop()
+
 
 if __name__ == "__main__":
     import logging

python/examples/core/two_agents_pub_sub.py

@@ -100,12 +100,14 @@ async def main() -> None:
         ),
     )
 
+    run_context = runtime.start()
+
     # Send a message to Jack to start the conversation.
     message = Message(content="Can you tell me something fun about SF?", source="User")
     await runtime.send_message(message, jack)
 
     # Process messages.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/demos/assistant.py

@@ -226,13 +226,12 @@ Type "exit" to exit the chat.
 """
     runtime = SingleThreadedAgentRuntime()
     user = assistant_chat(runtime)
+    _run_context = runtime.start()
     print(usage)
     # Request the user to start the conversation.
     await runtime.send_message(PublishNow(), user)
-    while True:
-        # TODO: have a way to exit the loop.
-        await runtime.process_next()
-        await asyncio.sleep(1)
+
+    # TODO: have a way to exit the loop.
 
 
 if __name__ == "__main__":

python/examples/demos/chat_room.py

@@ -17,7 +17,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from common.memory import BufferedChatMemory
 from common.types import Message, TextMessage
 from common.utils import convert_messages_to_llm_messages, get_chat_completion_client_from_envs
-from utils import TextualChatApp, TextualUserAgent, start_runtime
+from utils import TextualChatApp, TextualUserAgent
 
 
 # Define a custom agent that can handle chat room messages.
@@ -140,7 +140,7 @@ async def main() -> None:
     runtime = SingleThreadedAgentRuntime()
     app = TextualChatApp(runtime, user_name="You")
     chat_room(runtime, app)
-    asyncio.create_task(start_runtime(runtime))
+    _run_context = runtime.start()
     await app.run_async()
 
 

python/examples/demos/illustrator_critics.py

@@ -16,7 +16,7 @@ from common.agents import ChatCompletionAgent, ImageGenerationAgent
 from common.memory import BufferedChatMemory
 from common.patterns._group_chat_manager import GroupChatManager
 from common.utils import get_chat_completion_client_from_envs
-from utils import TextualChatApp, TextualUserAgent, start_runtime
+from utils import TextualChatApp, TextualUserAgent
 
 
 def illustrator_critics(runtime: AgentRuntime, app: TextualChatApp) -> None:
@@ -98,7 +98,7 @@ async def main() -> None:
     runtime = SingleThreadedAgentRuntime()
     app = TextualChatApp(runtime, user_name="You")
     illustrator_critics(runtime, app)
-    asyncio.create_task(start_runtime(runtime))
+    _run_context = runtime.start()
     await app.run_async()
 
 

python/examples/demos/software_consultancy.py

@@ -31,7 +31,7 @@ from common.agents import ChatCompletionAgent
 from common.memory import HeadAndTailChatMemory
 from common.patterns._group_chat_manager import GroupChatManager
 from common.utils import get_chat_completion_client_from_envs
-from utils import TextualChatApp, TextualUserAgent, start_runtime
+from utils import TextualChatApp, TextualUserAgent
 
 
 async def write_file(filename: str, content: str) -> str:
@@ -281,7 +281,7 @@ async def main() -> None:
     app = TextualChatApp(runtime, user_name="You")
     software_consultancy(runtime, app)
     # Start the runtime.
-    asyncio.create_task(start_runtime(runtime))
+    _run_context = runtime.start()
     # Start the app.
     await app.run_async()
 

python/examples/demos/utils.py

@@ -4,7 +4,6 @@ import random
 import sys
 from asyncio import Future
 
-from agnext.application import SingleThreadedAgentRuntime
 from agnext.components import Image, TypeRoutedAgent, message_handler
 from agnext.core import AgentRuntime, CancellationToken
 from textual.app import App, ComposeResult
@@ -189,9 +188,3 @@ class TextualUserAgent(TypeRoutedAgent):  # type: ignore
         self, message: ToolApprovalRequest, cancellation_token: CancellationToken
     ) -> ToolApprovalResponse:
         return await self._app.handle_tool_approval_request(message)
-
-
-async def start_runtime(runtime: SingleThreadedAgentRuntime) -> None:  # type: ignore
-    """Run the runtime in a loop."""
-    while True:
-        await runtime.process_next()

python/examples/patterns/coder_executor.py

@@ -182,12 +182,12 @@ async def main(task: str, temp_dir: str) -> None:
     # Register the agents.
     runtime.register("coder", lambda: Coder(model_client=get_chat_completion_client_from_envs(model="gpt-4-turbo")))
     runtime.register("executor", lambda: Executor(executor=LocalCommandLineCodeExecutor(work_dir=temp_dir)))
+    run_context = runtime.start()
 
     # Publish the task message.
     await runtime.publish_message(TaskMessage(content=task), namespace="default")
 
-    # Run the runtime until no more message.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/patterns/coder_reviewer.py

@@ -265,6 +265,7 @@ async def main() -> None:
             model_client=get_chat_completion_client_from_envs(model="gpt-3.5-turbo"),
         ),
     )
+    run_context = runtime.start()
     await runtime.publish_message(
         message=CodeWritingTask(
             task="Write a function to find the directory with the largest number of files using multi-processing."
@@ -273,7 +274,7 @@ async def main() -> None:
     )
 
     # Keep processing messages until idle.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/patterns/group_chat.py

@@ -148,11 +148,13 @@ async def main() -> None:
         ),
     )
 
+    # Start the runtime.
+    run_context = runtime.start()
+
     # Start the conversation.
     await runtime.publish_message(Message(content="Hello, everyone!", source="Moderator"), namespace="default")
 
-    # Run the runtime.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/patterns/mixture_of_agents.py

@@ -149,10 +149,11 @@ async def main() -> None:
             num_references=3,
         ),
     )
+    run_context = runtime.start()
     await runtime.publish_message(AggregatorTask(task="What are something fun to do in SF?"), namespace="default")
 
     # Keep processing messages.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/patterns/multi_agent_debate.py

@@ -246,11 +246,12 @@ async def main(question: str) -> None:
     # Register the aggregator agent.
     runtime.register("MathAggregator", lambda: MathAggregator(num_solvers=4))
 
+    run_context = runtime.start()
+
     # Send a math problem to the aggregator agent.
     await runtime.publish_message(Question(content=question), namespace="default")
 
-    # Run the runtime.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/patterns/orchestrator.py

@@ -130,10 +130,10 @@ def software_development(runtime: AgentRuntime) -> OrchestratorChat:  # type: ig
 async def run(message: str, user: str, scenario: Callable[[AgentRuntime], OrchestratorChat]) -> None:  # type: ignore
     runtime = SingleThreadedAgentRuntime()
     chat = scenario(runtime)
+    run_context = runtime.start()
     response = await runtime.send_message(TextMessage(content=message, source=user), chat.id)
-    while not response.done():
-        await runtime.process_next()
     print((await response).content)  # type: ignore
+    await run_context.stop()
 
 
 if __name__ == "__main__":

python/examples/tool-use/coding_one_agent_direct.py

@@ -140,14 +140,15 @@ async def main() -> None:
         ),
     )
 
+    run_context = runtime.start()
+
     # Send a task to the tool user.
     result = await runtime.send_message(
         UserRequest("Run the following Python code: print('Hello, World!')"), tool_agent
     )
 
     # Run the runtime until the task is completed.
-    while not result.done():
-        await runtime.process_next()
+    await run_context.stop()
 
     # Print the result.
     ai_response = result.result()

python/examples/tool-use/coding_two_agent_pub_sub.py

@@ -202,13 +202,14 @@ async def main() -> None:
         ),
     )
 
+    run_context = runtime.start()
+
     # Publish a task.
     await runtime.publish_message(
         UserRequest("Run the following Python code: print('Hello, World!')"), namespace="default"
     )
 
-    # Run the runtime.
-    await runtime.process_until_idle()
+    await run_context.stop_when_idle()
 
 
 if __name__ == "__main__":

python/examples/tool-use/custom_function_tool_one_agent_direct.py

@@ -49,12 +49,13 @@ async def main() -> None:
         ),
     )
 
+    run_context = runtime.start()
+
     # Send a task to the tool user.
     result = await runtime.send_message(UserRequest("What is the stock price of NVDA on 2024/06/01"), tool_agent)
 
     # Run the runtime until the task is completed.
-    while not result.done():
-        await runtime.process_next()
+    await run_context.stop()
 
     # Print the result.
     ai_response = result.result()