feat: add proper BuiltinTool support for both websearch providers

- Add web_search to BuiltinTool enum for consistency with brave_search
- Add WEB_SEARCH_BRAVE_TOOL constant following existing patterns
- Update agent logic to properly map both websearch tools:
  * web_search → BuiltinTool.web_search → WEB_SEARCH_TOOL → Tavily
  * web_search_brave → BuiltinTool.brave_search → WEB_SEARCH_BRAVE_TOOL → Brave
- Add tool encoding support for BuiltinTool.web_search in Llama model
- Replace hardcoded strings with constants for maintainability
- Fix tool conflicts that caused non-deterministic provider selection

This resolves the architectural issues where both websearch providers
exposed the same tool name, causing unpredictable behavior and making
it impossible for users to reliably choose their preferred provider.

Fixes #2606
This commit is contained in:
skamenan7 2025-07-07 15:41:41 -04:00
parent a0a8b86fc8
commit 533c1323b8
3 changed files with 14 additions and 4 deletions

View file

@ -24,6 +24,7 @@ class Role(Enum):
class BuiltinTool(Enum):
web_search = "web_search"
brave_search = "brave_search"
wolfram_alpha = "wolfram_alpha"
photogen = "photogen"

View file

@ -220,7 +220,10 @@ class ToolUtils:
@staticmethod
def encode_tool_call(t: ToolCall, tool_prompt_format: ToolPromptFormat) -> str:
if t.tool_name == BuiltinTool.brave_search:
if t.tool_name == BuiltinTool.web_search:
q = t.arguments["query"]
return f'web_search.call(query="{q}")'
elif t.tool_name == BuiltinTool.brave_search:
q = t.arguments["query"]
return f'brave_search.call(query="{q}")'
elif t.tool_name == BuiltinTool.wolfram_alpha:

View file

@ -80,6 +80,7 @@ def make_random_string(length: int = 8):
TOOLS_ATTACHMENT_KEY_REGEX = re.compile(r"__tools_attachment__=(\{.*?\})")
MEMORY_QUERY_TOOL = "knowledge_search"
WEB_SEARCH_TOOL = "web_search"
WEB_SEARCH_BRAVE_TOOL = "web_search_brave"
RAG_TOOL_GROUP = "builtin::rag"
logger = get_logger(name=__name__, category="agents")
@ -817,9 +818,12 @@ class ChatAgent(ShieldRunnerMixin):
for tool_def in tools.data:
if toolgroup_name.startswith("builtin") and toolgroup_name != RAG_TOOL_GROUP:
identifier: str | BuiltinTool | None = tool_def.identifier
if identifier == "web_search":
# Map tool identifiers to BuiltinTool enum values
if identifier == WEB_SEARCH_TOOL:
identifier = BuiltinTool.web_search
elif identifier == WEB_SEARCH_BRAVE_TOOL:
identifier = BuiltinTool.brave_search
else:
elif identifier in [e.value for e in BuiltinTool]:
identifier = BuiltinTool(identifier)
else:
# add if tool_name is unspecified or the tool_def identifier is the same as the tool_name
@ -880,8 +884,10 @@ class ChatAgent(ShieldRunnerMixin):
f"Tool {tool_name} not found in provided tools, registered tools: {', '.join([str(x) for x in registered_tool_names])}"
)
if isinstance(tool_name, BuiltinTool):
if tool_name == BuiltinTool.brave_search:
if tool_name == BuiltinTool.web_search:
tool_name_str = WEB_SEARCH_TOOL
elif tool_name == BuiltinTool.brave_search:
tool_name_str = WEB_SEARCH_BRAVE_TOOL
else:
tool_name_str = tool_name.value
else: