Add per-tool timeout_secs to register_function and register_direct_function

The default function call timeout (10s) causes silent failures for
long-running tools. This adds an optional timeout_secs parameter to
register_function() and register_direct_function() so individual tools
can override the global function_call_timeout_secs. The warning message
now mentions both the per-tool and global timeout options.
This commit is contained in:
Mark Backman
2026-03-04 09:37:56 -05:00
parent 9186f65952
commit df35ceca2c

View File

@@ -120,12 +120,15 @@ class FunctionCallRegistryItem:
function_name: The name of the function (None for catch-all handler).
handler: The handler for processing function call parameters.
cancel_on_interruption: Whether to cancel the call on interruption.
timeout_secs: Optional per-tool timeout in seconds. Overrides the global
``function_call_timeout_secs`` for this specific function.
"""
function_name: Optional[str]
handler: FunctionCallHandler | "DirectFunctionWrapper"
cancel_on_interruption: bool
handler_deprecated: bool
timeout_secs: Optional[float] = None
@dataclass
@@ -540,6 +543,7 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
start_callback=None,
*,
cancel_on_interruption: bool = True,
timeout_secs: Optional[float] = None,
):
"""Register a function handler for LLM function calls.
@@ -556,6 +560,9 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
cancel_on_interruption: Whether to cancel this function call when an
interruption occurs. Defaults to True.
timeout_secs: Optional per-tool timeout in seconds. Overrides the global
``function_call_timeout_secs`` for this specific function. Defaults to
None, which uses the global timeout.
"""
signature = inspect.signature(handler)
handler_deprecated = len(signature.parameters) > 1
@@ -574,6 +581,7 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
handler=handler,
cancel_on_interruption=cancel_on_interruption,
handler_deprecated=handler_deprecated,
timeout_secs=timeout_secs,
)
# Start callbacks are now deprecated.
@@ -592,6 +600,7 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
handler: DirectFunction,
*,
cancel_on_interruption: bool = True,
timeout_secs: Optional[float] = None,
):
"""Register a direct function handler for LLM function calls.
@@ -603,6 +612,9 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
handler: The direct function to register. Must follow DirectFunction protocol.
cancel_on_interruption: Whether to cancel this function call when an
interruption occurs. Defaults to True.
timeout_secs: Optional per-tool timeout in seconds. Overrides the global
``function_call_timeout_secs`` for this specific function. Defaults to
None, which uses the global timeout.
"""
wrapper = DirectFunctionWrapper(handler)
self._functions[wrapper.name] = FunctionCallRegistryItem(
@@ -610,6 +622,7 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
handler=wrapper,
cancel_on_interruption=cancel_on_interruption,
handler_deprecated=False,
timeout_secs=timeout_secs,
)
def unregister_function(self, function_name: Optional[str]):
@@ -837,9 +850,16 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
# Start a timeout task for deferred function calls
async def timeout_handler():
try:
await asyncio.sleep(self._function_call_timeout_secs)
effective_timeout = (
item.timeout_secs
if item.timeout_secs is not None
else self._function_call_timeout_secs
)
await asyncio.sleep(effective_timeout)
logger.warning(
f"{self} Function call [{runner_item.function_name}:{runner_item.tool_call_id}] timed out after {self._function_call_timeout_secs} seconds"
f"{self} Function call [{runner_item.function_name}:{runner_item.tool_call_id}] timed out after {effective_timeout} seconds."
f" You can increase this timeout by passing `timeout_secs` to `register_function()`,"
f" or set a global default via `function_call_timeout_secs` on the LLM constructor."
)
await function_call_result_callback(None)
except asyncio.CancelledError: