diff --git a/llama_stack/providers/remote/tool_runtime/tavily_search/config.py b/llama_stack/providers/remote/tool_runtime/tavily_search/config.py index ca4e615db..28f7c3d3c 100644 --- a/llama_stack/providers/remote/tool_runtime/tavily_search/config.py +++ b/llama_stack/providers/remote/tool_runtime/tavily_search/config.py @@ -18,10 +18,20 @@ class TavilySearchToolConfig(BaseModel): default=3, description="The maximum number of results to return", ) + timeout: float = Field( + default=30.0, + description="HTTP request timeout in seconds", + ) + connect_timeout: float = Field( + default=10.0, + description="HTTP connection timeout in seconds", + ) @classmethod def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]: return { - "api_key": "${env.TAVILY_SEARCH_API_KEY:=}", + "api_key": "${env.TAVILY_SEARCH_API_KEY:+}", "max_results": 3, + "timeout": 30.0, + "connect_timeout": 10.0, } diff --git a/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py b/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py index 1fe91fd7f..e8ca36cf4 100644 --- a/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py +++ b/llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py @@ -69,7 +69,11 @@ class TavilySearchToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsR async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult: api_key = self._get_api_key() - async with httpx.AsyncClient() as client: + + # Configure timeout for external API calls using config values + timeout = httpx.Timeout(timeout=self.config.timeout, connect=self.config.connect_timeout) + + async with httpx.AsyncClient(timeout=timeout) as client: response = await client.post( "https://api.tavily.com/search", json={"api_key": api_key, "query": kwargs["query"]},