use customtool's get_tool_definition to remove duplication in agentconfig

This commit is contained in:
Jeff Tang 2024-12-08 10:54:32 -08:00
parent a29013112f
commit ee65d95bb5
2 changed files with 12 additions and 37 deletions

View file

@ -286,6 +286,9 @@
" input_shields = [] if disable_safety else [\"llama_guard\"]\n",
" output_shields = [] if disable_safety else [\"llama_guard\"]\n",
"\n",
" # Initialize custom tool (ensure `WebSearchTool` is defined earlier in the notebook)\n",
" webSearchTool = WebSearchTool(api_key=BRAVE_SEARCH_API_KEY)\n",
" \n",
" # Define the agent configuration, including the model and tool setup\n",
" agent_config = AgentConfig(\n",
" model=MODEL_NAME,\n",
@ -296,18 +299,7 @@
" \"top_p\": 0.9,\n",
" },\n",
" tools=[\n",
" {\n",
" \"function_name\": \"web_search\", # Name of the tool being integrated\n",
" \"description\": \"Search the web for a given query\",\n",
" \"parameters\": {\n",
" \"query\": {\n",
" \"param_type\": \"str\",\n",
" \"description\": \"The query to search for\",\n",
" \"required\": True,\n",
" }\n",
" },\n",
" \"type\": \"function_call\",\n",
" },\n",
" webSearchTool.get_tool_definition()\n",
" ],\n",
" tool_choice=\"auto\",\n",
" tool_prompt_format=\"python_list\",\n",
@ -316,11 +308,8 @@
" enable_session_persistence=False,\n",
" )\n",
"\n",
" # Initialize custom tools (ensure `WebSearchTool` is defined earlier in the notebook)\n",
" custom_tools = [WebSearchTool(api_key=BRAVE_SEARCH_API_KEY)]\n",
"\n",
" # Create an agent instance with the client and configuration\n",
" agent = Agent(client, agent_config, custom_tools)\n",
" agent = Agent(client, agent_config, [webSearchTool])\n",
"\n",
" # Create a session for interaction and print the session ID\n",
" session_id = agent.create_session(\"test-session\")\n",