-
Notifications
You must be signed in to change notification settings - Fork 164
Search Tool Parameter updates #1112
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
4c7166f
16695d2
b19aa28
7a8f7f4
5c72de6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -40,18 +40,28 @@ class ExecutionResult: | |
| TAVILY_API_KEY: str | None = None | ||
|
|
||
| EXCLUDE_DOMAINS: list[str] | None = None | ||
| MAX_NUM_RESULTS: int = 20 | ||
|
|
||
|
|
||
| ## See docs https://docs.tavily.com/documentation/api-reference/endpoint/search | ||
| ## There is also a hosted MCP that can be used instead of this tool: https://github.com/tavily-ai/tavily-mcp?tab=readme-ov-file#remote-mcp-server | ||
| @mcp.tool(name="tavily-search") | ||
| @mcp.tool(name="web-search") | ||
| async def answer( | ||
| query: Annotated[str, Field(description="Search query.")], | ||
| exclude_domains: Annotated[list[str], Field(description="Domains to exclude from the search.")] = [], | ||
| num_results: Annotated[int, Field(description="Number of results to return.")] = 10, | ||
| answer_type: Annotated[ | ||
| str, | ||
| Field( | ||
| description='Type of results to return. Choose "answer" for a concise answer or "results" for a list of results.' | ||
| ), | ||
| ] = "answer", | ||
|
Comment on lines
+52
to
+58
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🧩 Analysis chain🏁 Script executed: cat -n nemo_skills/mcp/servers/tavily_search_tool.pyRepository: NVIDIA-NeMo/Skills Length of output: 6495 🌐 Web query:
💡 Result: The Search endpoint uses the parameter named max_results (default 5). Constraint: integer in the range 0 ≤ x ≤ 20. Also note related limits: chunks_per_source must be 1–3. Source: Tavily API docs. [1][2] Sources
Fix critical API payload parameter name and add type safety for The code sends Update the imports to include -from typing import Annotated, Any
+from typing import Annotated, Any, LiteralFix the payload parameter name: - "num_results": num_results,
+ "max_results": num_results,Use answer_type: Annotated[
- str,
+ Literal["answer", "results"],
Field(
description='Type of results to return. Choose "answer" for a concise answer or "results" for a list of results.'
),
] = "answer",Add validation for """Search the web for a query"""
+ if not (0 <= num_results <= 20):
+ return {"error": "num_results must be between 0 and 20"}
+
api_url = "https://api.tavily.com/search"
🤖 Prompt for AI Agents |
||
| ): | ||
| """Get a summary of search results from the web using Tavily.""" | ||
| """Search the web for a query""" | ||
|
|
||
| api_url = "https://api.tavily.com/search" | ||
| assert answer_type in ["answer", "results"], "Invalid answer type. Choose 'answer' or 'results'." | ||
| assert num_results <= MAX_NUM_RESULTS, f"Number of results must be less than or equal to {MAX_NUM_RESULTS}." | ||
|
|
||
| headers = { | ||
| "Authorization": f"Bearer {TAVILY_API_KEY}", | ||
|
|
@@ -63,6 +73,7 @@ async def answer( | |
| # "auto_parameters": False, | ||
| "search_depth": "basic", | ||
| "include_answer": "basic", ## or advanced. | ||
| "num_results": num_results, | ||
| # this should be statically set to the domains we want to exclude | ||
| "exclude_domains": exclude_domains, | ||
| } | ||
|
|
@@ -72,7 +83,7 @@ async def answer( | |
| if response.status_code != 200: | ||
| return {"error": response.json()["error"]} | ||
|
|
||
| result = response.json()["answer"] | ||
| result = response.json()[answer_type] | ||
gwarmstrong marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| return result | ||
|
|
||
|
|
@@ -99,7 +110,7 @@ def __init__(self) -> None: | |
| "args": ["-m", "nemo_skills.mcp.servers.tavily_search_tool"], | ||
| }, | ||
| "hide_args": { | ||
| "tavily-search": ["exclude_domains"], | ||
| "web-search": ["exclude_domains", "num_results", "answer_type"], | ||
| }, | ||
| "exclude_domains_config": None, | ||
| } | ||
|
|
@@ -120,6 +131,9 @@ async def execute(self, tool_name: str, arguments: dict[str, Any], extra_args: d | |
| if not hasattr(self, "exclude_domains"): | ||
| raise ValueError("exclude_domains_config is not set") | ||
| merged_extra["exclude_domains"] = self.exclude_domains | ||
| for key in ["num_results", "answer_type"]: | ||
| if key in self._config: | ||
| merged_extra[key] = self._config[key] | ||
| result = await self._client.call_tool(tool=tool_name, args=arguments, extra_args=merged_extra) | ||
| return result | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Replace mutable default argument with
None.Using a mutable list as a default argument can lead to unexpected behavior if the list is modified. Replace with
Noneand initialize within the function.Apply this diff:
Then update the function body to handle
None:🧰 Tools
🪛 Ruff (0.14.8)
50-50: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
🤖 Prompt for AI Agents