From bdede6d14e7b9fe652e8035197072f1f31d0de10 Mon Sep 17 00:00:00 2001 From: Hardik Shah Date: Mon, 9 Sep 2024 18:41:11 -0700 Subject: [PATCH] simplify search tool and enable configuration for search engine --- llama_toolchain/agentic_system/api/api.py | 12 ++++- llama_toolchain/agentic_system/client.py | 18 +------ .../meta_reference/agent_instance.py | 2 +- .../meta_reference/agentic_system.py | 22 ++++----- .../agentic_system/meta_reference/config.py | 1 + llama_toolchain/tools/builtin.py | 49 ++++++++++++------- 6 files changed, 56 insertions(+), 48 deletions(-) diff --git a/llama_toolchain/agentic_system/api/api.py b/llama_toolchain/agentic_system/api/api.py index 68ec980e6..b8be54861 100644 --- a/llama_toolchain/agentic_system/api/api.py +++ b/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, diff --git a/llama_toolchain/agentic_system/client.py b/llama_toolchain/agentic_system/client.py index 0120de44e..b47e402f0 100644 --- a/llama_toolchain/agentic_system/client.py +++ b/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( - remote_execution=RestAPIExecutionConfig( - 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(), + SearchToolDefinition(engine=SearchEngineType.bing), + WolframAlphaToolDefinition(), CodeInterpreterToolDefinition(), ] tool_definitions += [ diff --git a/llama_toolchain/agentic_system/meta_reference/agent_instance.py b/llama_toolchain/agentic_system/meta_reference/agent_instance.py index 4d38e0032..36c3d19e8 100644 --- a/llama_toolchain/agentic_system/meta_reference/agent_instance.py +++ b/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)) diff --git a/llama_toolchain/agentic_system/meta_reference/agentic_system.py b/llama_toolchain/agentic_system/meta_reference/agentic_system.py index 79d20fd26..9caa3a75b 100644 --- a/llama_toolchain/agentic_system/meta_reference/agentic_system.py +++ b/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): - if tool_defn.remote_execution is not None: - tool = RemoteSearchTool(tool_defn.remote_execution) - else: + elif isinstance(tool_defn, SearchToolDefinition): + key = None + if tool_defn.engine == SearchEngineType.brave: key = self.config.brave_search_api_key - if not key: - raise ValueError("Brave API key not defined in config") - tool = BraveSearchTool(key) + elif tool_defn.engine == SearchEngineType.bing: + key = self.config.bing_search_api_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( - dump_dir=tempfile.mkdtemp(), - ) + tool = PhotogenTool(dump_dir=tempfile.mkdtemp()) else: continue diff --git a/llama_toolchain/agentic_system/meta_reference/config.py b/llama_toolchain/agentic_system/meta_reference/config.py index 367ab17a5..f1a92f2e7 100644 --- a/llama_toolchain/agentic_system/meta_reference/config.py +++ b/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 diff --git a/llama_toolchain/tools/builtin.py b/llama_toolchain/tools/builtin.py index 1792fde57..56fda3723 100644 --- a/llama_toolchain/tools/builtin.py +++ b/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): - def __init__(self, config: RestAPIExecutionConfig) -> None: - self.config = config +class SearchTool(SingleMessageBuiltinTool): + def __init__(self, engine: SearchEngineType, api_key: str, **kwargs) -> None: + 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() - params["q"] = query - response = requests.get( - url=self.config.url, - params=params, - headers=self.config.headers, - ) + return await self.engine.search(query) + + +class BingSearch: + def __init__(self, api_key: str, top_k: int = 3, **kwargs) -> None: + 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: - return BuiltinTool.brave_search.value - - async def run_impl(self, query: str) -> str: + async def search(self, query: str) -> str: url = "https://api.search.brave.com/res/v1/web/search" headers = { "X-Subscription-Token": self.api_key,