API Keys passed from Client instead of distro configuration

This commit is contained in:
Hardik Shah 2024-09-10 12:36:30 -07:00
parent a11d92601b
commit 0df4d9c9bd
5 changed files with 16 additions and 21 deletions

View file

@ -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

View file

@ -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 += [

View file

@ -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):

View file

@ -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): ...

View file

@ -3,6 +3,7 @@ fire
httpx
huggingface-hub
llama-models>=0.0.13
python-dotenv
pydantic
requests
termcolor