From 9403a44f34fea7e9be3cddd8f8a7648fe835f282 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Tue, 6 Jan 2026 17:18:10 -0800 Subject: [PATCH] chore: Cleanup the workaround logic for streamlit Co-authored-by: Xiang (Sean) Zhou PiperOrigin-RevId: 852995899 --- .../adk/flows/llm_flows/base_llm_flow.py | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 824cd26b..3f8dd37a 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -63,7 +63,6 @@ logger = logging.getLogger('google_adk.' + __name__) _ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name' # Timing configuration -DEFAULT_REQUEST_QUEUE_TIMEOUT = 0.25 DEFAULT_TRANSFER_AGENT_DELAY = 1.0 DEFAULT_TASK_COMPLETION_DELAY = 1.0 @@ -238,29 +237,22 @@ class BaseLlmFlow(ABC): """Sends data to model.""" while True: live_request_queue = invocation_context.live_request_queue - try: - # Streamlit's execution model doesn't preemptively yield to the event - # loop. Therefore, we must explicitly introduce timeouts to allow the - # event loop to process events. - # TODO: revert back(remove timeout) once we move off streamlit. - live_request = await asyncio.wait_for( - live_request_queue.get(), timeout=DEFAULT_REQUEST_QUEUE_TIMEOUT - ) - # duplicate the live_request to all the active streams - logger.debug( - 'Sending live request %s to active streams: %s', - live_request, - invocation_context.active_streaming_tools, - ) - if invocation_context.active_streaming_tools: - for active_streaming_tool in ( - invocation_context.active_streaming_tools - ).values(): - if active_streaming_tool.stream: - active_streaming_tool.stream.send(live_request) - await asyncio.sleep(0) - except asyncio.TimeoutError: - continue + live_request = await live_request_queue.get() + # duplicate the live_request to all the active streams + logger.debug( + 'Sending live request %s to active streams: %s', + live_request, + invocation_context.active_streaming_tools, + ) + if invocation_context.active_streaming_tools: + for active_streaming_tool in ( + invocation_context.active_streaming_tools + ).values(): + if active_streaming_tool.stream: + active_streaming_tool.stream.send(live_request) + # Yield to event loop for cooperative multitasking + await asyncio.sleep(0) + if live_request.close: await llm_connection.close() return