mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-29 07:14:20 +00:00
tavily changes
This commit is contained in:
parent
51d9fd4808
commit
9b5f90eefa
2 changed files with 9 additions and 1 deletions
|
@ -18,10 +18,15 @@ class TavilySearchToolConfig(BaseModel):
|
|||
default=3,
|
||||
description="The maximum number of results to return",
|
||||
)
|
||||
timeout: float = Field(
|
||||
default=30.0,
|
||||
description="HTTP request timeout for the API",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
|
||||
return {
|
||||
"api_key": "${env.TAVILY_SEARCH_API_KEY:=}",
|
||||
"max_results": 3,
|
||||
"timeout": 30.0,
|
||||
}
|
||||
|
|
|
@ -69,7 +69,10 @@ 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:
|
||||
|
||||
timeout = httpx.Timeout(timeout=self.config.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"]},
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue