From 533c1323b8de8b007059f4de83fb267109a668ff Mon Sep 17 00:00:00 2001
From: skamenan7 <skamenan@redhat.com>
Date: Mon, 7 Jul 2025 15:41:41 -0400
Subject: [PATCH] feat: add proper BuiltinTool support for both websearch
 providers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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
---
 llama_stack/models/llama/datatypes.py                |  1 +
 llama_stack/models/llama/llama3/tool_utils.py        |  5 ++++-
 .../inline/agents/meta_reference/agent_instance.py   | 12 +++++++++---
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/llama_stack/models/llama/datatypes.py b/llama_stack/models/llama/datatypes.py
index 7f1ebed55..54075d170 100644
--- a/llama_stack/models/llama/datatypes.py
+++ b/llama_stack/models/llama/datatypes.py
@@ -24,6 +24,7 @@ class Role(Enum):
 
 
 class BuiltinTool(Enum):
+    web_search = "web_search"
     brave_search = "brave_search"
     wolfram_alpha = "wolfram_alpha"
     photogen = "photogen"
diff --git a/llama_stack/models/llama/llama3/tool_utils.py b/llama_stack/models/llama/llama3/tool_utils.py
index 574080184..10118aad0 100644
--- a/llama_stack/models/llama/llama3/tool_utils.py
+++ b/llama_stack/models/llama/llama3/tool_utils.py
@@ -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:
diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
index 4d2b9f8bf..df15fc2ab 100644
--- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
+++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
@@ -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: