adk-tool-search 0.2.1

Dynamic tool search for Google ADK — load tools on demand instead of all at once

adk-tool-search

Dynamic tool search for Google ADK — load tools on demand instead of all at once.

Implements the Anthropic Tool Search pattern for Google's Agent Development Kit (ADK). Instead of loading all tool definitions into context upfront, the agent discovers and loads tools on demand using BM25 search.

Primary integration target: standard ADK LlmAgent wiring with ToolRegistry + callbacks.

Why?

Problem	Impact
Context bloat	A typical multi-MCP setup can consume 50k+ tokens in tool definitions before the agent does any work
Tool selection accuracy	LLM ability to pick the right tool degrades past 30-50 tools
Gemini's 100-tool limit	Hard cap on function declarations in the Gemini API

This library reduces context usage by ~95% and keeps tool selection accurate across hundreds of tools.

How it works

┌─────────────────────────────────────────────────────┐
│  Startup                                            │
│  1. Fetch tools from MCP servers / register funcs   │
│  2. Index all tools in BM25 registry                │
│  3. Agent starts with only: search_tools, load_tool │
└─────────────────────────────────────────────────────┘
         │
         ▼
┌─────────────────────────────────────────────────────┐
│  Runtime (per user request)                         │
│  1. Agent calls search_tools("weather forecast")    │
│  2. Registry returns top-5 matches (name + snippet) │
│  3a. Option A: Load + execute in one turn           │
│     load_tool("get_forecast", args={"loc": "Tokyo"})│
│     → Returns the tool result immediately           │
│  3b. Option B: Load for subsequent turns            │
│     load_tool("get_forecast")                       │
│     → Tool is marked loaded for this session        │
│     → before_model_callback injects it next turn    │
│     → Agent calls get_forecast(location="Tokyo")   │
└─────────────────────────────────────────────────────┘

Inline execution (one-turn)

load_tool accepts an optional args dict. When provided, the tool is loaded and executed immediately within the after_tool_callback, returning the result in the same turn. This eliminates the extra round-trip:

# Three-turn flow (load, then call separately)
load_tool("get_weather")                          # Turn 2: "loaded, call next turn"
get_weather(location="Tokyo")                      # Turn 3: {"temp": 22, ...}

# Two-turn flow (load + execute inline)
load_tool("get_weather", args={"location":"Tokyo"})  # Turn 2: {"result": {"temp": 22, ...}}

Inline execution works for plain Python callables, ADK FunctionTool, and BaseTool subclasses that implement run_async (including MCP tools). If inline execution isn't possible for a tool type, the tool is still loaded for use on the next turn.

Loaded tools are session-scoped. A tool loaded in one session is not exposed to other sessions.

Persistence model:

• Loaded tool names are written to ADK session state (adk_tool_search.loaded_tools) on load_tool.
• before_model_callback reads that state and injects only those tools for the current session.
• With persistent session services (SQLite/DB/Vertex), loaded tools survive process restarts.
• With in-memory session services, restart continuity is not available.

Skills integration:

• If a use_skill tool call returns allowed_tools (or allowed_tools_raw / allowed-tools), adk-tool-search auto-loads matching registry tools into the same session state.
• This keeps skill-required tools deferred by default and activates them only when the skill is activated.
• Supported token forms include plain names (for example, get_weather) and function-like tokens (for example, Bash(git:*), resolved by base name when possible).

Install

pip install adk-tool-search

Development setup

git clone https://github.com/manojlds/adk-tool-search.git
cd adk-tool-search
uv sync --all-extras

Quick start

Recommended: use standard ADK LlmAgent

With plain Python functions

from google.adk.agents import LlmAgent
from adk_tool_search import (
    ToolRegistry,
    create_search_and_load_tools,
    create_session_scoped_loader_callbacks,
)

def get_weather(location: str) -> dict:
    """Get current weather for a location."""
    return {"location": location, "temp": 22, "condition": "sunny"}

def send_email(to: str, subject: str, body: str) -> dict:
    """Send an email."""
    return {"status": "sent"}

# 1. Register tools in the search index
registry = ToolRegistry()
registry.register_many([get_weather, send_email])

# 2. Create the lightweight discovery tools
search_tools, load_tool = create_search_and_load_tools(registry)

# 3. Create session-scoped loader callbacks
before_model_callback, after_tool_callback = create_session_scoped_loader_callbacks(registry)

# 4. Wire into a normal ADK LlmAgent
agent = LlmAgent(
    name="Assistant",
    model="gemini-2.5-flash",
    instruction="Use search_tools to find tools, load_tool to activate them, then call them.",
    tools=[search_tools, load_tool],
    before_model_callback=before_model_callback,
    after_tool_callback=after_tool_callback,
)

With MCP servers

from google.adk.agents import LlmAgent
from google.adk.tools.mcp import MCPToolset, StdioConnectionParams
from adk_tool_search import (
    ToolRegistry,
    create_search_and_load_tools,
    create_session_scoped_loader_callbacks,
)

# Fetch tools from MCP server (but don't give to agent)
mcp = MCPToolset(connection_params=StdioConnectionParams(command="npx", args=["-y", "@modelcontextprotocol/server-github"]))
mcp_tools = await mcp.get_tools()

# Index all MCP tools
registry = ToolRegistry()
registry.register_many(mcp_tools)

# Create search/load tools + callbacks
search_tools, load_tool = create_search_and_load_tools(registry)
before_model_callback, after_tool_callback = create_session_scoped_loader_callbacks(registry)

# Wire up a normal ADK LlmAgent
agent = LlmAgent(
    name="GitHubAssistant",
    model="gemini-2.5-flash",
    instruction="Use search_tools to find tools, load_tool to activate them, then call them.",
    tools=[search_tools, load_tool],
    before_model_callback=before_model_callback,
    after_tool_callback=after_tool_callback,
)

Optional helper factory

If you prefer less boilerplate, create_tool_search_agent wraps the above wiring:

from adk_tool_search import ToolRegistry, create_tool_search_agent

registry = ToolRegistry()
registry.register_many([get_weather, send_email])

agent = create_tool_search_agent(
    name="Assistant",
    model="gemini-2.5-flash",
    registry=registry,
)

# Tools loaded via load_tool are session-scoped.
# A tool loaded in one session is not visible to other sessions unless they load it too.

Examples

# Plain function tools demo
uv run python examples/function_tools_demo.py

# MCP server demo (requires GITHUB_TOKEN)
GITHUB_TOKEN=ghp_... uv run python examples/mcp_demo.py

API

ToolRegistry

• register(tool) — Register a single tool (function, ADK tool, or MCP tool)
• register_many(tools) — Register multiple tools (rebuilds index once)
• search(query, n=5) — BM25-first search with lexical fallback for tiny registries, returns ["name: snippet", ...]
• get_tool(name) — Get tool object by exact name
• tool_count / tool_names — Introspection properties

create_search_and_load_tools(registry)

Returns (search_tools, load_tool) — the two lightweight functions to give your agent.

load_tool accepts an optional args dict for inline execution (see above).

create_session_scoped_loader_callbacks(registry, *, auto_load_from_tool_names=..., auto_load_field_keys=..., auto_load_when=None, allowed_tool_token_resolver=None)

Returns (before_model_callback, after_tool_callback) that keep loaded tools scoped to each session.

Keyword arguments for auto-load behavior:

• auto_load_from_tool_names: set of tool names eligible for response-based auto-load (default: {"use_skill"}). Set to None for field-only mode (any tool response with matching fields triggers auto-load).
• auto_load_field_keys: ordered response keys to inspect for allowed-tools tokens (default: ("allowed_tools", "allowed_tools_raw", "allowed-tools"))
• auto_load_when: optional predicate (tool_name, args, tool_response) -> bool (overrides name-based matching)
• allowed_tool_token_resolver: optional custom token resolver (tokens, registry) -> (resolved_names, unresolved_tokens)

Examples:

# Default mode (only use_skill responses can auto-load)
before_model_callback, after_tool_callback = create_session_scoped_loader_callbacks(registry)

# Field-driven mode (any tool response containing allowed-tools fields)
before_model_callback, after_tool_callback = create_session_scoped_loader_callbacks(
    registry,
    auto_load_from_tool_names=None,
)

# Custom predicate
before_model_callback, after_tool_callback = create_session_scoped_loader_callbacks(
    registry,
    auto_load_when=lambda name, args, resp: name == "policy_router" and isinstance(resp, dict),
)

create_tool_search_agent(...) (optional helper)

Convenience wrapper around manual LlmAgent wiring.

• name, model — Standard Agent params
• registry — A populated ToolRegistry
• instruction — Optional custom instruction
• always_available_tools — Tools that skip deferred loading
• **agent_kwargs — Forwarded to Agent()