From 4c7166f10769c447c4229dd75249c6f2d7b3f8e2 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 15 Dec 2025 12:37:37 -0800 Subject: [PATCH 1/5] MAINT tavily updates for search engine redaction, num_results, and answer_type Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/tavily_search_tool.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/nemo_skills/mcp/servers/tavily_search_tool.py b/nemo_skills/mcp/servers/tavily_search_tool.py index ef4ca54ff5..4a33c3a2b3 100644 --- a/nemo_skills/mcp/servers/tavily_search_tool.py +++ b/nemo_skills/mcp/servers/tavily_search_tool.py @@ -41,15 +41,21 @@ class ExecutionResult: EXCLUDE_DOMAINS: list[str] | None = None +SUPPORTED_SEARCH_ANSWER_TYPES = ["results", "answer"] + ## 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.", choices=SUPPORTED_SEARCH_ANSWER_TYPES) + ] = "answer", ): - """Get a summary of search results from the web using Tavily.""" + """Search the web for a query""" api_url = "https://api.tavily.com/search" @@ -63,6 +69,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 +79,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 @@ -99,7 +106,7 @@ def __init__(self) -> None: "args": ["-m", "nemo_skills.mcp.servers.tavily_search_tool"], }, "hide_args": { - "tavily-search": ["exclude_domains"], + "tavily-search": ["exclude_domains", "num_results", "answer_type"], }, "exclude_domains_config": None, } @@ -120,6 +127,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 From 16695d2ba0f19bcd605fe22727f803bc3e06b8a2 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 15 Dec 2025 13:00:29 -0800 Subject: [PATCH 2/5] FIX tool name in hide_args Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/tavily_search_tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_skills/mcp/servers/tavily_search_tool.py b/nemo_skills/mcp/servers/tavily_search_tool.py index 4a33c3a2b3..aea390539a 100644 --- a/nemo_skills/mcp/servers/tavily_search_tool.py +++ b/nemo_skills/mcp/servers/tavily_search_tool.py @@ -106,7 +106,7 @@ def __init__(self) -> None: "args": ["-m", "nemo_skills.mcp.servers.tavily_search_tool"], }, "hide_args": { - "tavily-search": ["exclude_domains", "num_results", "answer_type"], + "web-search": ["exclude_domains", "num_results", "answer_type"], }, "exclude_domains_config": None, } From b19aa28335cacf372068282c1155f3dc92ebda34 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 15 Dec 2025 13:14:26 -0800 Subject: [PATCH 3/5] FIX enum type Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/tavily_search_tool.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/nemo_skills/mcp/servers/tavily_search_tool.py b/nemo_skills/mcp/servers/tavily_search_tool.py index aea390539a..5dea541088 100644 --- a/nemo_skills/mcp/servers/tavily_search_tool.py +++ b/nemo_skills/mcp/servers/tavily_search_tool.py @@ -17,6 +17,7 @@ import logging import os from dataclasses import dataclass +from enum import Enum from typing import Annotated, Any import httpx @@ -41,7 +42,10 @@ class ExecutionResult: EXCLUDE_DOMAINS: list[str] | None = None -SUPPORTED_SEARCH_ANSWER_TYPES = ["results", "answer"] + +class AnswerType(str, Enum): + results = "results" + answer = "answer" ## See docs https://docs.tavily.com/documentation/api-reference/endpoint/search @@ -51,9 +55,7 @@ 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.", choices=SUPPORTED_SEARCH_ANSWER_TYPES) - ] = "answer", + answer_type: Annotated[AnswerType, Field(description="Type of results to return.")] = AnswerType.answer, ): """Search the web for a query""" From 7a8f7f4f363ba65c8edcfff336312db7ff307e92 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 15 Dec 2025 13:19:08 -0800 Subject: [PATCH 4/5] FIX answer type annotation Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/tavily_search_tool.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/nemo_skills/mcp/servers/tavily_search_tool.py b/nemo_skills/mcp/servers/tavily_search_tool.py index 5dea541088..40bac9abf2 100644 --- a/nemo_skills/mcp/servers/tavily_search_tool.py +++ b/nemo_skills/mcp/servers/tavily_search_tool.py @@ -17,7 +17,6 @@ import logging import os from dataclasses import dataclass -from enum import Enum from typing import Annotated, Any import httpx @@ -43,11 +42,6 @@ class ExecutionResult: EXCLUDE_DOMAINS: list[str] | None = None -class AnswerType(str, Enum): - results = "results" - answer = "answer" - - ## 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="web-search") @@ -55,11 +49,17 @@ 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[AnswerType, Field(description="Type of results to return.")] = AnswerType.answer, + 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", ): """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'." headers = { "Authorization": f"Bearer {TAVILY_API_KEY}", From 5c72de6843f4391ca6afd44d2aaa6497877cfb5d Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 15 Dec 2025 13:36:14 -0800 Subject: [PATCH 5/5] add limit on num results Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/tavily_search_tool.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nemo_skills/mcp/servers/tavily_search_tool.py b/nemo_skills/mcp/servers/tavily_search_tool.py index 40bac9abf2..f72f6761dd 100644 --- a/nemo_skills/mcp/servers/tavily_search_tool.py +++ b/nemo_skills/mcp/servers/tavily_search_tool.py @@ -40,6 +40,7 @@ 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 @@ -60,6 +61,7 @@ async def answer( 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}",