simplify search tool and enable configuration for search engine

llama_toolchain/agentic_system/api/api.py

@@ -41,11 +41,19 @@ class ToolDefinitionCommon(BaseModel):
     output_shields: Optional[List[ShieldDefinition]] = Field(default_factory=list)
 
 
+class SearchEngineType(Enum):
+    bing = "bing"
+    brave = "brave"
+
+
 @json_schema_type
-class BraveSearchToolDefinition(ToolDefinitionCommon):
+class SearchToolDefinition(ToolDefinitionCommon):
+    # NOTE: brave_search is just a placeholder since model always uses
+    # brave_search as tool call name
     type: Literal[AgenticSystemTool.brave_search.value] = (
         AgenticSystemTool.brave_search.value
     )
+    engine: SearchEngineType = SearchEngineType.brave
     remote_execution: Optional[RestAPIExecutionConfig] = None
 
 
@@ -163,7 +171,7 @@ class MemoryToolDefinition(ToolDefinitionCommon):
 
 AgenticSystemToolDefinition = Annotated[
     Union[
-        BraveSearchToolDefinition,
+        SearchToolDefinition,
         WolframAlphaToolDefinition,
         PhotogenToolDefinition,
         CodeInterpreterToolDefinition,

llama_toolchain/agentic_system/client.py

@@ -16,7 +16,6 @@ from pydantic import BaseModel
 from termcolor import cprint
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_toolchain.common.deployment_types import RestAPIExecutionConfig
 from llama_toolchain.core.datatypes import RemoteProviderConfig
 
 from .api import *  # noqa: F403
@@ -135,21 +134,8 @@ async def run_main(host: str, port: int):
     api = AgenticSystemClient(f"http://{host}:{port}")
 
     tool_definitions = [
-        BraveSearchToolDefinition(
+        SearchToolDefinition(engine=SearchEngineType.bing),
-            remote_execution=RestAPIExecutionConfig(
+        WolframAlphaToolDefinition(),
-                url=URL(uri="https://api.bing.microsoft.com/v7.0/search"),
-                method=RestAPIMethod.GET,
-                headers={
-                    "Ocp-Apim-Subscription-Key": "2259b3f9e0cb4fc9b968bb3b02ab13e7"
-                },
-                params={
-                    "count": 3,
-                    "textDecorations": True,
-                    "textFormat": "HTML",
-                },
-            )
-        ),
-        # WolframAlphaToolDefinition(),
         CodeInterpreterToolDefinition(),
     ]
     tool_definitions += [

llama_toolchain/agentic_system/meta_reference/agent_instance.py

@@ -710,7 +710,7 @@ class ChatAgent(ShieldRunnerMixin):
     def _get_tools(self) -> List[ToolDefinition]:
         ret = []
         for t in self.agent_config.tools:
-            if isinstance(t, BraveSearchToolDefinition):
+            if isinstance(t, SearchToolDefinition):
                 ret.append(ToolDefinition(tool_name=BuiltinTool.brave_search))
             elif isinstance(t, WolframAlphaToolDefinition):
                 ret.append(ToolDefinition(tool_name=BuiltinTool.wolfram_alpha))

llama_toolchain/agentic_system/meta_reference/agentic_system.py

@@ -15,10 +15,9 @@ from llama_toolchain.memory.api import Memory
 from llama_toolchain.safety.api import Safety
 from llama_toolchain.agentic_system.api import *  # noqa: F403
 from llama_toolchain.tools.builtin import (
-    BraveSearchTool,
     CodeInterpreterTool,
     PhotogenTool,
-    RemoteSearchTool,
+    SearchTool,
     WolframAlphaTool,
 )
 from llama_toolchain.tools.safety import with_safety
@@ -63,20 +62,19 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem):
                 if not key:
                     raise ValueError("Wolfram API key not defined in config")
                 tool = WolframAlphaTool(key)
-            elif isinstance(tool_defn, BraveSearchToolDefinition):
+            elif isinstance(tool_defn, SearchToolDefinition):
-                if tool_defn.remote_execution is not None:
+                key = None
-                    tool = RemoteSearchTool(tool_defn.remote_execution)
+                if tool_defn.engine == SearchEngineType.brave:
-                else:
                     key = self.config.brave_search_api_key
-                    if not key:
+                elif tool_defn.engine == SearchEngineType.bing:
-                        raise ValueError("Brave API key not defined in config")
+                    key = self.config.bing_search_api_key
-                    tool = BraveSearchTool(key)
+                if not key:
+                    raise ValueError("API key not defined in config")
+                tool = SearchTool(tool_defn.engine, key)
             elif isinstance(tool_defn, CodeInterpreterToolDefinition):
                 tool = CodeInterpreterTool()
             elif isinstance(tool_defn, PhotogenToolDefinition):
-                tool = PhotogenTool(
+                tool = PhotogenTool(dump_dir=tempfile.mkdtemp())
-                    dump_dir=tempfile.mkdtemp(),
-                )
             else:
                 continue
 

llama_toolchain/agentic_system/meta_reference/config.py

@@ -11,4 +11,5 @@ from pydantic import BaseModel
 
 class MetaReferenceImplConfig(BaseModel):
     brave_search_api_key: Optional[str] = None
+    bing_search_api_key: Optional[str] = None
     wolfram_api_key: Optional[str] = None

llama_toolchain/tools/builtin.py

@@ -14,8 +14,6 @@ from typing import List, Optional
 import requests
 from termcolor import cprint
 
-from llama_toolchain.common.deployment_types import RestAPIExecutionConfig
-
 from .ipython_tool.code_execution import (
     CodeExecutionContext,
     CodeExecutionRequest,
@@ -85,21 +83,41 @@ class PhotogenTool(SingleMessageBuiltinTool):
         raise NotImplementedError()
 
 
-class RemoteSearchTool(SingleMessageBuiltinTool):
+class SearchTool(SingleMessageBuiltinTool):
-    def __init__(self, config: RestAPIExecutionConfig) -> None:
+    def __init__(self, engine: SearchEngineType, api_key: str, **kwargs) -> None:
-        self.config = config
+        self.api_key = api_key
+        if engine == SearchEngineType.bing:
+            self.engine = BingSearch(api_key, **kwargs)
+        elif engine == SearchEngineType.brave:
+            self.engine = BraveSearch(api_key, **kwargs)
+        else:
+            raise ValueError(f"Unknown search engine: {engine}")
 
     def get_name(self) -> str:
         return BuiltinTool.brave_search.value
 
     async def run_impl(self, query: str) -> str:
-        params = self.config.params.copy()
+        return await self.engine.search(query)
-        params["q"] = query
+
-        response = requests.get(
+
-            url=self.config.url,
+class BingSearch:
-            params=params,
+    def __init__(self, api_key: str, top_k: int = 3, **kwargs) -> None:
-            headers=self.config.headers,
+        self.api_key = api_key
-        )
+        self.top_k = top_k
+
+    async def search(self, query: str) -> str:
+        url = "https://api.bing.microsoft.com/v7.0/search"
+        headers = {
+            "Ocp-Apim-Subscription-Key": self.api_key,
+        }
+        params = {
+            "count": self.top_k,
+            "textDecorations": True,
+            "textFormat": "HTML",
+            "q": query,
+        }
+
+        response = requests.get(url=url, params=params, headers=headers)
         response.raise_for_status()
         clean = self._clean_response(response.json())
         return json.dumps(clean)
@@ -126,14 +144,11 @@ class RemoteSearchTool(SingleMessageBuiltinTool):
         return {"query": query, "top_k": clean_response}
 
 
-class BraveSearchTool(SingleMessageBuiltinTool):
+class BraveSearch:
     def __init__(self, api_key: str) -> None:
         self.api_key = api_key
 
-    def get_name(self) -> str:
+    async def search(self, query: str) -> str:
-        return BuiltinTool.brave_search.value
-
-    async def run_impl(self, query: str) -> str:
         url = "https://api.search.brave.com/res/v1/web/search"
         headers = {
             "X-Subscription-Token": self.api_key,