diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 1fa78d1f..ae6dda1a 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -290,7 +290,7 @@ async def _stream_events_with_auto_handoff( ) return - next_prompt = await self.framework.continue_prompt(tape=tape, state=state) + next_prompt = await self.framework.continue_prompt(prompt=next_prompt, tape=tape, state=state) await tape.append_event( "loop.step", { diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index bd6fc4f9..336c3d44 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -219,8 +219,8 @@ async def run_model_stream(self, prompt: str | list[dict], session_id: str, stat ) @hookimpl - def continue_prompt(self, tape: Tape, state: StreamState) -> str: - del state + def continue_prompt(self, prompt: str | list[dict], tape: Tape, state: StreamState) -> str: + del prompt, state if "context" in tape.context.state: return f"{DEFAULT_CONTINUE_PROMPT} [context: {tape.context.state['context']}]" return DEFAULT_CONTINUE_PROMPT diff --git a/src/bub/framework.py b/src/bub/framework.py index 5d94660c..4227f266 100644 --- a/src/bub/framework.py +++ b/src/bub/framework.py @@ -124,11 +124,11 @@ async def build_prompt( prompt = content_of(message) return cast("str | list[dict[str, Any]]", prompt) - async def continue_prompt(self, tape: Tape, state: StreamState) -> str: + async def continue_prompt(self, prompt: str | list[dict], tape: Tape, state: StreamState) -> str: """Build the prompt for the next step of an agent loop.""" - prompt = await self._hook_runtime.call_first("continue_prompt", tape=tape, state=state) - if isinstance(prompt, str): - return prompt + next_prompt = await self._hook_runtime.call_first("continue_prompt", prompt=prompt, tape=tape, state=state) + if isinstance(next_prompt, str): + return next_prompt raise TypeError("hook.continue_prompt must return str") async def build_state(self, message: Envelope, session_id: str) -> TurnState: diff --git a/src/bub/hooks/specs.py b/src/bub/hooks/specs.py index e85e7e67..d5ad5ba7 100644 --- a/src/bub/hooks/specs.py +++ b/src/bub/hooks/specs.py @@ -66,7 +66,7 @@ def run_model_stream(self, prompt: str | list[dict], session_id: str, state: Tur raise NotImplementedError @hookspec(firstresult=True) - def continue_prompt(self, tape: Tape, state: StreamState) -> str: + def continue_prompt(self, prompt: str | list[dict], tape: Tape, state: StreamState) -> str: """Build the prompt used to continue an agent loop. Implementations may be synchronous or asynchronous. The first diff --git a/tests/test_builtin_agent.py b/tests/test_builtin_agent.py index f08e1f50..73b3d1b3 100644 --- a/tests/test_builtin_agent.py +++ b/tests/test_builtin_agent.py @@ -257,6 +257,7 @@ async def test_agent_loop_awaits_continue_prompt_hook_with_stream_state() -> Non agent = _make_agent() tape = _FakeTape(_ForkCapture()) prompts: list[str | list[dict]] = [] + continuation_prompts: list[str | list[dict]] = [] observed_usage: list[dict[str, Any] | None] = [] async def run_once(**kwargs: Any) -> AsyncStreamEvents: @@ -268,7 +269,8 @@ async def iterator() -> AsyncIterator[StreamEvent]: return AsyncStreamEvents(iterator(), state=StreamState(usage={"step": len(prompts)})) - async def continue_prompt(*, tape: _FakeTape, state: StreamState) -> str: + async def continue_prompt(*, prompt: str | list[dict], tape: _FakeTape, state: StreamState) -> str: + continuation_prompts.append(prompt) observed_usage.append(state.usage) return "custom continuation" @@ -286,6 +288,7 @@ async def continue_prompt(*, tape: _FakeTape, state: StreamState) -> str: assert [event.kind for event in events] == ["final", "final"] assert prompts == ["initial prompt", "custom continuation"] + assert continuation_prompts == ["initial prompt"] assert observed_usage == [{"step": 1}] diff --git a/tests/test_builtin_hook_impl.py b/tests/test_builtin_hook_impl.py index f026ee5c..12c639ad 100644 --- a/tests/test_builtin_hook_impl.py +++ b/tests/test_builtin_hook_impl.py @@ -96,7 +96,7 @@ def test_continue_prompt_includes_tape_context(tmp_path: Path) -> None: _, impl, _ = _build_impl(tmp_path) tape = _fake_tape(tmp_path).with_context(TapeContext(state={"context": "telegram metadata"})) - prompt = impl.continue_prompt(tape=tape, state=StreamState()) + prompt = impl.continue_prompt(prompt="current prompt", tape=tape, state=StreamState()) assert prompt == f"{DEFAULT_CONTINUE_PROMPT} [context: telegram metadata]" diff --git a/tests/test_framework.py b/tests/test_framework.py index 2bb00601..42f048f6 100644 --- a/tests/test_framework.py +++ b/tests/test_framework.py @@ -127,21 +127,22 @@ async def test_continue_prompt_awaits_high_priority_async_hook() -> None: class SyncPlugin: @hookimpl - def continue_prompt(self, tape, state): + def continue_prompt(self, prompt, tape, state): called.append("sync") return "sync prompt" class AsyncPlugin: @hookimpl - async def continue_prompt(self, tape: Any, state: StreamState) -> str: + async def continue_prompt(self, prompt: str, tape: Any, state: StreamState) -> str: called.append("async") + assert prompt == "current prompt" assert state.usage == {"total_tokens": 42} return "async prompt" framework._plugin_manager.register(SyncPlugin(), name="sync") framework._plugin_manager.register(AsyncPlugin(), name="async") - prompt = await framework.continue_prompt(tape=tape, state=state) + prompt = await framework.continue_prompt(prompt="current prompt", tape=tape, state=state) assert prompt == "async prompt" assert called == ["async"] diff --git a/website/src/content/docs/docs/reference/hooks.mdx b/website/src/content/docs/docs/reference/hooks.mdx index 781398a2..fc4f9877 100644 --- a/website/src/content/docs/docs/reference/hooks.mdx +++ b/website/src/content/docs/docs/reference/hooks.mdx @@ -20,7 +20,7 @@ For the *why* and *how* of each stage see [Turn pipeline](/docs/concepts/turn-pi | `build_prompt` | firstresult | `(message: Envelope, session_id: str, state: TurnState) -> str \| list[dict]` | prompt or content parts | `BubFramework.process_inbound` | Falls back to `content_of(message)` when no impl returns a non-`None` value, or when the selected firstresult is falsy. A falsy selected value does not cause lower-priority impls to run. | | `run_model` | firstresult | `(prompt, session_id, state) -> str` | model text | `HookRuntime.run_model` | Legacy text path. Implement either `run_model` or `run_model_stream`, not both. | | `run_model_stream` | firstresult | `(prompt, session_id, state) -> AsyncStreamEvents` | async stream | `HookRuntime.run_model_stream` | Preferred. Falls back to wrapping a `run_model` result in a one-chunk stream when no streaming impl exists. | -| `continue_prompt` | firstresult | `(tape: Tape, state: StreamState) -> str` | next-step prompt | `Agent._stream_events_with_auto_handoff` | Runs before each continued agent-loop step. Implementations may be sync or async; the builtin preserves the default prompt and optional tape context. | +| `continue_prompt` | firstresult | `(prompt: str \| list[dict], tape: Tape, state: StreamState) -> str` | next-step prompt | `Agent._stream_events_with_auto_handoff` | Runs before each continued agent-loop step. `prompt` is the prompt used by the completed step. Implementations may be sync or async; the builtin preserves the default prompt and optional tape context. | | `save_state` | broadcast | `(session_id, state, message, model_output) -> None` | none | `process_inbound` model-stage finally block | Runs after prompt resolution for model-stage success or failure; failures before prompt/model execution skip it. | | `render_outbound` | broadcast | `(message, session_id, state, model_output) -> list[Envelope]` | outbound batch | `BubFramework._collect_outbounds` | All batches are concatenated via `unpack_batch`. Empty results trigger a default echo envelope. | | `dispatch_outbound` | broadcast | `(message: Envelope) -> bool` | sent flag | `process_inbound` per outbound | Each outbound is fanned out to every impl. | diff --git a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx index 52e84ba3..51685fe2 100644 --- a/website/src/content/docs/zh-cn/docs/reference/hooks.mdx +++ b/website/src/content/docs/zh-cn/docs/reference/hooks.mdx @@ -20,7 +20,7 @@ description: BubHookSpecs 中每个钩子的类型、签名、返回值与调用 | `build_prompt` | firstresult | `(message: Envelope, session_id: str, state: TurnState) -> str \| list[dict]` | prompt or content parts | `BubFramework.process_inbound` | 没有实现返回非 `None` 值，或选中的 firstresult 为 falsy 时，回退到 `content_of(message)`。已选中的 falsy 值不会触发继续尝试低优先级实现。 | | `run_model` | firstresult | `(prompt, session_id, state) -> str` | model text | `HookRuntime.run_model` | 旧版纯文本路径。`run_model` 与 `run_model_stream` 选其一实现,不要同时实现。 | | `run_model_stream` | firstresult | `(prompt, session_id, state) -> AsyncStreamEvents` | async stream | `HookRuntime.run_model_stream` | 推荐路径。无流式实现时框架将 `run_model` 结果包装为单个 chunk 的流。 | -| `continue_prompt` | firstresult | `(tape: Tape, state: StreamState) -> str` | 下一步 prompt | `Agent._stream_events_with_auto_handoff` | 每次 agent loop 继续前调用。实现可同步或异步；builtin 保留默认 prompt 与可选 tape context。 | +| `continue_prompt` | firstresult | `(prompt: str \| list[dict], tape: Tape, state: StreamState) -> str` | 下一步 prompt | `Agent._stream_events_with_auto_handoff` | 每次 agent loop 继续前调用。`prompt` 是刚完成 step 使用的 prompt。实现可同步或异步；builtin 保留默认 prompt 与可选 tape context。 | | `save_state` | broadcast | `(session_id, state, message, model_output) -> None` | none | `process_inbound` model-stage finally block | prompt 解析完成后，在模型阶段成功或失败时执行；若失败发生在 prompt/model 执行前则跳过。 | | `render_outbound` | broadcast | `(message, session_id, state, model_output) -> list[Envelope]` | outbound batch | `BubFramework._collect_outbounds` | 所有批次通过 `unpack_batch` 拼接;结果为空时使用默认回声 envelope。 | | `dispatch_outbound` | broadcast | `(message: Envelope) -> bool` | sent flag | `process_inbound` per outbound | 每个 outbound 都会广播给所有实现。 |