Metadata-Version: 2.4
Name: adk-tool-search
Version: 0.2.1
Summary: Dynamic tool search for Google ADK — load tools on demand instead of all at once
Author: ADK Tool Search Contributors
License: Apache-2.0
Project-URL: Homepage, https://github.com/manojlds/adk-tool-search
Project-URL: Documentation, https://github.com/manojlds/adk-tool-search
Project-URL: Repository, https://github.com/manojlds/adk-tool-search
Project-URL: Issues, https://github.com/manojlds/adk-tool-search/issues
Keywords: adk,agent,tools,ai,llm,gemini,mcp,tool-search
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Requires-Dist: google-adk>=1.0.0
Requires-Dist: rank-bm25>=0.2.2
Provides-Extra: dev
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: pytest-asyncio>=0.24; extra == "dev"
Requires-Dist: pytest-cov>=4.0; extra == "dev"
Requires-Dist: pytest-timeout>=2.0; extra == "dev"
Requires-Dist: ruff>=0.1; extra == "dev"
Requires-Dist: mypy>=1.0; extra == "dev"
Requires-Dist: litellm>=1.81; extra == "dev"
Requires-Dist: python-dotenv>=1.0; extra == "dev"

# adk-tool-search

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

Implements the [Anthropic Tool Search](https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool) 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:

```python
# 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

```bash
pip install adk-tool-search
```

### Development setup

```bash
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

```python
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

```python
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:

```python
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

```bash
# 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:

```python
# 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()`
