simplify toolgroups registration

llama_stack/providers/tests/agents/test_agents.py

@@ -45,8 +45,7 @@ def common_params(inference_model):
         sampling_params=SamplingParams(temperature=0.7, top_p=0.95),
         input_shields=[],
         output_shields=[],
-        available_tools=[],
-        preprocessing_tools=[],
+        toolgroups=[],
         max_infer_iters=5,
     )
 
@@ -83,27 +82,27 @@ def query_attachment_messages():
     ]
 
 
-async def create_agent_turn_with_search_tool(
+async def create_agent_turn_with_toolgroup(
     agents_stack: Dict[str, object],
     search_query_messages: List[object],
     common_params: Dict[str, str],
-    tool_name: str,
+    toolgroup_name: str,
 ) -> None:
     """
-    Create an agent turn with a search tool.
+    Create an agent turn with a toolgroup.
 
     Args:
         agents_stack (Dict[str, object]): The agents stack.
         search_query_messages (List[object]): The search query messages.
         common_params (Dict[str, str]): The common parameters.
-        search_tool_definition (SearchToolDefinition): The search tool definition.
+        toolgroup_name (str): The name of the toolgroup.
     """
 
-    # Create an agent with the search tool
+    # Create an agent with the toolgroup
     agent_config = AgentConfig(
         **{
             **common_params,
-            "tools": [tool_name],
+            "toolgroups": [toolgroup_name],
         }
     )
 
@@ -249,7 +248,7 @@ class TestAgents:
         agent_config = AgentConfig(
             **{
                 **common_params,
-                "tools": ["memory"],
+                "toolgroups": ["builtin::memory"],
                 "tool_choice": ToolChoice.auto,
             }
         )
@@ -289,13 +288,58 @@ class TestAgents:
         if "TAVILY_SEARCH_API_KEY" not in os.environ:
             pytest.skip("TAVILY_SEARCH_API_KEY not set, skipping test")
 
-        await create_agent_turn_with_search_tool(
-            agents_stack,
-            search_query_messages,
-            common_params,
-            "brave_search",
+        # Create an agent with the toolgroup
+        agent_config = AgentConfig(
+            **{
+                **common_params,
+                "toolgroups": ["builtin::web_search"],
+            }
         )
 
+        agent_id, session_id = await create_agent_session(
+            agents_stack.impls[Api.agents], agent_config
+        )
+        turn_request = dict(
+            agent_id=agent_id,
+            session_id=session_id,
+            messages=search_query_messages,
+            stream=True,
+        )
+
+        turn_response = [
+            chunk
+            async for chunk in await agents_stack.impls[Api.agents].create_agent_turn(
+                **turn_request
+            )
+        ]
+
+        assert len(turn_response) > 0
+        assert all(
+            isinstance(chunk, AgentTurnResponseStreamChunk) for chunk in turn_response
+        )
+
+        check_event_types(turn_response)
+
+        # Check for tool execution events
+        tool_execution_events = [
+            chunk
+            for chunk in turn_response
+            if isinstance(chunk.event.payload, AgentTurnResponseStepCompletePayload)
+            and chunk.event.payload.step_details.step_type
+            == StepType.tool_execution.value
+        ]
+        assert len(tool_execution_events) > 0, "No tool execution events found"
+
+        # Check the tool execution details
+        tool_execution = tool_execution_events[0].event.payload.step_details
+        assert isinstance(tool_execution, ToolExecutionStep)
+        assert len(tool_execution.tool_calls) > 0
+        actual_tool_name = tool_execution.tool_calls[0].tool_name
+        assert actual_tool_name == "web_search"
+        assert len(tool_execution.tool_responses) > 0
+
+        check_turn_complete_event(turn_response, session_id, search_query_messages)
+
 
 def check_event_types(turn_response):
     event_types = [chunk.event.payload.event_type for chunk in turn_response]

llama_stack/providers/tests/tools/fixtures.py

@@ -8,16 +8,9 @@ import os
 
 import pytest
 import pytest_asyncio
-from llama_models.llama3.api.datatypes import BuiltinTool
 
 from llama_stack.apis.models import ModelInput, ModelType
-from llama_stack.apis.tools import (
-    BuiltInToolDef,
-    ToolGroupInput,
-    ToolParameter,
-    UserDefinedToolDef,
-    UserDefinedToolGroupDef,
-)
+from llama_stack.apis.tools import ToolGroupInput
 from llama_stack.distribution.datatypes import Api, Provider
 from llama_stack.providers.tests.resolver import construct_stack_for_test
 
@@ -47,30 +40,7 @@ def tool_runtime_memory_and_search() -> ProviderFixture:
 @pytest.fixture(scope="session")
 def tool_group_input_memory() -> ToolGroupInput:
     return ToolGroupInput(
-        tool_group_id="memory_group",
-        tool_group=UserDefinedToolGroupDef(
-            tools=[
-                UserDefinedToolDef(
-                    name="memory",
-                    description="Query the memory bank",
-                    parameters=[
-                        ToolParameter(
-                            name="input_messages",
-                            description="The input messages to search for in memory",
-                            parameter_type="list",
-                            required=True,
-                        ),
-                    ],
-                    metadata={
-                        "config": {
-                            "memory_bank_configs": [
-                                {"bank_id": "test_bank", "type": "vector"}
-                            ]
-                        }
-                    },
-                )
-            ],
-        ),
+        toolgroup_id="builtin::memory",
         provider_id="memory-runtime",
     )
 
@@ -78,10 +48,7 @@ def tool_group_input_memory() -> ToolGroupInput:
 @pytest.fixture(scope="session")
 def tool_group_input_tavily_search() -> ToolGroupInput:
     return ToolGroupInput(
-        tool_group_id="tavily_search_group",
-        tool_group=UserDefinedToolGroupDef(
-            tools=[BuiltInToolDef(built_in_type=BuiltinTool.brave_search, metadata={})],
-        ),
+        toolgroup_id="builtin::web_search",
         provider_id="tavily-search",
     )
 

llama_stack/providers/tests/tools/test_tools.py

@@ -43,8 +43,8 @@ def sample_documents():
 
 class TestTools:
     @pytest.mark.asyncio
-    async def test_brave_search_tool(self, tools_stack, sample_search_query):
-        """Test the Brave search tool functionality."""
+    async def test_web_search_tool(self, tools_stack, sample_search_query):
+        """Test the web search tool functionality."""
         if "TAVILY_SEARCH_API_KEY" not in os.environ:
             pytest.skip("TAVILY_SEARCH_API_KEY not set, skipping test")
 
@@ -52,7 +52,7 @@ class TestTools:
 
         # Execute the tool
         response = await tools_impl.invoke_tool(
-            tool_name="brave_search", args={"query": sample_search_query}
+            tool_name="web_search", args={"query": sample_search_query}
         )
 
         # Verify the response
@@ -89,11 +89,12 @@ class TestTools:
         response = await tools_impl.invoke_tool(
             tool_name="memory",
             args={
-                "input_messages": [
+                "messages": [
                     UserMessage(
                         content="What are the main topics covered in the documentation?",
                     )
                 ],
+                "memory_bank_ids": ["test_bank"],
             },
         )