LLMService: make sure function call timeout handler is started

This commit is contained in:
Aleix Conchillo Flaqué
2026-01-30 11:38:26 -08:00
parent 3f77da627d
commit a98ca9b65b

View File

@@ -601,15 +601,7 @@ class LLMService(AIService):
cancel_on_interruption=item.cancel_on_interruption,
)
# Start a timeout task for deferred function calls
async def timeout_handler():
await asyncio.sleep(self._function_call_timeout_secs)
logger.warning(
f"{self} Function call [{runner_item.function_name}:{runner_item.tool_call_id}] timed out after {self._function_call_timeout_secs} seconds"
)
await function_call_result_callback(None)
timeout_task = self.create_task(timeout_handler())
timeout_task: Optional[asyncio.Task] = None
# Define a callback function that pushes a FunctionCallResultFrame upstream & downstream.
async def function_call_result_callback(
@@ -631,6 +623,24 @@ class LLMService(AIService):
properties=properties,
)
# Start a timeout task for deferred function calls
async def timeout_handler():
try:
await asyncio.sleep(self._function_call_timeout_secs)
logger.warning(
f"{self} Function call [{runner_item.function_name}:{runner_item.tool_call_id}] timed out after {self._function_call_timeout_secs} seconds"
)
await function_call_result_callback(None)
except asyncio.CancelledError:
raise
timeout_task = self.create_task(timeout_handler())
# Yield to the event loop so the timeout task coroutine gets entered
# before it could be cancelled. Without this, cancelling the task before
# it starts would leave the coroutine in a "never awaited" state.
await asyncio.sleep(0)
try:
if isinstance(item.handler, DirectFunctionWrapper):
# Handler is a DirectFunctionWrapper
@@ -667,6 +677,9 @@ class LLMService(AIService):
)
await item.handler(params)
except Exception as e:
# Cancel timeout task if it exists
if timeout_task and not timeout_task.done():
await self.cancel_task(timeout_task)
error_message = f"Error executing function call [{runner_item.function_name}]: {e}"
logger.error(f"{self} {error_message}")
await self.push_error(error_msg=error_message, exception=e, fatal=False)