diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 2249f93a..43047247 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -7,7 +7,7 @@ import re import shlex import time -from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Collection, Iterable from contextlib import AsyncExitStack, aclosing from dataclasses import dataclass, replace from datetime import UTC, datetime @@ -29,15 +29,13 @@ from bub.store import AsyncTapeStore, AsyncTapeStoreAdapter, InMemoryTapeStore, TapeStore, is_async_tape_store from bub.streaming import AsyncStreamEvents, StreamEvent, StreamState from bub.tape import Tape -from bub.tools import REGISTRY, Tool, ToolContext, ToolProvider, model_tools +from bub.tools import REGISTRY, Tool, ToolContext, model_tools from bub.tracing import Span, current_span from bub.turn import TurnState from bub.utils import workspace_from_state HINT_RE = re.compile(r"\$([A-Za-z0-9_.-]+)") MAX_AUTO_HANDOFF_RETRIES = 1 -# Model aliases of the tools prepared for the current request, used to recover unknown tool calls. -REQUEST_TOOLS_STATE_KEY = "_runtime_request_tools" class Agent: @@ -75,22 +73,10 @@ def __init__( ) self.framework = framework self.tools = {tool.name: tool for tool in tools} if tools is not None else REGISTRY.copy() - self.tool_sources: dict[object, Mapping[str, Tool]] = {} - self.tool_providers: list[ToolProvider] = [] self.tape_store = tape_store self.skill_dirs = skill_dirs self.model_runner = ModelRunner(self.settings, hooks=framework.get_agent_hooks()) - @property - def known_tools(self) -> dict[str, Tool]: - """Known runtime tools; later registered sources take precedence.""" - tools = self.tools.copy() - for source in self.tool_sources.values(): - for name, item in source.items(): - tools.pop(name, None) - tools[name] = item - return tools - @cached_property def tape(self) -> Tape: """Return the lazily constructed, cached tape factory for this agent. @@ -260,14 +246,13 @@ async def _run_command(self, tape: Tape, *, line: str) -> str: output = "" status = "ok" try: - known_tools = self.known_tools - if name not in known_tools: - if "bash" not in known_tools: + if name not in self.tools: + if "bash" not in self.tools: raise ValueError("bash tool is not available") # noqa: TRY301 - bash_tool = known_tools["bash"] + bash_tool = self.tools["bash"] output = bash_tool.render(await bash_tool.run(context=context, command=line)) else: - command_tool = known_tools[name] + command_tool = self.tools[name] args = _parse_args(arg_tokens) if command_tool.context: args.kwargs["context"] = context @@ -459,27 +444,50 @@ async def _run_once( prompt_text = "" else: prompt_text = _extract_text_from_parts(prompt) - known_tools = self.known_tools - if allowed_tools is not None: - from bub.builtin.tools import resolve_tool_names - - allowed_tools = resolve_tool_names(allowed_tools, all_names=known_tools) if allowed_skills is not None: allowed_skills = {name.casefold() for name in allowed_skills} tape.context.state["allowed_skills"] = list(allowed_skills) - if allowed_tools is not None: - tools = [tool for tool in known_tools.values() if tool.name in allowed_tools] - else: - tools = list(known_tools.values()) return await self._run_once_stream( tape=tape, prompt=prompt, prompt_text=prompt_text, model=model, allowed_skills=allowed_skills, - tools=tools, + tools=self._allowed_tools(tape, allowed_tools), + ) + + def _allowed_tools(self, tape: Tape, allowed_tools: Collection[str] | None) -> list[Tool]: + """Filter direct and deferred tools alike by ``allowed_tools`` and record the scope in state.""" + from bub.builtin.tools import ALLOWED_TOOLS_STATE_KEY, resolve_tool_names + + state = tape.context.state + if allowed_tools is None: + state.pop(ALLOWED_TOOLS_STATE_KEY, None) + return list(self.tools.values()) + allowed = resolve_tool_names(allowed_tools, all_names=self.tools) + state[ALLOWED_TOOLS_STATE_KEY] = sorted(allowed) + return [tool for name, tool in self.tools.items() if name in allowed] + + async def _prepare_deferred_tools(self, tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]: + """Return direct tools followed by deferred tools loaded on the tape, plus the deferred tools prompt. + + ``tool.describe`` is exposed exactly when a deferred tool remains, since it is the only way to load one. + """ + from bub.builtin.tools import ( + TOOL_DESCRIBE_TOOL_NAME, + loaded_tool_names, + render_deferred_tools_prompt, + tool_describe, ) + direct = [tool for tool in tools if not tool.deferred and tool.name != TOOL_DESCRIBE_TOOL_NAME] + deferred = {tool.name: tool for tool in tools if tool.deferred and tool.agent_use} + if not deferred: + return direct, "" + direct.append(self.tools.get(TOOL_DESCRIBE_TOOL_NAME, tool_describe)) + loaded = [deferred[name] for name in await loaded_tool_names(tape) if name in deferred] + return [*direct, *loaded], render_deferred_tools_prompt(deferred.values()) + async def _run_once_stream( self, *, @@ -490,20 +498,14 @@ async def _run_once_stream( allowed_skills: set[str] | None, tools: list[Tool], ) -> AsyncStreamEvents: - tools_prompts: list[str] = [] - for provider in self.tool_providers: - tools, tools_prompt = await provider(tools, tape) - if tools_prompt: - tools_prompts.append(tools_prompt) - tape.context.state[REQUEST_TOOLS_STATE_KEY] = [item.name for item in model_tools(tools)] - tools, tools_prompt = await self._prepare_code_mode(tools, tape) - if tools_prompt: - tools_prompts.append(tools_prompt) + # Code mode sees every allowed tool, deferred ones included; deferral applies to model-facing tools. + tools, code_mode_prompt = await self._prepare_code_mode(tools, tape) + tools, deferred_prompt = await self._prepare_deferred_tools(tools, tape) system_prompt = self._system_prompt( prompt_text, state=tape.context.state, allowed_skills=allowed_skills, - tools_prompt="\n\n".join(tools_prompts), + tools_prompt="\n\n".join(block for block in (code_mode_prompt, deferred_prompt) if block), ) resolved_model = model or self.settings.model diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index 2ec79708..6592fb5d 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -9,7 +9,7 @@ from loguru import logger from bub import inquirer as bub_inquirer -from bub.builtin.agent import REQUEST_TOOLS_STATE_KEY, Agent +from bub.builtin.agent import Agent from bub.builtin.commands import strip_command_prefix from bub.builtin.context import default_tape_context from bub.builtin.onboarding import collect_model_config @@ -356,16 +356,29 @@ async def before_tool_call( replace it with a guidance ``tool_result`` so the model can re-issue a valid call on the next step. """ + from bub.builtin.codemode import CODE_TOOLS_STATE_KEY + from bub.builtin.tools import TOOL_DESCRIBE_TOOL_NAME from bub.tools import model_tools - if (request_tools := state.get(REQUEST_TOOLS_STATE_KEY)) is not None: - available_tools = tuple(request_tools) + agent = self._get_agent(state) + + agent_tools = model_tools(agent.tools.values()) + if "_runtime_tool_names" in state: + code_tools = state.get(CODE_TOOLS_STATE_KEY) or () + available_tools = (*state["_runtime_tool_names"], *(tool_item.name for tool_item in code_tools)) else: - agent = self._get_agent(state) - available_tools = tuple(tool_item.name for tool_item in model_tools(agent.known_tools.values())) + available_tools = tuple(tool_item.name for tool_item in agent_tools) if call.tool in available_tools: return None + describe_name = TOOL_DESCRIBE_TOOL_NAME.replace(".", "_") + if describe_name in available_tools and any( + tool_item.deferred and tool_item.name == call.tool for tool_item in agent_tools + ): + return ToolCallDecision.replace( + f"Tool `{call.tool}` is not loaded. Call `{describe_name}` with its name first." + ) + matches = get_close_matches(call.tool, available_tools, n=3, cutoff=0.6) if matches: suggestions = "\n".join(f"- {name}" for name in matches) diff --git a/src/bub/builtin/tools.py b/src/bub/builtin/tools.py index ff6740c2..31b5a062 100644 --- a/src/bub/builtin/tools.py +++ b/src/bub/builtin/tools.py @@ -4,7 +4,7 @@ import contextlib import json import uuid -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from contextlib import aclosing from dataclasses import asdict from typing import TYPE_CHECKING, Any, TypedDict, cast, final @@ -19,10 +19,14 @@ if TYPE_CHECKING: from bub.builtin.agent import Agent + from bub.tape import Tape DEFAULT_COMMAND_TIMEOUT_SECONDS = 30 DEFAULT_HEADERS = {"accept": "text/markdown"} DEFAULT_REQUEST_TIMEOUT_SECONDS = 10 +TOOL_DESCRIBE_TOOL_NAME = "tool.describe" +TOOL_LOADED_EVENT = "tool.loaded" +ALLOWED_TOOLS_STATE_KEY = "_runtime_allowed_tools" def _to_model_name(name: str) -> str: @@ -104,6 +108,31 @@ def render_tools_prompt(tools: Iterable[Tool]) -> str: return f"\n{'\n'.join(lines)}\n" +def render_deferred_tools_prompt(tools: Iterable[Tool]) -> str: + """List deferred tools by model-facing name so the model knows what it can load.""" + names = "\n".join(f"- {_to_model_name(tool_item.name)}" for tool_item in tools) + if not names: + return "" + return ( + "\n" + f"These tools are not loaded yet. Call {_to_model_name(TOOL_DESCRIBE_TOOL_NAME)} with their names to load " + "their definitions; loaded tools become callable from the next step.\n" + f"{names}\n" + ) + + +async def loaded_tool_names(tape: Tape) -> list[str]: + """Return runtime names of tools loaded on this tape, in load order.""" + names: dict[str, None] = {} + for entry in await tape.search(tape.query().kinds("event")): + if entry.payload.get("name") != TOOL_LOADED_EVENT: + continue + data = entry.payload.get("data") + if isinstance(data, Mapping): + names.update(dict.fromkeys(str(name) for name in data.get("names") or ())) + return list(names) + + def _raise_for_failed_shell(returncode: int | None, output: str) -> None: if returncode in (None, 0): return @@ -167,6 +196,17 @@ class ToolFailure(TypedDict): error: str +class ToolDefinition(TypedDict): + name: str + description: str + parameters: dict[str, Any] + + +class ToolDescriptions(TypedDict): + tools: list[ToolDefinition] + unknown: list[str] + + class TapeInfoResult(TypedDict): name: str entries: int @@ -358,6 +398,40 @@ def skill_describe(name: str | None = None, *, context: ToolContext) -> SkillLis return {"name": skill.name, "location": str(skill.location), "content": skill.body() or ""} +@tool(context=True, name=TOOL_DESCRIBE_TOOL_NAME, preserve=True) +async def tool_describe(names: list[str], *, context: ToolContext) -> ToolDescriptions: + """Load tools by name and return their definitions. Deferred tools become callable from the next step.""" + agent = _get_agent(context) + allowed_tools = context.state.get(ALLOWED_TOOLS_STATE_KEY) + available = { + name: tool_item + for name, tool_item in agent.tools.items() + if tool_item.agent_use and (allowed_tools is None or name in allowed_tools) + } + index = _tool_name_index(available) + loaded = set(await loaded_tool_names(context.tape)) + definitions: list[ToolDefinition] = [] + unknown: list[str] = [] + newly_loaded: list[str] = [] + for name in names: + resolved = index.get(name.strip().casefold()) + if resolved is None: + unknown.append(name) + continue + tool_item = available[resolved] + definitions.append({ + "name": _to_model_name(resolved), + "description": tool_item.description, + "parameters": tool_item.parameters, + }) + if tool_item.deferred and resolved not in loaded: + loaded.add(resolved) + newly_loaded.append(resolved) + if newly_loaded: + await context.tape.append_event(TOOL_LOADED_EVENT, {"names": newly_loaded}) + return {"tools": definitions, "unknown": unknown} + + @tool(context=True, name="tape.info", renderer=_render_tape_info) async def tape_info(context: ToolContext) -> TapeInfoResult: """Get information about the current tape, such as number of entries and anchors.""" @@ -434,7 +508,7 @@ async def run_subagent(param: SubAgentInput, *, context: ToolContext) -> SubAgen else: subagent_session = param.session state = {**context.state, "session_id": subagent_session} - allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.known_tools) + allowed_tools = resolve_tool_names(param.allowed_tools or None, exclude={"subagent"}, all_names=agent.tools) output = "" errors: list[str] = [] stream = await agent.run_stream( diff --git a/src/bub/tools.py b/src/bub/tools.py index c9346417..1b873fb3 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -149,6 +149,8 @@ class Tool: """Keep the tool directly callable by the model in code mode; others are reachable only from code.""" output_schema: dict[str, Any] | None = None """JSON schema of the structured result, used to describe the tool to model-written code.""" + deferred: bool = False + """Load the tool on demand: only its name is listed until ``tool.describe`` loads its definition.""" @property def code_use(self) -> bool: @@ -189,6 +191,7 @@ def from_callable( agent_use: bool = True, renderer: Callable[[Any], str] | None = None, preserve: bool = False, + deferred: bool = False, ) -> Tool: signature = inspect.signature(func) if context and "context" not in signature.parameters: @@ -216,13 +219,10 @@ def validated(*args: Any, **kwargs: Any) -> Any: renderer=renderer, preserve=preserve, output_schema=_output_schema(func), + deferred=deferred, ) -type ToolProvider = Callable[[list[Tool], Tape], Awaitable[tuple[list[Tool], str]]] -"""Prepare registered tools and a prompt fragment for one model request.""" - - def model_tools(tools: Iterable[Tool]) -> list[Tool]: """Convert agent-enabled runtime tools into model-safe aliases.""" return [replace(tool_item, name=tool_item.name.replace(".", "_")) for tool_item in tools if tool_item.agent_use] @@ -562,6 +562,7 @@ def tool( agent_use: bool = ..., renderer: Callable[[Any], str] | None = ..., preserve: bool = ..., + deferred: bool = ..., ) -> Tool: ... @@ -576,6 +577,7 @@ def tool( agent_use: bool = ..., renderer: Callable[[Any], str] | None = ..., preserve: bool = ..., + deferred: bool = ..., ) -> Callable[[Callable], Tool]: ... @@ -589,11 +591,13 @@ def tool( agent_use: bool = True, renderer: Callable[[Any], str] | None = None, preserve: bool = False, + deferred: bool = False, ) -> Tool | Callable[[Callable], Tool]: """Decorator to convert a function into a Tool instance. Tools should return structured results; ``renderer`` turns such a result into the plain text shown to the model outside code mode (defaults to JSON for non-strings). + ``deferred`` tools are loaded on demand through ``tool.describe``. """ def decorator(func: Callable) -> Tool: @@ -618,6 +622,7 @@ def handler(*args: Any, **kwargs: Any) -> Any: renderer=renderer, preserve=preserve, output_schema=_output_schema(func), + deferred=deferred, ) else: result = Tool.from_callable( @@ -628,6 +633,7 @@ def handler(*args: Any, **kwargs: Any) -> Any: agent_use=agent_use, renderer=renderer, preserve=preserve, + deferred=deferred, ) tool_instance = _add_logging(result) REGISTRY[tool_instance.name] = tool_instance diff --git a/tests/test_builtin_agent.py b/tests/test_builtin_agent.py index 836aa86c..f36b2eb1 100644 --- a/tests/test_builtin_agent.py +++ b/tests/test_builtin_agent.py @@ -55,8 +55,6 @@ async def build_prompt(message: dict[str, Any], session_id: str, state: dict[str agent.command_prefix = agent.settings.command_prefix agent.framework = framework agent.tools = REGISTRY.copy() - agent.tool_sources = {} - agent.tool_providers = [] agent.tape_store = None agent.skill_dirs = None agent.model_runner = _FakeModelRunner(agent.settings) diff --git a/tests/test_builtin_hook_impl.py b/tests/test_builtin_hook_impl.py index 031130ef..a2cc26f6 100644 --- a/tests/test_builtin_hook_impl.py +++ b/tests/test_builtin_hook_impl.py @@ -41,7 +41,6 @@ def __init__(self, home: Path, *, tape: Tape | None = None) -> None: self.command_prefix = "," self.settings = SimpleNamespace(home=home) self.tools = REGISTRY.copy() - self.known_tools = self.tools # A real in-memory async tape so load_state's recovery path runs against # the same store the tests write `model_switch` events to. self.tape = tape if tape is not None else _fake_tape(home) @@ -500,10 +499,9 @@ def test_before_tool_call_recovers_tool_outside_current_request(tmp_path: Path) _, impl, _ = _build_impl(tmp_path) import asyncio - from bub.builtin.agent import REQUEST_TOOLS_STATE_KEY from bub.hooks.interception import ToolCall - state = {REQUEST_TOOLS_STATE_KEY: ["bash"]} + state = {"_runtime_tool_names": ("bash",)} async def _do(name: str): return await impl.before_tool_call(ToolCall(run_id="r", tool=name, arguments={}), state=state) diff --git a/tests/test_subagent_tool.py b/tests/test_subagent_tool.py index d77f5930..1826de24 100644 --- a/tests/test_subagent_tool.py +++ b/tests/test_subagent_tool.py @@ -21,7 +21,6 @@ def __init__(self, state: dict[str, Any]) -> None: class FakeAgent: def __init__(self) -> None: self.tools = REGISTRY.copy() - self.known_tools = self.tools self.run_stream = AsyncMock(side_effect=self._run_stream) async def _run_stream(self, **kwargs: Any) -> AsyncStreamEvents: diff --git a/tests/test_tool_loading.py b/tests/test_tool_loading.py new file mode 100644 index 00000000..a0008a82 --- /dev/null +++ b/tests/test_tool_loading.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from any_llm.types.completion import ChatCompletion + +from bub.builtin.agent import Agent +from bub.framework import BubFramework +from bub.tools import Tool + + +def _completion(message: dict[str, Any]) -> ChatCompletion: + return ChatCompletion.model_validate({ + "id": "reply", + "model": "test-model", + "created": 0, + "object": "chat.completion", + "choices": [ + {"index": 0, "finish_reason": "tool_calls" if "tool_calls" in message else "stop", "message": message} + ], + }) + + +def _tool_call(call_id: str, name: str, arguments: dict[str, Any]) -> dict[str, Any]: + return { + "role": "assistant", + "tool_calls": [ + {"id": call_id, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}} + ], + } + + +def _tool_names(request: dict[str, Any]) -> list[str]: + return [item["function"]["name"] for item in request.get("tools") or []] + + +def _system_prompt(request: dict[str, Any]) -> str: + return "\n".join(message["content"] for message in request["messages"] if message["role"] == "system") + + +@pytest.fixture +def framework(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> BubFramework: + monkeypatch.setenv("BUB_HOME", str(tmp_path)) + framework = BubFramework(config_file=tmp_path / "config.yml") + framework.workspace = tmp_path + framework.load_builtin_hooks() + return framework + + +def _scripted_provider(monkeypatch: pytest.MonkeyPatch, replies: list[dict[str, Any]]) -> list[dict[str, Any]]: + requests: list[dict[str, Any]] = [] + + class Provider: + SUPPORTS_COMPLETION_STREAMING = False + + async def acompletion(self, **kwargs: Any) -> ChatCompletion: + requests.append(kwargs) + return _completion(replies[len(requests) - 1]) + + monkeypatch.setattr("bub.builtin.model_runner.AnyLLM.create", lambda *args, **kwargs: Provider()) + return requests + + +@pytest.mark.asyncio +async def test_deferred_tool_is_loaded_through_tool_describe_and_persists_on_tape( + framework: BubFramework, monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[str] = [] + + def lookup(name: str) -> str: + calls.append(name) + return f"Hello {name}" + + direct = Tool.from_callable(lambda: "direct", name="direct") + supplied = Tool.from_callable(lookup, name="provider.lookup", deferred=True) + other = Tool.from_callable(lambda: "other", name="other", deferred=True) + requests = _scripted_provider( + monkeypatch, + [ + _tool_call("describe", "tool_describe", {"names": ["provider_lookup"]}), + _tool_call("lookup", "provider_lookup", {"name": "Ada"}), + {"role": "assistant", "content": "Hello Ada"}, + {"role": "assistant", "content": "Again"}, + ], + ) + agent = Agent(framework, tools=[direct, supplied, other], skill_dirs=[]) + + stream = await agent.run_stream(session_id="loading", prompt="Greet Ada.", model="openrouter:test-model") + events = [event async for event in stream] + + assert any(event.data.get("text") == "Hello Ada" for event in events if event.kind == "final") + assert calls == ["Ada"] + assert _tool_names(requests[0]) == ["direct", "tool_describe"] + assert _tool_names(requests[1]) == ["direct", "tool_describe", "provider_lookup"] + describe_result = next(message for message in requests[1]["messages"] if message["role"] == "tool") + assert json.loads(describe_result["content"])["tools"][0]["name"] == "provider_lookup" + prompt = _system_prompt(requests[0]) + assert "" in prompt + assert "- provider_lookup" in prompt + assert "- other" in prompt + + stream = await agent.run_stream(session_id="loading", prompt="Again.", model="openrouter:test-model") + _ = [event async for event in stream] + assert _tool_names(requests[3]) == ["direct", "tool_describe", "provider_lookup"] + + +@pytest.mark.asyncio +async def test_allowed_tools_filter_covers_direct_and_deferred_tools( + framework: BubFramework, monkeypatch: pytest.MonkeyPatch +) -> None: + direct = Tool.from_callable(lambda: "direct", name="direct") + denied = Tool.from_callable(lambda: "denied", name="denied") + allowed_deferred = Tool.from_callable(lambda: "lookup", name="provider.lookup", deferred=True) + denied_deferred = Tool.from_callable(lambda: "secret", name="secret", deferred=True) + requests = _scripted_provider( + monkeypatch, + [ + _tool_call("describe", "tool_describe", {"names": ["provider_lookup", "secret"]}), + {"role": "assistant", "content": "done"}, + ], + ) + agent = Agent(framework, tools=[direct, denied, allowed_deferred, denied_deferred], skill_dirs=[]) + + stream = await agent.run_stream( + session_id="allowed", + prompt="Load tools.", + model="openrouter:test-model", + allowed_tools=["direct", "provider_lookup"], + ) + _ = [event async for event in stream] + + assert _tool_names(requests[0]) == ["direct", "tool_describe"] + assert _tool_names(requests[1]) == ["direct", "tool_describe", "provider_lookup"] + prompt = _system_prompt(requests[0]) + assert "- provider_lookup" in prompt + assert "secret" not in prompt + describe_result = next(message for message in requests[1]["messages"] if message["role"] == "tool") + assert json.loads(describe_result["content"])["unknown"] == ["secret"] + + +@pytest.mark.asyncio +async def test_tool_describe_is_hidden_without_deferred_tools( + framework: BubFramework, monkeypatch: pytest.MonkeyPatch +) -> None: + from bub.builtin.tools import tool_describe + + direct = Tool.from_callable(lambda: "direct", name="direct") + requests = _scripted_provider(monkeypatch, [{"role": "assistant", "content": "done"}]) + agent = Agent(framework, tools=[direct, tool_describe], skill_dirs=[]) + + stream = await agent.run_stream(session_id="plain", prompt="Hi.", model="openrouter:test-model") + _ = [event async for event in stream] + + assert _tool_names(requests[0]) == ["direct"] + assert "" not in _system_prompt(requests[0]) + + +@pytest.mark.asyncio +async def test_calling_unloaded_deferred_tool_returns_load_guidance( + framework: BubFramework, monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[str] = [] + supplied = Tool.from_callable(lambda: calls.append("lookup") or "ok", name="provider.lookup", deferred=True) + requests = _scripted_provider( + monkeypatch, + [_tool_call("lookup", "provider_lookup", {}), {"role": "assistant", "content": "done"}], + ) + agent = Agent(framework, tools=[supplied], skill_dirs=[]) + + stream = await agent.run_stream(session_id="unloaded", prompt="Look up.", model="openrouter:test-model") + _ = [event async for event in stream] + + assert calls == [] + result = next(message for message in requests[1]["messages"] if message["role"] == "tool") + assert result["content"] == "Tool `provider_lookup` is not loaded. Call `tool_describe` with its name first." + + +@pytest.mark.asyncio +async def test_code_mode_exposes_deferred_tools_to_code_without_loading( + framework: BubFramework, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from bub.builtin.codemode import CODE_TOOLS_STATE_KEY, run_code + + direct = Tool.from_callable(lambda: "direct", name="direct") + supplied = Tool.from_callable(lambda: "lookup", name="provider.lookup", deferred=True) + requests = _scripted_provider(monkeypatch, [{"role": "assistant", "content": "done"}]) + agent = Agent(framework, tools=[direct, supplied, run_code], skill_dirs=[]) + state: dict[str, Any] = {"code_mode": True, "_runtime_workspace": str(tmp_path)} + + stream = await agent.run_stream(session_id="code", prompt="Hi.", model="openrouter:test-model", state=state) + _ = [event async for event in stream] + + assert _tool_names(requests[0]) == ["run_code"] + assert [tool.name for tool in state[CODE_TOOLS_STATE_KEY]] == ["direct", "provider_lookup"] + assert "" not in _system_prompt(requests[0]) diff --git a/tests/test_tool_providers.py b/tests/test_tool_providers.py deleted file mode 100644 index a5e5ea4b..00000000 --- a/tests/test_tool_providers.py +++ /dev/null @@ -1,114 +0,0 @@ -from __future__ import annotations - -import json -from dataclasses import replace -from pathlib import Path -from typing import Any - -import pytest -from any_llm.types.completion import ChatCompletion - -from bub.builtin.agent import Agent -from bub.builtin.codemode import run_code -from bub.framework import BubFramework -from bub.tape import Tape -from bub.tools import Tool - - -@pytest.mark.asyncio -@pytest.mark.parametrize("code_mode", [False, True]) -@pytest.mark.parametrize("reverse_providers", [False, True]) -async def test_discovery_precedence_is_independent_of_provider_order_for_native_and_code_calls( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, code_mode: bool, reverse_providers: bool -) -> None: - monkeypatch.setenv("BUB_HOME", str(tmp_path)) - framework = BubFramework(config_file=tmp_path / "config.yml") - framework.workspace = tmp_path - framework.load_builtin_hooks() - direct = Tool.from_callable(lambda: "direct", name="direct") - denied = Tool.from_callable(lambda: "denied", name="denied") - pending = Tool.from_callable(lambda: "pending", name="provider.pending") - calls: list[str] = [] - - def lookup(name: str) -> str: - calls.append(name) - return f"Hello {name}" - - supplied = Tool.from_callable(lookup, name="provider.lookup") - - async def select(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]: - return [item for item in tools if item is not pending], "Use provider_lookup to greet the requested person." - - async def guide(tools: list[Tool], tape: Tape) -> tuple[list[Tool], str]: - return [ - replace(item, renderer=str.upper) if item.name == supplied.name else item for item in tools - ], "Keep the greeting brief." - - requests: list[dict[str, Any]] = [] - - class Provider: - SUPPORTS_COMPLETION_STREAMING = False - - async def acompletion(self, **kwargs: Any) -> ChatCompletion: - requests.append(kwargs) - message: dict[str, Any] = {"role": "assistant", "content": "Hello Ada"} - if len(requests) == 1: - name = "run_code" if code_mode else "provider_lookup" - arguments = {"code": "print(await tools.provider_lookup(name='Ada'))"} if code_mode else {"name": "Ada"} - message = { - "role": "assistant", - "tool_calls": [ - { - "id": "lookup", - "type": "function", - "function": {"name": name, "arguments": json.dumps(arguments)}, - } - ], - } - return ChatCompletion.model_validate({ - "id": "reply", - "model": "test-model", - "created": 0, - "object": "chat.completion", - "choices": [ - { - "index": 0, - "finish_reason": "tool_calls" if "tool_calls" in message else "stop", - "message": message, - } - ], - }) - - monkeypatch.setattr("bub.builtin.model_runner.AnyLLM.create", lambda *args, **kwargs: Provider()) - agent = Agent(framework, tools=[direct, run_code], skill_dirs=[]) - agent.tool_sources["earlier"] = {supplied.name: Tool.from_callable(lambda: "wrong", name=supplied.name)} - agent.tool_sources["later"] = {supplied.name: supplied, denied.name: denied, pending.name: pending} - agent.tool_providers = [guide, select] if reverse_providers else [select, guide] - stream = await agent.run_stream( - session_id="provider", - prompt="Greet Ada.", - model="openrouter:test-model", - allowed_tools=["direct", "provider_lookup", "provider_pending", "run_code"], - state={"code_mode": code_mode}, - ) - events = [event async for event in stream] - assert any(event.data.get("text") == "Hello Ada" for event in events if event.kind == "final") - assert calls == ["Ada"] - # Request preparation does not register discovered tools on the agent. - assert agent.tools.keys() == {"direct", "run_code"} - definitions = {item["function"]["name"]: item["function"] for item in requests[0]["tools"]} - assert definitions.keys() == ({"run_code"} if code_mode else {"direct", "provider_lookup"}) - if code_mode: - stub = next((tmp_path / "codemode").rglob("*.pyi")).read_text() - assert "provider_pending" not in stub - system = "\n".join(message["content"] for message in requests[0]["messages"] if message["role"] == "system") - guidance = ( - ["Keep the greeting brief.", "Use provider_lookup"] - if reverse_providers - else ["Use provider_lookup", "Keep the greeting brief."] - ) - assert system.index(guidance[0]) < system.index(guidance[1]) - expected = "Hello Ada" if code_mode else "HELLO ADA" - assert any( - expected in message.get("content", "") for message in requests[1]["messages"] if message["role"] == "tool" - ) diff --git a/website/src/content/docs/docs/build/tools.mdx b/website/src/content/docs/docs/build/tools.mdx index 1b73791c..b0f1a580 100644 --- a/website/src/content/docs/docs/build/tools.mdx +++ b/website/src/content/docs/docs/build/tools.mdx @@ -64,13 +64,15 @@ If `run_code` is not allowed (for example, a subagent restricted with `allowed_t `run_code` hands the code to the session [environment](#run-tools-in-an-environment)'s `run_code`, together with a `call_tool` callback. Tool calls always run on the host through that callback, so hooks still see every call. `timeout_seconds` cancels the environment call. The builtin `LocalEnvironment` starts a fresh Python process for every call (`sys.executable`, working directory is the workspace) and forwards tool calls as JSON lines over its stdin and stdout, so arguments and results must be JSON-serializable. When the code finishes, fails or times out, it kills the process and everything it started. That process runs on the host with the same permissions as `bash`, so enable code mode only where `bash` would be acceptable. The return annotation of a tool function (or `Tool.output_schema`, for tools built by hand) determines the result type shown in the stub. -### Tool discovery and request preparation +### Load tools on demand -Register directly available tools in `Agent.tools`. For a separate discovery inventory, set `agent.tool_sources[source]` to a mapping of runtime names to `Tool` instances. `source` is a plugin-owned key used to update or remove that inventory. `Agent.known_tools` combines direct tools and source inventories for commands and `allowed_tools`; later registered sources win duplicate names, independently of provider order. Updating an existing source preserves its priority. +Mark a tool `deferred=True` (on `@tool`, `Tool.from_callable`, or with `dataclasses.replace(tool, deferred=True)`) to keep its definition out of model requests until it is needed. Tools without the flag are loaded directly. -Append request-time preparation callables to `agent.tool_providers`. The `ToolProvider` contract in `bub.tools` is `async (tools, tape) -> (tools, prompt_fragment)`. Providers receive the filtered toolset, run in their supplied order, and pass their result to the next provider. They may select definitions and contribute discovery guidance, but must not reintroduce tools excluded by the scope. Builtin preparation runs last, so code mode uses the final selected toolset. - -Request preparation does not change `Agent.tools`. Known tools that pass the scope reach the providers, so a provider that defers discovered tools must remove them from the request itself. The tools left after the providers form the request toolset: only those calls run, and calls to other names receive a guidance result. Plugins own discovery state, source removal and connection cleanup. Full native definitions are sent through tool schemas instead of repeated in system. Keep discovery summaries stable and append native definitions in the conversation's first-discovery order; request serialization determines how much of the prefix remains cacheable. +- The system prompt lists the model-facing names of all allowed deferred tools in a `` block, and the model gets `tool_describe(names)`. It is exposed only when at least one allowed deferred tool exists. +- `tool_describe` returns the name, description, and parameters of each requested tool and appends a `tool.loaded` tape event with the runtime names of newly loaded deferred tools. Unknown or disallowed names are returned under `unknown`. +- On every model request Bub reads all `tool.loaded` events from the tape and appends those tools, in load order, after the direct tools. Loaded tools stay available for later steps and turns of the session until the tape is reset. +- `allowed_tools` filters direct and deferred tools alike: a deferred tool outside the scope is neither listed nor loadable. +- Calling a deferred tool before loading it returns guidance to call `tool_describe` first. Code mode sees every allowed tool: deferred tools that are not preserved appear in the stub without loading, and deferral applies only to the remaining model-facing tools. ## Run tools in an environment @@ -93,7 +95,7 @@ The `REGISTRY` lives in [`bub.tools`](https://github.com/bubbuild/bub/blob/main/ REGISTRY: dict[str, Tool] = {} ``` -Every `@tool` call mutates this dict at **import time**. `Agent(tools=None)` snapshots it when created; later imports do not update existing agents. Add tools to an instance's `Agent.tools`, or register a discovery inventory in `agent.tool_sources` and a preparation callable in `agent.tool_providers`. +Every `@tool` call mutates this dict at **import time**. `Agent(tools=None)` snapshots it when created; later imports do not update existing agents. Add tools to an instance's `Agent.tools`, and mark them `deferred=True` to [load them on demand](#load-tools-on-demand). ## Import the tools module from your plugin diff --git a/website/src/content/docs/zh-cn/docs/build/tools.mdx b/website/src/content/docs/zh-cn/docs/build/tools.mdx index 6475f87f..dbf841c8 100644 --- a/website/src/content/docs/zh-cn/docs/build/tools.mdx +++ b/website/src/content/docs/zh-cn/docs/build/tools.mdx @@ -65,13 +65,15 @@ Code mode 让模型在 Python 中调用工具,而不必一次只发起一个 `run_code` 会把代码和一个 `call_tool` 回调一起交给当前 session [执行环境](#在执行环境中运行工具)的 `run_code` 执行。工具调用始终经这个回调在宿主上执行,所以 hook 仍能看到每一次调用。超过 `timeout_seconds` 时会取消执行环境里的执行。builtin 的 `LocalEnvironment` 每次调用都会启动一个新的 Python 进程(使用 `sys.executable`,工作目录为 workspace),工具调用以 JSON 行的形式经进程的 stdin/stdout 转发,因此参数和结果都必须能 JSON 序列化。代码执行完、出错或超时后,它会杀掉该进程及其启动的所有子进程。这个进程直接跑在宿主上,权限与 `bash` 相同,因此只应在允许使用 `bash` 的场景下开启 code mode。stub 中的结果类型来自工具函数的返回注解(手动构造的工具则来自 `Tool.output_schema`)。 -### 工具发现与请求准备 +### 按需加载工具 -直接可用的工具注册在 `Agent.tools`。独立的发现清单通过 `agent.tool_sources[source]` 登记,值为运行时名称到 `Tool` 实例的映射。`source` 是插件管理的键,用于更新或移除自己的清单。`Agent.known_tools` 合并直接工具和各来源清单,供命令查找及 `allowed_tools` 使用;后注册的来源优先处理重名,与 provider 顺序无关。更新已有来源不改变优先级。 +给工具标记 `deferred=True`(可用于 `@tool`、`Tool.from_callable`,或 `dataclasses.replace(tool, deferred=True)`),它的定义就不会进入模型请求,直到被加载。未标记的工具直接加载。 -向 `agent.tool_providers` 追加请求准备函数。`bub.tools` 中的 `ToolProvider` 契约是 `async (tools, tape) -> (tools, prompt_fragment)`。Provider 接收过滤后的工具集合,按提供的顺序执行,并将结果交给下一项。它们可以选择完整定义并贡献发现提示,但不得重新加入被 scope 排除的工具。Builtin preparation 最后执行,因此 code mode 使用最终选中的工具集合。 - -请求准备不会修改 `Agent.tools`。通过 scope 的已知工具都会交给 provider,因此需要延迟暴露的发现工具必须由 provider 自己从本次请求中移除。provider 处理后剩下的工具构成本次请求的工具集合:只有这些调用会执行,调用其他名称会收到引导结果。发现状态、来源移除和连接清理由插件管理。完整原生定义通过工具 schema 发送,系统提示词不再重复列出。保持发现摘要稳定,并按当前会话的首次发现顺序追加原生定义;此前前缀有多少可以命中缓存,取决于请求序列化方式。 +- 系统提示词的 `` 块列出所有允许范围内按需加载工具的模型名称,模型可调用 `tool_describe(names)` 加载它们。只有存在至少一个允许的按需加载工具时才会暴露该工具。 +- `tool_describe` 返回所请求工具的名称、描述和参数,并把新加载的按需工具的运行时名称作为 `tool.loaded` 事件写入 tape。未知或不允许的名称放在 `unknown` 中返回。 +- 每次请求模型时,Bub 从 tape 读取所有 `tool.loaded` 事件,按加载顺序把这些工具追加到直接加载的工具之后。已加载的工具在该 session 后续的步骤和轮次中持续可用,直到 tape 被重置。 +- `allowed_tools` 同时过滤直接加载和按需加载的工具:范围外的按需工具既不会被列出,也无法加载。 +- 未加载就调用按需工具时,会返回提示,要求先调用 `tool_describe`。code mode 能看到所有允许的工具:非 preserve 的按需工具无需加载即出现在 stub 中,按需加载只作用于剩下的模型可见工具。 ## 在执行环境中运行工具 @@ -94,7 +96,7 @@ Code mode 让模型在 Python 中调用工具,而不必一次只发起一个 REGISTRY: dict[str, Tool] = {} ``` -每一次 `@tool` 调用都会在**导入时**修改这个字典。`Agent(tools=None)` 在创建时取得其快照,之后导入的工具不会自动更新已有 Agent。可直接向实例的 `Agent.tools` 添加工具,或分别通过 `agent.tool_sources` 登记发现清单、通过 `agent.tool_providers` 登记请求准备函数。 +每一次 `@tool` 调用都会在**导入时**修改这个字典。`Agent(tools=None)` 在创建时取得其快照,之后导入的工具不会自动更新已有 Agent。可直接向实例的 `Agent.tools` 添加工具,标记 `deferred=True` 即可[按需加载](#按需加载工具)。 ## 在插件里导入工具模块