Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions nemo_skills/mcp/servers/tavily_search_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.")] = [],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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 None and initialize within the function.

Apply this diff:

-    exclude_domains: Annotated[list[str], Field(description="Domains to exclude from the search.")] = [],
+    exclude_domains: Annotated[list[str] | None, Field(description="Domains to exclude from the search.")] = None,

Then update the function body to handle None:

+    if exclude_domains is None:
+        exclude_domains = []
+
     api_url = "https://api.tavily.com/search"

Committable suggestion skipped: line range outside the PR's diff.

🧰 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
In nemo_skills/mcp/servers/tavily_search_tool.py around line 50, replace the
mutable default argument "exclude_domains: Annotated[list[str], Field(...)] =
[]" with "exclude_domains: Annotated[list[str] | None, Field(...)] = None" (or
simply set to None while keeping typing), and then in the function body check if
exclude_domains is None and set it to an empty list (e.g., exclude_domains = []
if exclude_domains is None else exclude_domains) so the default list is created
per-call rather than shared across calls.

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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

cat -n nemo_skills/mcp/servers/tavily_search_tool.py

Repository: NVIDIA-NeMo/Skills

Length of output: 6495


🌐 Web query:

Tavily API search endpoint num_results parameter constraints documentation

💡 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

  • Tavily Search API reference (parameters and max_results range). [1]
  • Tavily SDK / docs (parameters; chunks_per_source limits). [2]

Fix critical API payload parameter name and add type safety for answer_type.

The code sends "num_results" to the Tavily API payload (line 74), but the API expects "max_results". Additionally, use Literal type for answer_type to enforce valid options at the type level, and add validation for num_results according to Tavily API constraints (0–20).

Update the imports to include Literal:

-from typing import Annotated, Any
+from typing import Annotated, Any, Literal

Fix the payload parameter name:

-        "num_results": num_results,
+        "max_results": num_results,

Use Literal type for answer_type:

     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 num_results (API constraint: 0–20):

     """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"

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents
nemo_skills/mcp/servers/tavily_search_tool.py lines 51-57: the API parameter
name and types are incorrect — change the payload key sent to Tavily from
"num_results" to "max_results" where the request is built (around line 74),
update the function parameter annotation to use typing.Literal for answer_type
(Literal["answer","results"]) and import Literal at top, and add explicit
validation for num_results to ensure it's an int within Tavily's allowed range
0–20 (raise/convert and clamp or raise ValueError) before adding it to the
payload.

):
"""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}",
Expand All @@ -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,
}
Expand All @@ -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]

return result

Expand All @@ -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,
}
Expand All @@ -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

Expand Down