|
63 | 63 | _ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name' |
64 | 64 |
|
65 | 65 | # Timing configuration |
66 | | -DEFAULT_REQUEST_QUEUE_TIMEOUT = 0.25 |
67 | 66 | DEFAULT_TRANSFER_AGENT_DELAY = 1.0 |
68 | 67 | DEFAULT_TASK_COMPLETION_DELAY = 1.0 |
69 | 68 |
|
@@ -238,29 +237,22 @@ async def _send_to_model( |
238 | 237 | """Sends data to model.""" |
239 | 238 | while True: |
240 | 239 | 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 | + |
264 | 256 | if live_request.close: |
265 | 257 | await llm_connection.close() |
266 | 258 | return |
|
0 commit comments