diff --git a/llama_toolchain/agentic_system/api/api.py b/llama_toolchain/agentic_system/api/api.py index b8be54861..fd4cd9aa0 100644 --- a/llama_toolchain/agentic_system/api/api.py +++ b/llama_toolchain/agentic_system/api/api.py @@ -53,6 +53,7 @@ class SearchToolDefinition(ToolDefinitionCommon): type: Literal[AgenticSystemTool.brave_search.value] = ( AgenticSystemTool.brave_search.value ) + api_key: str engine: SearchEngineType = SearchEngineType.brave remote_execution: Optional[RestAPIExecutionConfig] = None @@ -62,6 +63,7 @@ class WolframAlphaToolDefinition(ToolDefinitionCommon): type: Literal[AgenticSystemTool.wolfram_alpha.value] = ( AgenticSystemTool.wolfram_alpha.value ) + api_key: str remote_execution: Optional[RestAPIExecutionConfig] = None diff --git a/llama_toolchain/agentic_system/client.py b/llama_toolchain/agentic_system/client.py index b47e402f0..c38cd03ef 100644 --- a/llama_toolchain/agentic_system/client.py +++ b/llama_toolchain/agentic_system/client.py @@ -6,11 +6,12 @@ import asyncio import json +import os from typing import AsyncGenerator import fire - import httpx +from dotenv import load_dotenv from pydantic import BaseModel from termcolor import cprint @@ -22,6 +23,9 @@ from .api import * # noqa: F403 from .event_logger import EventLogger +load_dotenv() + + async def get_client_impl(config: RemoteProviderConfig, _deps): return AgenticSystemClient(config.url) @@ -134,8 +138,11 @@ async def run_main(host: str, port: int): api = AgenticSystemClient(f"http://{host}:{port}") tool_definitions = [ - SearchToolDefinition(engine=SearchEngineType.bing), - WolframAlphaToolDefinition(), + SearchToolDefinition( + engine=SearchEngineType.bing, + api_key=os.getenv("BING_SEARCH_API_KEY"), + ), + WolframAlphaToolDefinition(api_key=os.getenv("WOLFRAM_ALPHA_API_KEY")), CodeInterpreterToolDefinition(), ] tool_definitions += [ diff --git a/llama_toolchain/agentic_system/meta_reference/agentic_system.py b/llama_toolchain/agentic_system/meta_reference/agentic_system.py index 9caa3a75b..052d87029 100644 --- a/llama_toolchain/agentic_system/meta_reference/agentic_system.py +++ b/llama_toolchain/agentic_system/meta_reference/agentic_system.py @@ -58,19 +58,9 @@ class MetaReferenceAgenticSystemImpl(AgenticSystem): builtin_tools = [] for tool_defn in agent_config.tools: if isinstance(tool_defn, WolframAlphaToolDefinition): - key = self.config.wolfram_api_key - if not key: - raise ValueError("Wolfram API key not defined in config") - tool = WolframAlphaTool(key) + tool = WolframAlphaTool(tool_defn.api_key) elif isinstance(tool_defn, SearchToolDefinition): - key = None - if tool_defn.engine == SearchEngineType.brave: - key = self.config.brave_search_api_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) + tool = SearchTool(tool_defn.engine, tool_defn.api_key) elif isinstance(tool_defn, CodeInterpreterToolDefinition): tool = CodeInterpreterTool() elif isinstance(tool_defn, PhotogenToolDefinition): diff --git a/llama_toolchain/agentic_system/meta_reference/config.py b/llama_toolchain/agentic_system/meta_reference/config.py index f1a92f2e7..17beb348e 100644 --- a/llama_toolchain/agentic_system/meta_reference/config.py +++ b/llama_toolchain/agentic_system/meta_reference/config.py @@ -4,12 +4,7 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from typing import Optional - 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 +class MetaReferenceImplConfig(BaseModel): ... diff --git a/requirements.txt b/requirements.txt index 720f84b79..1fd61b886 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ fire httpx huggingface-hub llama-models>=0.0.13 +python-dotenv pydantic requests termcolor