agents to use tools api (#673)

# What does this PR do?

PR #639 introduced the notion of Tools API and ability to invoke tools
through API just as any resource. This PR changes the Agents to start
using the Tools API to invoke tools. Major changes include:
1) Ability to specify tool groups with AgentConfig
2) Agent gets the corresponding tool definitions for the specified tools
and pass along to the model
3) Attachements are now named as Documents and their behavior is mostly
unchanged from user perspective
4) You can specify args that can be injected to a tool call through
Agent config. This is especially useful in case of memory tool, where
you want the tool to operate on a specific memory bank.
5) You can also register tool groups with args, which lets the agent
inject these as well into the tool call.
6) All tests have been migrated to use new tools API and fixtures
including client SDK tests
7) Telemetry just works with tools API because of our trace protocol
decorator


## Test Plan
```
pytest -s -v -k fireworks llama_stack/providers/tests/agents/test_agents.py  \
   --safety-shield=meta-llama/Llama-Guard-3-8B \
   --inference-model=meta-llama/Llama-3.1-8B-Instruct

pytest -s -v -k together  llama_stack/providers/tests/tools/test_tools.py \
   --safety-shield=meta-llama/Llama-Guard-3-8B \
   --inference-model=meta-llama/Llama-3.1-8B-Instruct

LLAMA_STACK_CONFIG="/Users/dineshyv/.llama/distributions/llamastack-together/together-run.yaml" pytest -v tests/client-sdk/agents/test_agents.py
```
run.yaml:
https://gist.github.com/dineshyv/0365845ad325e1c2cab755788ccc5994

Notebook:
https://colab.research.google.com/drive/1ck7hXQxRl6UvT-ijNRZ-gMZxH1G3cN2d?usp=sharing

llama_stack/distribution/datatypes.py

@@ -20,7 +20,7 @@ from llama_stack.apis.safety import Safety
 from llama_stack.apis.scoring import Scoring
 from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnInput
 from llama_stack.apis.shields import Shield, ShieldInput
-from llama_stack.apis.tools import Tool, ToolGroup, ToolRuntime
+from llama_stack.apis.tools import Tool, ToolGroup, ToolGroupInput, ToolRuntime
 from llama_stack.providers.datatypes import Api, ProviderSpec
 from llama_stack.providers.utils.kvstore.config import KVStoreConfig
 
@@ -161,6 +161,7 @@ a default SQLite store will be used.""",
     datasets: List[DatasetInput] = Field(default_factory=list)
     scoring_fns: List[ScoringFnInput] = Field(default_factory=list)
     eval_tasks: List[EvalTaskInput] = Field(default_factory=list)
+    tool_groups: List[ToolGroupInput] = Field(default_factory=list)
 
 
 class BuildConfig(BaseModel):

llama_stack/distribution/library_client.py

@@ -267,6 +267,7 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
                 self.config, self.custom_provider_registry
             )
         except ModuleNotFoundError as _e:
+            cprint(_e.msg, "red")
             cprint(
                 "Using llama-stack as a library requires installing dependencies depending on the template (providers) you choose.\n",
                 "yellow",

llama_stack/distribution/resolver.py

@@ -5,9 +5,7 @@
 # the root directory of this source tree.
 import importlib
 import inspect
-
 import logging
-
 from typing import Any, Dict, List, Set
 
 from llama_stack.apis.agents import Agents
@@ -28,7 +26,6 @@ from llama_stack.apis.shields import Shields
 from llama_stack.apis.telemetry import Telemetry
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.distribution.client import get_client_impl
-
 from llama_stack.distribution.datatypes import (
     AutoRoutedProviderSpec,
     Provider,
@@ -38,7 +35,6 @@ from llama_stack.distribution.datatypes import (
 from llama_stack.distribution.distribution import builtin_automatically_routed_apis
 from llama_stack.distribution.store import DistributionRegistry
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
-
 from llama_stack.providers.datatypes import (
     Api,
     DatasetsProtocolPrivate,

llama_stack/distribution/routers/routers.py

@@ -6,7 +6,7 @@
 
 from typing import Any, AsyncGenerator, Dict, List, Optional
 
-from llama_stack.apis.common.content_types import InterleavedContent
+from llama_stack.apis.common.content_types import InterleavedContent, URL
 from llama_stack.apis.datasetio import DatasetIO, PaginatedRowsResult
 from llama_stack.apis.eval import (
     AppEvalTaskConfig,
@@ -38,7 +38,7 @@ from llama_stack.apis.scoring import (
     ScoringFnParams,
 )
 from llama_stack.apis.shields import Shield
-from llama_stack.apis.tools import Tool, ToolGroupDef, ToolRuntime
+from llama_stack.apis.tools import ToolDef, ToolRuntime
 from llama_stack.providers.datatypes import RoutingTable
 
 
@@ -417,7 +417,9 @@ class ToolRuntimeRouter(ToolRuntime):
             args=args,
         )
 
-    async def discover_tools(self, tool_group: ToolGroupDef) -> List[Tool]:
-        return await self.routing_table.get_provider_impl(
-            tool_group.name
-        ).discover_tools(tool_group)
+    async def list_runtime_tools(
+        self, tool_group_id: Optional[str] = None, mcp_endpoint: Optional[URL] = None
+    ) -> List[ToolDef]:
+        return await self.routing_table.get_provider_impl(tool_group_id).list_tools(
+            tool_group_id, mcp_endpoint
+        )

llama_stack/distribution/routers/routing_tables.py

@@ -6,7 +6,7 @@
 
 from typing import Any, Dict, List, Optional
 
-from pydantic import parse_obj_as
+from pydantic import TypeAdapter
 
 from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.type_system import ParamType
@@ -26,20 +26,12 @@ from llama_stack.apis.scoring_functions import (
     ScoringFunctions,
 )
 from llama_stack.apis.shields import Shield, Shields
-from llama_stack.apis.tools import (
-    MCPToolGroupDef,
-    Tool,
-    ToolGroup,
-    ToolGroupDef,
-    ToolGroups,
-    UserDefinedToolGroupDef,
-)
+from llama_stack.apis.tools import Tool, ToolGroup, ToolGroups, ToolHost
 from llama_stack.distribution.datatypes import (
     RoutableObject,
     RoutableObjectWithProvider,
     RoutedProtocol,
 )
-
 from llama_stack.distribution.store import DistributionRegistry
 from llama_stack.providers.datatypes import Api, RoutingTable
 
@@ -361,7 +353,7 @@ class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
             memory_bank_data["embedding_dimension"] = model.metadata[
                 "embedding_dimension"
             ]
-        memory_bank = parse_obj_as(MemoryBank, memory_bank_data)
+        memory_bank = TypeAdapter(MemoryBank).validate_python(memory_bank_data)
         await self.register_object(memory_bank)
         return memory_bank
 
@@ -496,54 +488,45 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
     async def list_tools(self, tool_group_id: Optional[str] = None) -> List[Tool]:
         tools = await self.get_all_with_type("tool")
         if tool_group_id:
-            tools = [tool for tool in tools if tool.tool_group == tool_group_id]
+            tools = [tool for tool in tools if tool.toolgroup_id == tool_group_id]
         return tools
 
     async def list_tool_groups(self) -> List[ToolGroup]:
         return await self.get_all_with_type("tool_group")
 
-    async def get_tool_group(self, tool_group_id: str) -> ToolGroup:
-        return await self.get_object_by_identifier("tool_group", tool_group_id)
+    async def get_tool_group(self, toolgroup_id: str) -> ToolGroup:
+        return await self.get_object_by_identifier("tool_group", toolgroup_id)
 
     async def get_tool(self, tool_name: str) -> Tool:
         return await self.get_object_by_identifier("tool", tool_name)
 
     async def register_tool_group(
         self,
-        tool_group_id: str,
-        tool_group: ToolGroupDef,
-        provider_id: Optional[str] = None,
+        toolgroup_id: str,
+        provider_id: str,
+        mcp_endpoint: Optional[URL] = None,
+        args: Optional[Dict[str, Any]] = None,
     ) -> None:
         tools = []
-        tool_defs = []
-        if provider_id is None:
-            if len(self.impls_by_provider_id.keys()) > 1:
-                raise ValueError(
-                    f"No provider_id specified and multiple providers available. Please specify a provider_id. Available providers: {', '.join(self.impls_by_provider_id.keys())}"
-                )
-            provider_id = list(self.impls_by_provider_id.keys())[0]
-
-        if isinstance(tool_group, MCPToolGroupDef):
-            tool_defs = await self.impls_by_provider_id[provider_id].discover_tools(
-                tool_group
-            )
-
-        elif isinstance(tool_group, UserDefinedToolGroupDef):
-            tool_defs = tool_group.tools
-        else:
-            raise ValueError(f"Unknown tool group: {tool_group}")
+        tool_defs = await self.impls_by_provider_id[provider_id].list_runtime_tools(
+            toolgroup_id, mcp_endpoint
+        )
+        tool_host = (
+            ToolHost.model_context_protocol if mcp_endpoint else ToolHost.distribution
+        )
 
         for tool_def in tool_defs:
             tools.append(
                 Tool(
                     identifier=tool_def.name,
-                    tool_group=tool_group_id,
-                    description=tool_def.description,
-                    parameters=tool_def.parameters,
+                    toolgroup_id=toolgroup_id,
+                    description=tool_def.description or "",
+                    parameters=tool_def.parameters or [],
                     provider_id=provider_id,
                     tool_prompt_format=tool_def.tool_prompt_format,
                     provider_resource_id=tool_def.name,
                     metadata=tool_def.metadata,
+                    tool_host=tool_host,
                 )
             )
         for tool in tools:
@@ -561,9 +544,11 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
 
         await self.dist_registry.register(
             ToolGroup(
-                identifier=tool_group_id,
+                identifier=toolgroup_id,
                 provider_id=provider_id,
-                provider_resource_id=tool_group_id,
+                provider_resource_id=toolgroup_id,
+                mcp_endpoint=mcp_endpoint,
+                args=args,
             )
         )
 

llama_stack/distribution/stack.py

@@ -12,7 +12,6 @@ from typing import Any, Dict, Optional
 
 import pkg_resources
 import yaml
-
 from termcolor import colored
 
 from llama_stack.apis.agents import Agents
@@ -33,14 +32,13 @@ from llama_stack.apis.scoring_functions import ScoringFunctions
 from llama_stack.apis.shields import Shields
 from llama_stack.apis.synthetic_data_generation import SyntheticDataGeneration
 from llama_stack.apis.telemetry import Telemetry
-
+from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.distribution.datatypes import StackRunConfig
 from llama_stack.distribution.distribution import get_provider_registry
 from llama_stack.distribution.resolver import ProviderRegistry, resolve_impls
 from llama_stack.distribution.store.registry import create_dist_registry
 from llama_stack.providers.datatypes import Api
 
-
 log = logging.getLogger(__name__)
 
 LLAMA_STACK_API_VERSION = "alpha"
@@ -65,6 +63,8 @@ class LlamaStack(
     Models,
     Shields,
     Inspect,
+    ToolGroups,
+    ToolRuntime,
 ):
     pass
 
@@ -81,6 +81,7 @@ RESOURCES = [
         "list_scoring_functions",
     ),
     ("eval_tasks", Api.eval_tasks, "register_eval_task", "list_eval_tasks"),
+    ("tool_groups", Api.tool_groups, "register_tool_group", "list_tool_groups"),
 ]
 
 

llama_stack/distribution/store/registry.py

@@ -12,7 +12,6 @@ import pydantic
 
 from llama_stack.distribution.datatypes import KVStoreConfig, RoutableObjectWithProvider
 from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
-
 from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
 from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
 
@@ -36,7 +35,7 @@ class DistributionRegistry(Protocol):
 
 
 REGISTER_PREFIX = "distributions:registry"
-KEY_VERSION = "v3"
+KEY_VERSION = "v4"
 KEY_FORMAT = f"{REGISTER_PREFIX}:{KEY_VERSION}::" + "{type}:{identifier}"