From 1859d8516c4090bed04d50e6dd6b15f0e87b9974 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Thu, 6 Aug 2026 15:55:22 +0800 Subject: [PATCH] feat: implement reasoning effort management in session handling and tools Signed-off-by: Frost Ming --- src/bub/builtin/agent.py | 3 +- src/bub/builtin/hook_impl.py | 17 ++++-- src/bub/builtin/model_runner.py | 15 ++++-- src/bub/builtin/tools.py | 44 ++++++---------- src/bub/tools.py | 33 +++++++++++- tests/test_agent_hooks.py | 2 +- tests/test_builtin_agent.py | 35 +++++++++++++ tests/test_builtin_hook_impl.py | 13 +++++ tests/test_builtin_model_runner.py | 26 ++++++++++ tests/test_builtin_tools.py | 83 ++++++++++++++---------------- tests/test_tools.py | 65 ++++++++++++++++++++++- 11 files changed, 253 insertions(+), 83 deletions(-) diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 874c99ef..8406d111 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -32,6 +32,7 @@ REGISTRY, Tool, ToolContext, + model_tools, ) from bub.turn import TurnState from bub.utils import workspace_from_state @@ -352,8 +353,6 @@ async def _run_once_stream( allowed_skills: set[str] | None, tools: list[Tool], ) -> AsyncStreamEvents: - from bub.builtin.tools import model_tools - system_prompt = self._system_prompt( prompt_text, state=tape.context.state, allowed_skills=allowed_skills, tools=tools ) diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index 6494cf4d..b74e5519 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -90,6 +90,16 @@ async def _recover_session_model(self, session_id: str) -> str | None: return str(model) if model else None return None + async def _recover_session_reasoning_effort(self, session_id: str) -> str | None: + """Recover the latest per-session reasoning effort override.""" + session = self._get_agent().tape.session_tape(session_id, self.framework.workspace) + entries = list(await session.store.fetch_all(session.query().kinds("event"))) + for entry in reversed(entries): + if entry.kind == "event" and entry.payload.get("name") == "reasoning_effort_switch": + reasoning_effort = (entry.payload.get("data") or {}).get("reasoning_effort") + return str(reasoning_effort) if reasoning_effort else None + return None + @staticmethod async def _discard_message(_: ChannelMessage) -> None: return @@ -149,6 +159,8 @@ async def load_state(self, message: ChannelMessage, session_id: str) -> TurnStat # fresh/unknown session never inherits another session's model. if model := await self._recover_session_model(session_id): state["model"] = model + if reasoning_effort := await self._recover_session_reasoning_effort(session_id): + state["reasoning_effort"] = reasoning_effort if model := field_of(message, "context", {}).get("model"): state["model"] = model if thread_id := field_of(message, "context", {}).get("thread_id"): @@ -161,9 +173,8 @@ async def save_state(self, session_id: str, state: TurnState, message: ChannelMe lifespan = field_of(message, "lifespan") if lifespan is not None: await lifespan.__aexit__(tp, value, traceback) - # The per-session model override is persisted on the session tape by the - # ``model`` tool itself (a ``model_switch`` event, merged back at end of - # turn), so nothing to write here — this hook only closes the lifespan. + # Per-session completion overrides are persisted by their tools as tape + # events, so nothing to write here — this hook only closes the lifespan. @hookimpl async def build_prompt(self, message: ChannelMessage, session_id: str, state: TurnState) -> str | list[dict]: diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index fbbb4a71..b0301a37 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -81,11 +81,15 @@ def create_llm_client(candidate: ModelCandidate, client_kwargs: dict[str, Any]) return AnyLLM.create(candidate.provider, **client_kwargs) async def completion_response( - self, *, model: str, messages: list[dict[str, Any]], tools: list[Tool], max_tokens: int | None = None + self, + *, + model: str, + messages: list[dict[str, Any]], + tools: list[Tool], + max_tokens: int | None = None, + reasoning_effort: str | None = None, ) -> CompletionResult: - from bub.builtin.tools import completion_tools - - tool_payloads = completion_tools(tools) or None + tool_payloads = [tool.to_schema() for tool in tools] or None completion_messages: list[dict[str, Any] | ChatCompletionMessage] = list(messages) clients = list(self.iter_llm_clients(model)) completion_error: Exception | None = None @@ -101,6 +105,8 @@ async def completion_response( "max_tokens": max_tokens if max_tokens is not None else self.settings.max_tokens, "stream": streaming, } + if reasoning_effort is not None: + completion_kwargs["reasoning_effort"] = reasoning_effort return cast("CompletionResult", await llm.acompletion(**completion_kwargs)) except Exception as exc: if completion_error is None: @@ -175,6 +181,7 @@ async def fire_after(error: Exception | None = None) -> None: messages=list(request.messages), tools=tools, max_tokens=request.max_tokens, + reasoning_effort=tape.context.state.get("reasoning_effort"), ) async for event in self._completion_events(completion, state, output): yield event diff --git a/src/bub/builtin/tools.py b/src/bub/builtin/tools.py index 9527bf7c..3e4cb916 100644 --- a/src/bub/builtin/tools.py +++ b/src/bub/builtin/tools.py @@ -4,11 +4,9 @@ import json import uuid from collections.abc import Iterable -from dataclasses import replace from pathlib import Path from typing import TYPE_CHECKING, cast -from openai.types.chat import ChatCompletionToolParam from pydantic import BaseModel, Field from bub.builtin.shell_manager import shell_manager @@ -72,11 +70,6 @@ def resolve_tool_names(names: Iterable[str] | None = None, *, exclude: Iterable[ return resolved - excluded -def model_tools(tools: Iterable[Tool]) -> list[Tool]: - """Convert runtime tool names into model-safe aliases.""" - return [replace(tool_item, name=_to_model_name(tool_item.name)) for tool_item in tools] - - def _tool_signature(tool_item: Tool) -> str: properties = tool_item.parameters.get("properties", {}) if not isinstance(properties, dict) or not properties: @@ -90,10 +83,11 @@ def _tool_signature(tool_item: Tool) -> str: def render_tools_prompt(tools: Iterable[Tool]) -> str: """Render a human-readable description of tools for builtin agent prompts.""" - if not tools: + agent_tools = [tool_item for tool_item in tools if tool_item.agent_use] + if not agent_tools: return "" lines = [] - for tool_item in tools: + for tool_item in agent_tools: line = f"- {_tool_signature(tool_item)}" if tool_item.description: line += f": {tool_item.description}" @@ -101,21 +95,6 @@ def render_tools_prompt(tools: Iterable[Tool]) -> str: return f"\n{'\n'.join(lines)}\n" -def completion_tools(tools: Iterable[Tool]) -> list[ChatCompletionToolParam]: - """Build any-llm completion tool payloads from Bub tools.""" - return [ - { - "type": "function", - "function": { - "name": tool_item.name, - "description": tool_item.description, - "parameters": tool_item.parameters, - }, - } - for tool_item in tools - ] - - def _raise_for_failed_shell(returncode: int | None, output: str) -> None: if returncode in (None, 0): return @@ -376,7 +355,7 @@ async def run_subagent(param: SubAgentInput, *, context: ToolContext) -> str: return output -@tool(name="help") +@tool(name="help", agent_use=False) def show_help() -> str: """Show a help message.""" return ( @@ -399,7 +378,7 @@ def show_help() -> str: ) -@tool(name="quit", context=True) +@tool(name="quit", context=True, agent_use=False) async def quit_tool(*, context: ToolContext) -> str: """Abort the tasks of the current session. DO NOT use it in a normal workflow.""" agent = _get_agent(context) @@ -409,7 +388,7 @@ async def quit_tool(*, context: ToolContext) -> str: return "Session tasks stopped." -@tool(name="model", context=True) +@tool(name="model", context=True, agent_use=False) async def set_model(model_id: str, *, context: ToolContext) -> str: """Switch the model for THIS session. Invoke as the `,model ` command. @@ -425,6 +404,17 @@ async def set_model(model_id: str, *, context: ToolContext) -> str: return f"Session model set to {model_id} (applies from the next turn)." +@tool(name="reasoning_effort", context=True, agent_use=False) +async def set_reasoning_effort(reasoning_effort: str, *, context: ToolContext) -> str: + """Set the reasoning effort for this session starting from the next turn.""" + reasoning_effort = reasoning_effort.strip() + if not reasoning_effort: + raise ValueError("reasoning_effort must not be empty") + context.state["reasoning_effort"] = reasoning_effort + await context.tape.append_event("reasoning_effort_switch", {"reasoning_effort": reasoning_effort}) + return f"Session reasoning effort set to {reasoning_effort} (applies from the next turn)." + + def _resolve_path(context: ToolContext, raw_path: str) -> Path: workspace = context.state.get("_runtime_workspace") path = Path(raw_path).expanduser() diff --git a/src/bub/tools.py b/src/bub/tools.py index 071c6e62..9b3bc01e 100644 --- a/src/bub/tools.py +++ b/src/bub/tools.py @@ -6,7 +6,7 @@ import inspect import json import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Iterable, Sequence from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Protocol, overload @@ -94,10 +94,22 @@ class Tool: description: str = "" parameters: dict[str, Any] = field(default_factory=dict) context: bool = False + agent_use: bool = True def run(self, *args: Any, **kwargs: Any) -> Any: return self.handler(*args, **kwargs) + def to_schema(self) -> dict[str, Any]: + """Build an any-llm completion tool payload.""" + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + @classmethod def from_callable( cls, @@ -106,6 +118,7 @@ def from_callable( name: str | None = None, description: str | None = None, context: bool = False, + agent_use: bool = True, ) -> Tool: signature = inspect.signature(func) if context and "context" not in signature.parameters: @@ -129,9 +142,15 @@ def validated(*args: Any, **kwargs: Any) -> Any: parameters=parameters, handler=validated, context=context, + agent_use=agent_use, ) +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] + + @dataclass(frozen=True) class ToolExecution: tool_results: list[Any] = field(default_factory=list) @@ -405,6 +424,7 @@ def tool( model: type[BaseModel] | None = ..., description: str | None = ..., context: bool = ..., + agent_use: bool = ..., ) -> Tool: ... @@ -416,6 +436,7 @@ def tool( model: type[BaseModel] | None = ..., description: str | None = ..., context: bool = ..., + agent_use: bool = ..., ) -> Callable[[Callable], Tool]: ... @@ -426,6 +447,7 @@ def tool( model: type[BaseModel] | None = None, description: str | None = None, context: bool = False, + agent_use: bool = True, ) -> Tool | Callable[[Callable], Tool]: """Decorator to convert a function into a Tool instance.""" @@ -447,9 +469,16 @@ def handler(*args: Any, **kwargs: Any) -> Any: parameters=model.model_json_schema(), handler=handler, context=context, + agent_use=agent_use, ) else: - result = Tool.from_callable(func, name=name, description=description, context=context) + result = Tool.from_callable( + func, + name=name, + description=description, + context=context, + agent_use=agent_use, + ) tool_instance = _add_logging(result) REGISTRY[tool_instance.name] = tool_instance return tool_instance diff --git a/tests/test_agent_hooks.py b/tests/test_agent_hooks.py index ceb448f6..a4a3ae6c 100644 --- a/tests/test_agent_hooks.py +++ b/tests/test_agent_hooks.py @@ -254,7 +254,7 @@ def _runner_and_tape(self, hooks: AgentHooks, captured: dict): from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, TapeContext class FakeRunner(ModelRunner): - async def completion_response(self, *, model, messages, tools, max_tokens=None): + async def completion_response(self, *, model, messages, tools, max_tokens=None, reasoning_effort=None): captured.update(model=model, max_tokens=max_tokens) async def chunks(): diff --git a/tests/test_builtin_agent.py b/tests/test_builtin_agent.py index 3c8d5572..7d8a01e0 100644 --- a/tests/test_builtin_agent.py +++ b/tests/test_builtin_agent.py @@ -361,6 +361,41 @@ def denied_agent_tool() -> str: assert "tests_denied_agent_tool" not in system_prompt +@pytest.mark.asyncio +async def test_agent_run_excludes_tools_disabled_for_agent_use() -> None: + visible_name = "tests.visible_agent_tool" + internal_name = "tests.internal_agent_tool" + REGISTRY.pop(visible_name, None) + REGISTRY.pop(internal_name, None) + + @tool(name=visible_name, description="Visible tool") + def visible_agent_tool() -> str: + return "visible" + + @tool(name=internal_name, description="Internal tool", agent_use=False) + def internal_agent_tool() -> str: + return "internal" + + agent = _make_agent() + fork_capture = _ForkCapture() + agent.tape = _FakeTapeFactory(fork_capture) # type: ignore[assignment] + + result = await agent.run_stream( + session_id="user/s1", + prompt="hello", + state={"_runtime_workspace": "/tmp"}, # noqa: S108 + allowed_tools=[visible_name, internal_name], + ) + [event async for event in result] + + completion_kwargs = _model_runner(agent).completion_kwargs + assert completion_kwargs is not None + assert [tool.name for tool in completion_kwargs["tools"]] == ["tests_visible_agent_tool"] + system_prompt = completion_kwargs["messages"][0]["content"] + assert "tests_visible_agent_tool" in system_prompt + assert "tests_internal_agent_tool" not in system_prompt + + @pytest.mark.asyncio async def test_agent_run_rejects_unknown_allowed_tools() -> None: agent = _make_agent() diff --git a/tests/test_builtin_hook_impl.py b/tests/test_builtin_hook_impl.py index 9efb58e6..7fb6bb7d 100644 --- a/tests/test_builtin_hook_impl.py +++ b/tests/test_builtin_hook_impl.py @@ -142,6 +142,19 @@ async def test_load_state_injects_model_recorded_on_session_tape(tmp_path: Path) assert state["model"] == "openai:gpt-4o" +@pytest.mark.asyncio +async def test_load_state_injects_reasoning_effort_recorded_on_session_tape(tmp_path: Path) -> None: + _, impl, agent = _build_impl(tmp_path) + session = agent.tape.session_tape("resolved-session", impl.framework.workspace) + await session.append_event("reasoning_effort_switch", {"reasoning_effort": "high"}) + + message = ChannelMessage(session_id="session", channel="cli", chat_id="room", content="hello") + + state = await impl.load_state(message=message, session_id="resolved-session") + + assert state["reasoning_effort"] == "high" + + @pytest.mark.asyncio async def test_load_state_does_not_inject_model_for_unknown_session(tmp_path: Path) -> None: """A session with nothing recorded on its tape must not inherit any model (no leakage).""" diff --git a/tests/test_builtin_model_runner.py b/tests/test_builtin_model_runner.py index 4dc1b4db..ca7be09f 100644 --- a/tests/test_builtin_model_runner.py +++ b/tests/test_builtin_model_runner.py @@ -138,6 +138,32 @@ async def test_anthropic_prompt_caching_is_requested() -> None: assert "stream_options" not in llm.completion_kwargs +@pytest.mark.asyncio +async def test_run_applies_reasoning_effort_from_tape_state(tmp_path: Path) -> None: + tape = Tape( + tmp_path, + AsyncTapeStoreAdapter(InMemoryTapeStore()), + TapeContext(state={"reasoning_effort": "high"}), + ).scoped("test-tape") + llm = _FakeStreamingOpenAIProvider() + runner = _FakeOpenAIModelRunner( + AgentSettings.model_construct( + model="openai:gpt-test", + max_tokens=100, + model_timeout_seconds=None, + completion_args={"reasoning_effort": "low"}, + ), + llm, + ) + + await tape.ensure_bootstrap_anchor() + events = runner.run(tape=tape, model="gpt-test", tools=[], system_prompt=None, prompt="hello") + [event async for event in events] + + assert llm.completion_kwargs is not None + assert llm.completion_kwargs["reasoning_effort"] == "high" + + @pytest.mark.asyncio async def test_completion_args_are_forwarded_without_overriding_managed_args() -> None: llm = _FakeStreamingOpenAIProvider() diff --git a/tests/test_builtin_tools.py b/tests/test_builtin_tools.py index 8e4100b9..873219b0 100644 --- a/tests/test_builtin_tools.py +++ b/tests/test_builtin_tools.py @@ -14,13 +14,12 @@ from bub.builtin.tools import ( bash, bash_output, - completion_tools, kill_bash, - model_tools, quit_tool, render_tools_prompt, resolve_tool_names, set_model, + set_reasoning_effort, tape_info, ) from bub.errors import ErrorKind @@ -58,47 +57,6 @@ async def test_tape_info_formats_token_cache_hit_rate(tmp_path) -> None: assert "last_token_cache_hit_rate: 37.50%" in result -def test_completion_tools_builds_any_llm_payload() -> None: - parameters = { - "type": "object", - "properties": {"value": {"type": "string"}}, - "required": ["value"], - } - sample_tool = Tool( - name="tests_sample_tool", - description="Sample tool", - parameters=parameters, - handler=lambda value: value, - ) - - assert completion_tools([sample_tool]) == [ - { - "type": "function", - "function": { - "name": "tests_sample_tool", - "description": "Sample tool", - "parameters": parameters, - }, - } - ] - - -def test_model_tools_rewrites_dotted_names_without_mutating_original() -> None: - tool_name = "tests.rename_me" - REGISTRY.pop(tool_name, None) - - @tool(name=tool_name, description="rename") - def rename_me(value: str) -> str: - return "ok" - - rewritten = model_tools([rename_me]) - - assert [item.name for item in rewritten] == ["tests_rename_me"] - assert rewritten[0].parameters == rename_me.parameters - assert rename_me.name == tool_name - assert "additionalProperties" not in rename_me.parameters - - def test_render_tools_prompt_renders_available_tools_block() -> None: first_name = "tests.prompt_one" second_name = "tests.prompt_two" @@ -135,6 +93,12 @@ def test_render_tools_prompt_returns_empty_string_for_empty_input() -> None: assert render_tools_prompt([]) == "" +def test_render_tools_prompt_excludes_tools_disabled_for_agent_use() -> None: + internal_tool = Tool(name="tests.internal", handler=lambda: None, agent_use=False) + + assert render_tools_prompt([internal_tool]) == "" + + def test_resolve_tool_names_accepts_runtime_names_and_model_aliases() -> None: dotted_name = "tests.resolve_alias" underscored_name = "tests_with_underscore" @@ -212,6 +176,39 @@ async def test_set_model_overwrites_previous_model(tmp_path) -> None: assert context.state["model"] == "anthropic:claude-3" +def test_set_reasoning_effort_is_registered_for_internal_use() -> None: + assert REGISTRY["reasoning_effort"] is set_reasoning_effort + assert set_reasoning_effort.context is True + assert set_reasoning_effort.agent_use is False + assert set_reasoning_effort.parameters == { + "type": "object", + "properties": {"reasoning_effort": {"type": "string"}}, + "required": ["reasoning_effort"], + } + + +@pytest.mark.asyncio +async def test_set_reasoning_effort_writes_state_and_records_on_tape(tmp_path) -> None: + context = _tool_context(tmp_path) + + result = await set_reasoning_effort.run(reasoning_effort=" high ", context=context) + + assert context.state["reasoning_effort"] == "high" + assert result == "Session reasoning effort set to high (applies from the next turn)." + entries = list(await context.tape.store.fetch_all(context.tape.query().kinds("event"))) + switches = [ + entry for entry in entries if entry.kind == "event" and entry.payload.get("name") == "reasoning_effort_switch" + ] + assert len(switches) == 1 + assert switches[0].payload.get("data") == {"reasoning_effort": "high"} + + +@pytest.mark.asyncio +async def test_set_reasoning_effort_rejects_empty_value(tmp_path) -> None: + with pytest.raises(ValueError, match="must not be empty"): + await set_reasoning_effort.run(reasoning_effort=" ", context=_tool_context(tmp_path)) + + @pytest.mark.asyncio async def test_bash_returns_stdout_for_foreground_command(tmp_path) -> None: result = await bash.run(cmd=_python_shell("print('hello')"), context=_tool_context(tmp_path)) diff --git a/tests/test_tools.py b/tests/test_tools.py index e1aa0c2c..21f38da5 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -6,13 +6,61 @@ from loguru import logger from pydantic import BaseModel -from bub.tools import REGISTRY, tool, tool_call_reporter +from bub.tools import REGISTRY, Tool, model_tools, tool, tool_call_reporter class EchoInput(BaseModel): value: str +def test_tool_builds_completion_payload() -> None: + parameters = { + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + } + sample_tool = Tool( + name="tests_sample_tool", + description="Sample tool", + parameters=parameters, + handler=lambda value: value, + ) + + assert sample_tool.to_schema() == { + "type": "function", + "function": { + "name": "tests_sample_tool", + "description": "Sample tool", + "parameters": parameters, + }, + } + + +def test_model_tools_rewrites_dotted_names_without_mutating_original() -> None: + tool_name = "tests.rename_me" + REGISTRY.pop(tool_name, None) + + @tool(name=tool_name, description="rename") + def rename_me(value: str) -> str: + return "ok" + + rewritten = model_tools([rename_me]) + + assert [item.name for item in rewritten] == ["tests_rename_me"] + assert rewritten[0].parameters == rename_me.parameters + assert rename_me.name == tool_name + assert "additionalProperties" not in rename_me.parameters + + +def test_model_tools_excludes_tools_disabled_for_agent_use() -> None: + visible_tool = Tool(name="tests.visible", handler=lambda: None) + internal_tool = Tool(name="tests.internal", handler=lambda: None, agent_use=False) + + rewritten = model_tools([visible_tool, internal_tool]) + + assert [item.name for item in rewritten] == ["tests_visible"] + + @pytest.mark.asyncio async def test_tool_decorator_registers_tool_and_preserves_metadata() -> None: tool_name = "tests.sync_tool" @@ -24,10 +72,25 @@ def sync_tool(payload: EchoInput) -> str: assert sync_tool.name == tool_name assert sync_tool.description == "Sync test tool" + assert sync_tool.agent_use is True assert REGISTRY[tool_name] is sync_tool assert await sync_tool.run(value="hello") == "HELLO" +@pytest.mark.asyncio +async def test_tool_decorator_can_disable_agent_use_without_disabling_direct_calls() -> None: + tool_name = "tests.internal_tool" + REGISTRY.pop(tool_name, None) + + @tool(name=tool_name, agent_use=False) + def internal_tool(value: str) -> str: + return value.upper() + + assert internal_tool.agent_use is False + assert REGISTRY[tool_name] is internal_tool + assert await internal_tool.run("hello") == "HELLO" + + @pytest.mark.asyncio async def test_tool_wrapper_logs_and_omits_context_from_log_payload(monkeypatch: pytest.MonkeyPatch) -> None: tool_name = "tests.async_tool"