Skip to content

Commit 9403a44

Browse files
seanzhougooglecopybara-github
authored andcommitted
chore: Cleanup the workaround logic for streamlit
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 852995899
1 parent 35308d3 commit 9403a44

File tree

1 file changed

+16
-24
lines changed

1 file changed

+16
-24
lines changed

src/google/adk/flows/llm_flows/base_llm_flow.py

Lines changed: 16 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,6 @@
6363
_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'
6464

6565
# Timing configuration
66-
DEFAULT_REQUEST_QUEUE_TIMEOUT = 0.25
6766
DEFAULT_TRANSFER_AGENT_DELAY = 1.0
6867
DEFAULT_TASK_COMPLETION_DELAY = 1.0
6968

@@ -238,29 +237,22 @@ async def _send_to_model(
238237
"""Sends data to model."""
239238
while True:
240239
live_request_queue = invocation_context.live_request_queue
241-
try:
242-
# Streamlit's execution model doesn't preemptively yield to the event
243-
# loop. Therefore, we must explicitly introduce timeouts to allow the
244-
# event loop to process events.
245-
# TODO: revert back(remove timeout) once we move off streamlit.
246-
live_request = await asyncio.wait_for(
247-
live_request_queue.get(), timeout=DEFAULT_REQUEST_QUEUE_TIMEOUT
248-
)
249-
# duplicate the live_request to all the active streams
250-
logger.debug(
251-
'Sending live request %s to active streams: %s',
252-
live_request,
253-
invocation_context.active_streaming_tools,
254-
)
255-
if invocation_context.active_streaming_tools:
256-
for active_streaming_tool in (
257-
invocation_context.active_streaming_tools
258-
).values():
259-
if active_streaming_tool.stream:
260-
active_streaming_tool.stream.send(live_request)
261-
await asyncio.sleep(0)
262-
except asyncio.TimeoutError:
263-
continue
240+
live_request = await live_request_queue.get()
241+
# duplicate the live_request to all the active streams
242+
logger.debug(
243+
'Sending live request %s to active streams: %s',
244+
live_request,
245+
invocation_context.active_streaming_tools,
246+
)
247+
if invocation_context.active_streaming_tools:
248+
for active_streaming_tool in (
249+
invocation_context.active_streaming_tools
250+
).values():
251+
if active_streaming_tool.stream:
252+
active_streaming_tool.stream.send(live_request)
253+
# Yield to event loop for cooperative multitasking
254+
await asyncio.sleep(0)
255+
264256
if live_request.close:
265257
await llm_connection.close()
266258
return

0 commit comments

Comments
 (0)