diff --git a/src/bub/builtin/commands.py b/src/bub/builtin/commands.py new file mode 100644 index 00000000..132dd98a --- /dev/null +++ b/src/bub/builtin/commands.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.commands`.""" + +import sys + +from bub.core import commands as _module +from bub.core.commands import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/src/bub/builtin/context.py b/src/bub/builtin/context.py new file mode 100644 index 00000000..510f507c --- /dev/null +++ b/src/bub/builtin/context.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.context`.""" + +import sys + +from bub.core import context as _module +from bub.core.context import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/src/bub/core/agent.py b/src/bub/core/agent.py index 78e8502a..1b1bc5c2 100644 --- a/src/bub/core/agent.py +++ b/src/bub/core/agent.py @@ -73,6 +73,8 @@ class Agent: :func:`default_tape_context`, which replays the recorded chat. sidecars: Sidecar capabilities mounted on every tape. max_steps: Maximum model calls in one run. + max_tool_concurrency: Maximum tool calls from one model response that run at + once. None runs them all concurrently; 1 runs them one at a time, in order. name: Agent name, reported in ``invoke_agent`` traces. command_prefix: When set, a text prompt starting with it runs as a command (see :meth:`run_command`) instead of the agent loop. None disables commands. @@ -97,6 +99,7 @@ def __init__( context: TapeContext | None = None, sidecars: Sequence[TapeSidecar] = (), max_steps: int = sys.maxsize, + max_tool_concurrency: int | None = None, name: str = "agent", command_prefix: str | None = None, ) -> None: @@ -117,6 +120,9 @@ def __init__( sidecars=tuple(sidecars), ) self.max_steps = max_steps + if max_tool_concurrency is not None and max_tool_concurrency < 1: + raise ValueError("max_tool_concurrency must be a positive integer or None") + self.max_tool_concurrency = max_tool_concurrency self.name = name self.command_prefix = None if command_prefix is None else validate_command_prefix(command_prefix) self.events = EventBus() @@ -352,7 +358,7 @@ async def run_once( yield StreamEvent("tool_call", {"tool_calls": serialized_tool_calls}) turn_state["_runtime_tool_names"] = tuple(tool_map) context = ToolContext(tape=tape, run_id=run_id, state=turn_state, events=self.events) - execution = await ToolExecutor(events=self.events).execute_many( + execution = await ToolExecutor(events=self.events, max_concurrency=self.max_tool_concurrency).execute_many( tool_invocations, context=context, call_ids=[call.id for call in tool_calls], diff --git a/src/bub/core/tools.py b/src/bub/core/tools.py index 10fdbe9a..76eec3bf 100644 --- a/src/bub/core/tools.py +++ b/src/bub/core/tools.py @@ -5,6 +5,7 @@ import json import time from collections.abc import Callable, Iterable, Sequence +from contextlib import AbstractAsyncContextManager, nullcontext from dataclasses import dataclass, field, replace from typing import Any, Literal, get_args, get_type_hints, overload @@ -286,12 +287,20 @@ class ToolExecutor: With ``render`` (the default) successful results are rendered to model-facing text with ``Tool.render``; ``render=False`` returns tools' structured results unchanged, as needed when the caller is code rather than the model. ``ToolCallStarted`` and ``ToolCallFinished`` - events are emitted on ``events``. + events are emitted on ``events``. ``max_concurrency`` limits how many of this executor's + calls run at once; None leaves them unlimited. """ - def __init__(self, events: EventBus | None = None, *, render: bool = True) -> None: + def __init__( + self, events: EventBus | None = None, *, render: bool = True, max_concurrency: int | None = None + ) -> None: + if max_concurrency is not None and max_concurrency < 1: + raise ValueError("max_concurrency must be a positive integer or None") self._events = events self._render = render + self._limiter: AbstractAsyncContextManager[Any] = ( + nullcontext() if max_concurrency is None else asyncio.Semaphore(max_concurrency) + ) async def execute( self, @@ -347,33 +356,36 @@ async def _trace_tool_response( context: ToolContext | None, call_id: str | None, ) -> Any: - span = Span( - f"execute_tool {tool.name}", - { - "gen_ai.operation.name": "execute_tool", - "gen_ai.tool.name": tool.name, - "gen_ai.tool.type": "function", - "gen_ai.tool.call.id": call_id, - "bub.run_id": context.run_id if context else None, - "gen_ai.conversation.id": context.state.get("session_id") if context else None, - }, - ) - try: - with span.activate(): - result = await self._handle_tool_response_async(tool, arguments, context, span=span) - if isinstance(result, _FailedToolResult): - span.fail(result.error) - span.set(**{ - "gen_ai.tool.call.result": result.error.as_dict() if result.result is None else result.result - }) - else: - span.set(**{"gen_ai.tool.call.result": result}) - return result - except BaseException as exc: - span.fail(exc) - raise - finally: - span.end() + async with self._limiter: + span = Span( + f"execute_tool {tool.name}", + { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": tool.name, + "gen_ai.tool.type": "function", + "gen_ai.tool.call.id": call_id, + "bub.run_id": context.run_id if context else None, + "gen_ai.conversation.id": context.state.get("session_id") if context else None, + }, + ) + try: + with span.activate(): + result = await self._handle_tool_response_async(tool, arguments, context, span=span) + if isinstance(result, _FailedToolResult): + span.fail(result.error) + span.set(**{ + "gen_ai.tool.call.result": result.error.as_dict() + if result.result is None + else result.result + }) + else: + span.set(**{"gen_ai.tool.call.result": result}) + return result + except BaseException as exc: + span.fail(exc) + raise + finally: + span.end() def _invoke_tool( self, diff --git a/src/bub/hooks/interception.py b/src/bub/hooks/interception.py index bcc56c37..57588a36 100644 --- a/src/bub/hooks/interception.py +++ b/src/bub/hooks/interception.py @@ -27,12 +27,20 @@ "LLMCallDecision", "LLMCallRequest", "LLMCallResult", + "LlmCallDecision", + "LlmCallRequest", + "LlmCallResult", "PluginAgentHooks", "ToolCall", "ToolCallDecision", "ToolCallResult", ] +# Compatibility aliases for the names used before 0.6.0a2. +LlmCallDecision = LLMCallDecision +LlmCallRequest = LLMCallRequest +LlmCallResult = LLMCallResult + class PluginAgentHooks: """Fault-isolated adapter running the framework's pluggy interception hooks for agent events. diff --git a/src/bub/sidecars.py b/src/bub/sidecars.py new file mode 100644 index 00000000..88e59846 --- /dev/null +++ b/src/bub/sidecars.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.sidecars`.""" + +import sys + +from bub.core import sidecars as _module +from bub.core.sidecars import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/src/bub/store.py b/src/bub/store.py new file mode 100644 index 00000000..3fd8aeca --- /dev/null +++ b/src/bub/store.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.store`.""" + +import sys + +from bub.core import store as _module +from bub.core.store import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/src/bub/tape.py b/src/bub/tape.py new file mode 100644 index 00000000..965d1d18 --- /dev/null +++ b/src/bub/tape.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.tape`.""" + +import sys + +from bub.core import tape as _module +from bub.core.tape import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/src/bub/tools.py b/src/bub/tools.py new file mode 100644 index 00000000..76d4ee66 --- /dev/null +++ b/src/bub/tools.py @@ -0,0 +1,8 @@ +"""Compatibility alias for :mod:`bub.core.tools`.""" + +import sys + +from bub.core import tools as _module +from bub.core.tools import * # noqa: F403 + +sys.modules[__name__] = _module diff --git a/tests/test_compat_aliases.py b/tests/test_compat_aliases.py new file mode 100644 index 00000000..14b535a5 --- /dev/null +++ b/tests/test_compat_aliases.py @@ -0,0 +1,39 @@ +import importlib +import subprocess +import sys + +import pytest + +from bub.core import events +from bub.hooks import interception + +MOVED_MODULES = { + "bub.tape": "bub.core.tape", + "bub.tools": "bub.core.tools", + "bub.store": "bub.core.store", + "bub.sidecars": "bub.core.sidecars", + "bub.builtin.commands": "bub.core.commands", + "bub.builtin.context": "bub.core.context", +} + + +@pytest.mark.parametrize(("old", "new"), MOVED_MODULES.items()) +def test_moved_module_alias_is_the_new_module(old: str, new: str) -> None: + assert importlib.import_module(old) is importlib.import_module(new) + + +def test_top_level_module_aliases_load_no_app_dependencies() -> None: + script = ( + "import sys, bub.tape, bub.tools, bub.store, bub.sidecars\n" + "loaded = {name.split('.')[0] for name in sys.modules}\n" + "print(sorted(loaded & {'typer', 'pluggy', 'pydantic_settings', 'yaml', 'rapidfuzz'}))\n" + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + + assert result.stdout.splitlines() == ["[]"] + + +@pytest.mark.parametrize("suffix", ["Request", "Result", "Decision"]) +def test_interception_keeps_llm_call_payload_aliases(suffix: str) -> None: + assert getattr(interception, f"LlmCall{suffix}") is getattr(events, f"LLMCall{suffix}") + assert f"LlmCall{suffix}" in interception.__all__ diff --git a/tests/test_core_agent.py b/tests/test_core_agent.py index e2e22327..e6448325 100644 --- a/tests/test_core_agent.py +++ b/tests/test_core_agent.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import asyncio from dataclasses import replace from pathlib import Path from typing import Any @@ -181,6 +182,48 @@ async def test_core_agent_stops_at_max_steps(provider_service: ProviderService) assert statuses == ["model=openai:test", "tool=ok", "continue", "run=error:1:max_steps_reached=1"] +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "expected_peak"), [(None, 3), (2, 2), (1, 1)]) +async def test_max_tool_concurrency_limits_tool_calls_of_one_response( + provider_service: ProviderService, limit: int | None, expected_peak: int +) -> None: + calls = [ + {"index": i, "id": f"call-{i}", "function": {"name": "wait", "arguments": f'{{"n": {i}}}'}} for i in range(3) + ] + provider_service.reply(sse(tool_events(calls))) + provider_service.reply(sse(chat_events("done"))) + running = 0 + peak = 0 + started: list[int] = [] + + async def wait(n: int) -> str: + nonlocal running, peak + running += 1 + peak = max(peak, running) + await asyncio.sleep(0.01) + running -= 1 + return "ok" + + async with provider_service.client() as client: + agent = Agent( + model=_model(client), + tools=[Tool.from_callable(wait, name="wait")], + max_tool_concurrency=limit, + ) + agent.events.on(ToolCallStarted, lambda event: started.append(event.call.arguments["n"])) + async with agent.run(["go"], tape="concurrency") as stream: + [event async for event in stream] + + assert peak == expected_peak + assert started == [0, 1, 2] + + +@pytest.mark.parametrize("limit", [0, -1]) +def test_max_tool_concurrency_must_be_positive(limit: int) -> None: + with pytest.raises(ValueError, match="max_tool_concurrency"): + Agent(max_tool_concurrency=limit) + + @pytest.mark.asyncio async def test_redirected_call_reuses_the_provider_of_a_republic_model(provider_service: ProviderService) -> None: provider_service.reply(sse(chat_events())) diff --git a/website/src/content/docs/docs/build/sdk.md b/website/src/content/docs/docs/build/sdk.md index 1e2fc35d..f8b7c614 100644 --- a/website/src/content/docs/docs/build/sdk.md +++ b/website/src/content/docs/docs/build/sdk.md @@ -88,6 +88,7 @@ continues the conversation stored under `sessions/`. | `context` | A `TapeContext` choosing which tape entries become model context. Defaults to `default_tape_context()`. | | `sidecars` | `TapeSidecar` capabilities mounted on every tape. | | `max_steps` | Maximum model calls in one run. | +| `max_tool_concurrency` | Maximum tool calls from one model response that run at once. `None` (the default) runs them all concurrently; `1` runs them one at a time, in order. | | `name` | Agent name reported in `invoke_agent` traces. Defaults to `"agent"`. | | `command_prefix` | Prefix marking command prompts, such as `","`; see [Commands](#commands). `None` (the default) disables commands. | diff --git a/website/src/content/docs/zh-cn/docs/build/sdk.md b/website/src/content/docs/zh-cn/docs/build/sdk.md index 3097e792..8b3fdbb9 100644 --- a/website/src/content/docs/zh-cn/docs/build/sdk.md +++ b/website/src/content/docs/zh-cn/docs/build/sdk.md @@ -86,6 +86,7 @@ if __name__ == "__main__": | `context` | 决定哪些 tape 条目进入模型上下文的 `TapeContext`。默认 `default_tape_context()`。 | | `sidecars` | 挂载到每条 tape 上的 `TapeSidecar` 能力。 | | `max_steps` | 一次运行中最多的模型调用次数。 | +| `max_tool_concurrency` | 同一次模型回复中可同时执行的工具调用数上限。默认 `None` 即全部并发；`1` 则按顺序逐个执行。 | | `name` | 在 `invoke_agent` trace 中上报的 agent 名称。默认 `"agent"`。 | | `command_prefix` | 标记命令 prompt 的前缀，例如 `","`；见 [命令](#命令)。默认 `None`，即不启用命令。 |