feat(core): implement Prompt Assembly system for structured prompt composition#7172
feat(core): implement Prompt Assembly system for structured prompt composition#7172whatevertogo wants to merge 6 commits intoAstrBotDevs:masterfrom
Conversation
…dering - Add models for Prompt Assembly including SystemBlock, UserAppendPart, ContextContribution, and PromptAssembly. - Create a renderer for assembling prompts into ProviderRequest, ensuring proper order and rendering logic. - Introduce tracing functionality to capture structured debug information during prompt assembly. - Implement event registration for plugins to modify prompt assembly before rendering. - Enhance unit tests to cover new prompt assembly features, including ordering and mutation handling.
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The
build_prompt_trace_snapshotdocstring says the structured trace only includes core-owned blocks and excludes plugin hook modifications, butbuild_main_agentcurrently calls it after runningOnPromptAssemblyEventon the samePromptAssembly; consider either taking a pre-hook snapshot or updating the docstring to match the actual behavior. - In
register_on_prompt_assembly's docstring example, the decorator is written as@on_prompt_assembly()but the actual exported decorator isregister_on_prompt_assembly; updating the example to use the correct name will avoid confusion for plugin authors.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `build_prompt_trace_snapshot` docstring says the structured trace only includes core-owned blocks and excludes plugin hook modifications, but `build_main_agent` currently calls it after running `OnPromptAssemblyEvent` on the same `PromptAssembly`; consider either taking a pre-hook snapshot or updating the docstring to match the actual behavior.
- In `register_on_prompt_assembly`'s docstring example, the decorator is written as `@on_prompt_assembly()` but the actual exported decorator is `register_on_prompt_assembly`; updating the example to use the correct name will avoid confusion for plugin authors.
## Individual Comments
### Comment 1
<location path="tests/unit/test_astr_main_agent.py" line_range="1280-1289" />
<code_context>
+ async def test_build_main_agent_prompt_assembly_hook_can_add_blocks(
</code_context>
<issue_to_address>
**suggestion (testing):** Also test the case where the OnPromptAssembly hook cancels the request by returning True.
There's another important branch in `build_main_agent` where `call_event_hook(..., OnPromptAssemblyEvent, ...)` returns `True` and triggers an early `return None`. Please add a sibling test where the patched `call_event_hook` returns `True` for `OnPromptAssemblyEvent`, and assert that `build_main_agent` returns `None` and no rendering or trace recording occurs. This will cover the new cancellation behavior and help avoid regressions.
Suggested implementation:
```python
async def test_build_main_agent_prompt_assembly_hook_can_add_blocks(
self, mock_event, mock_context, mock_provider
):
module = ama
mock_event.platform_meta.support_proactive_message = False
mock_context.get_provider_by_id.return_value = None
mock_context.get_using_provider.return_value = mock_provider
mock_context.get_config.return_value = {}
_setup_conversation_for_build(mock_context.conversation_manager)
tool_manager = MagicMock()
tool_manager.get_full_tool_set.return_value = ToolSet()
@pytest.mark.asyncio
async def test_build_main_agent_prompt_assembly_hook_can_cancel_request(
self, mock_event, mock_context, mock_provider, monkeypatch
):
module = ama
mock_event.platform_meta.support_proactive_message = False
mock_context.get_provider_by_id.return_value = None
mock_context.get_using_provider.return_value = mock_provider
mock_context.get_config.return_value = {}
_setup_conversation_for_build(mock_context.conversation_manager)
tool_manager = MagicMock()
tool_manager.get_full_tool_set.return_value = ToolSet()
# Patch call_event_hook so that OnPromptAssemblyEvent cancels the request
async def fake_call_event_hook(*args, **kwargs):
# Depending on how call_event_hook is called, the event class/type
# will be in either args or kwargs; we check both.
event_cls = None
if len(args) >= 2:
event_cls = args[1]
event_cls = kwargs.get("event_cls", event_cls)
if event_cls is getattr(module, "OnPromptAssemblyEvent"):
# Returning True should trigger the early return in build_main_agent
return True
return False
monkeypatch.setattr(module, "call_event_hook", fake_call_event_hook)
# Track any potential rendering / provider calls so we can assert they don't happen
mock_provider.create_chat_completion = MagicMock()
mock_provider.__call__ = MagicMock()
# Act: build_main_agent should see the cancellation and return None
result = await module.build_main_agent(
event=mock_event,
context=mock_context,
tool_manager=tool_manager,
)
assert result is None
assert not mock_provider.create_chat_completion.called
assert not mock_provider.__call__.called
```
To fully align this test with your codebase and the reviewer’s intent, you should:
1. Adjust the `build_main_agent` call to match the actual function signature (e.g. additional parameters like `telemetry`, `settings`, etc., if required in other tests in this file).
2. Replace the `mock_provider.create_chat_completion` / `mock_provider.__call__` assertions with assertions against the *actual* rendering call(s) used in `build_main_agent` (for example, if you have a helper like `render_blocks_to_provider_request` or a specific provider method).
3. If your tracing/telemetry logic is implemented via a function or object (e.g. `module.record_trace`, `module.telemetry.record`, etc.), patch that in this test similarly to `call_event_hook` and assert that it was **not** called when the hook returns `True`.
4. If `call_event_hook` in your codebase uses a different positional/keyword argument layout for the event class/type, update `fake_call_event_hook` to inspect the correct argument position or keyword (you can copy the pattern from other tests that patch `call_event_hook` in this file).
</issue_to_address>
### Comment 2
<location path="astrbot/core/astr_main_agent.py" line_range="180" />
<code_context>
reset_coro: Coroutine | None = None
+def _resolve_prompt_assembly(
+ assembly: PromptAssembly | None,
+) -> tuple[PromptAssembly, bool]:
</code_context>
<issue_to_address>
**issue (complexity):** Consider simplifying the PromptAssembly lifecycle by passing a single shared assembly through internal helpers and rendering only once in build_main_agent, instead of resolving/finalizing it in each helper.
The new `PromptAssembly` wiring is adding avoidable complexity in this module.
You already construct a single `PromptAssembly` in `build_main_agent` and render once at the end, so the `assembly: PromptAssembly | None = None` + `_resolve_prompt_assembly` + `_finalize_prompt_assembly` pattern in all helpers is unnecessary here and makes control flow harder to follow (extra params, early-return finalize calls, “should_render” state).
You can simplify while keeping all behavior by:
1. Treating `PromptAssembly` as required for these internal helpers.
2. Removing `_resolve_prompt_assembly` / `_finalize_prompt_assembly`.
3. Keeping the single `render_prompt_assembly` call in `build_main_agent`.
All affected functions in this file are internal (prefixed with `_`), so changing their signatures doesn’t affect external API.
### 1. Remove `_resolve_prompt_assembly` / `_finalize_prompt_assembly`
```py
- def _resolve_prompt_assembly(
- assembly: PromptAssembly | None,
- ) -> tuple[PromptAssembly, bool]:
- if assembly is None:
- return PromptAssembly(), True
- return assembly, False
-
-
- def _finalize_prompt_assembly(
- req: ProviderRequest,
- assembly: PromptAssembly,
- should_render: bool,
- ) -> None:
- if should_render:
- render_prompt_assembly(req, assembly)
```
### 2. Make `PromptAssembly` required in helpers
Example with `_apply_kb` and `_apply_file_extract` (same pattern applies to all others):
```py
- async def _apply_kb(
- event: AstrMessageEvent,
- req: ProviderRequest,
- plugin_context: Context,
- config: MainAgentBuildConfig,
- assembly: PromptAssembly | None = None,
- ) -> None:
- prompt_assembly, should_render = _resolve_prompt_assembly(assembly)
+ async def _apply_kb(
+ event: AstrMessageEvent,
+ req: ProviderRequest,
+ plugin_context: Context,
+ config: MainAgentBuildConfig,
+ assembly: PromptAssembly,
+ ) -> None:
if not config.kb_agentic_mode:
if req.prompt is None:
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
return
try:
kb_result = await retrieve_knowledge_base(
query=req.prompt,
umo=event.unified_msg_origin,
context=plugin_context,
)
if not kb_result:
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
return
- add_system_block(
- prompt_assembly,
+ add_system_block(
+ assembly,
source="knowledge_base",
order=SYSTEM_BLOCK_ORDER_KB,
content=f"\n\n[Related Knowledge Base Results]:\n{kb_result}",
)
except Exception as exc: # noqa: BLE001
logger.error("Error occurred while retrieving knowledge base: %s", exc)
else:
...
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
```
```py
- async def _apply_file_extract(
- event: AstrMessageEvent,
- req: ProviderRequest,
- config: MainAgentBuildConfig,
- assembly: PromptAssembly | None = None,
- ) -> None:
- prompt_assembly, should_render = _resolve_prompt_assembly(assembly)
+ async def _apply_file_extract(
+ event: AstrMessageEvent,
+ req: ProviderRequest,
+ config: MainAgentBuildConfig,
+ assembly: PromptAssembly,
+ ) -> None:
...
if not file_paths:
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
return
...
for file_content, file_name in zip(file_contents, file_names):
- add_context_suffix(
- prompt_assembly,
+ add_context_suffix(
+ assembly,
source="file_extract",
order=CONTEXT_ORDER_FILE_EXTRACT,
messages=[{...}],
)
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
```
Apply the same simplification to:
- `_apply_local_env_tools`
- `_ensure_persona_and_skills`
- `_ensure_img_caption`
- `_process_quote_message`
- `_append_system_reminders`
- `_decorate_llm_request`
- `_apply_llm_safety_mode`
- `_apply_sandbox_tools`
i.e. remove `PromptAssembly | None`, remove `should_render`, and remove all `_finalize_prompt_assembly` calls in normal/early-return/exception paths. Just mutate `req` directly for non-prompt fields, and use `assembly` for prompt-related blocks.
Example for `_decorate_llm_request`:
```py
- async def _decorate_llm_request(..., assembly: PromptAssembly | None = None) -> None:
- prompt_assembly, should_render = _resolve_prompt_assembly(assembly)
+ async def _decorate_llm_request(..., assembly: PromptAssembly) -> None:
...
if req.conversation:
- await _ensure_persona_and_skills(req, cfg, plugin_context, event, prompt_assembly)
+ await _ensure_persona_and_skills(req, cfg, plugin_context, event, assembly)
...
if img_cap_prov_id and req.image_urls:
- await _ensure_img_caption(..., prompt_assembly)
+ await _ensure_img_caption(..., assembly)
...
- await _process_quote_message(..., prompt_assembly)
+ await _process_quote_message(..., assembly)
...
- _append_system_reminders(event, req, cfg, tz, prompt_assembly)
- _finalize_prompt_assembly(req, prompt_assembly, should_render)
+ _append_system_reminders(event, req, cfg, tz, assembly)
```
### 3. Pass the shared `PromptAssembly` from `build_main_agent`
You already create a single `PromptAssembly` and render once at the end; just pass this instance everywhere:
```py
- prompt_assembly = PromptAssembly()
+ prompt_assembly = PromptAssembly()
...
if config.file_extract_enabled:
try:
- await _apply_file_extract(event, req, config, prompt_assembly)
+ await _apply_file_extract(event, req, config, prompt_assembly)
except Exception as exc:
...
...
- await _decorate_llm_request(event, req, plugin_context, config, prompt_assembly)
- await _apply_kb(event, req, plugin_context, config, prompt_assembly)
+ await _decorate_llm_request(event, req, plugin_context, config, prompt_assembly)
+ await _apply_kb(event, req, plugin_context, config, prompt_assembly)
...
if config.llm_safety_mode:
- _apply_llm_safety_mode(config, req, prompt_assembly)
+ _apply_llm_safety_mode(config, req, prompt_assembly)
if config.computer_use_runtime == "sandbox":
- _apply_sandbox_tools(config, req, req.session_id, prompt_assembly)
+ _apply_sandbox_tools(config, req, req.session_id, prompt_assembly)
elif config.computer_use_runtime == "local":
- _apply_local_env_tools(req, prompt_assembly)
+ _apply_local_env_tools(req, prompt_assembly)
```
Leave the final lifecycle as-is and fully centralized:
```py
prompt_assembly.metadata["base_request"] = summarize_provider_request_base(req)
if await call_event_hook(..., PromptMutation(prompt_assembly)):
return None
render_prompt_assembly(req, prompt_assembly)
try:
event.trace.record("core_prompt_assembly", **build_prompt_trace_snapshot(prompt_assembly))
except Exception:
logger.debug("Failed to record core_prompt_assembly trace", exc_info=True)
_sanitize_context_by_modalities(config, provider, req)
```
This keeps all functionality (PromptAssembly ordering, hooks, tracing, etc.) but removes the distributed “resolve/finalize” lifecycle logic, simplifies helper signatures, and makes it clear that helpers only **register** into a shared assembly while `build_main_agent` is solely responsible for rendering.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
Pull request overview
Implements a new “Prompt Assembly” pipeline in AstrBot core to structure prompt composition across system/user/context channels, enabling safer plugin prompt injection via a dedicated on_prompt_assembly hook and improving traceability of prompt construction.
Changes:
- Added
astrbot/core/prompt/module (models, registration/facade, renderer, tracing) for structured prompt composition. - Refactored
astrbot/core/astr_main_agent.pyto register prompt blocks intoPromptAssembly, run a plugin hook, then render intoProviderRequestand record a trace snapshot. - Extended the plugin event system with
EventType.OnPromptAssemblyEventand exposedon_prompt_assemblyfor plugin authors; updated dashboard labeling and added unit tests.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
astrbot/core/prompt/models.py |
Defines the 3-channel data model and ordering constants for prompt assembly. |
astrbot/core/prompt/assembly.py |
Adds registration helpers and the PromptMutation facade for plugin hooks. |
astrbot/core/prompt/renderer.py |
Renders a PromptAssembly into ProviderRequest fields (idempotent). |
astrbot/core/prompt/tracing.py |
Builds a structured trace snapshot for prompt assembly debugging/visibility controls. |
astrbot/core/prompt/__init__.py |
Exposes prompt assembly public API exports. |
astrbot/core/astr_main_agent.py |
Migrates inline prompt concatenation to the new prompt assembly lifecycle and hook. |
astrbot/core/star/star_handler.py |
Adds EventType.OnPromptAssemblyEvent and typing overloads for handler lookup. |
astrbot/core/star/register/star_handler.py |
Adds register_on_prompt_assembly decorator for plugin hook registration. |
astrbot/core/star/register/__init__.py |
Re-exports the new registration decorator. |
astrbot/api/event/filter/__init__.py |
Exposes on_prompt_assembly to plugin API consumers. |
astrbot/dashboard/routes/plugin.py |
Adds dashboard UI translation for the new event type. |
tests/unit/test_prompt_renderer.py |
New unit tests covering ordering/idempotency/trace snapshot behavior and mutation warnings. |
tests/unit/test_internal_agent_trace.py |
New focused test ensuring internal agent trace preserves string system_prompt. |
tests/unit/test_astr_main_agent.py |
Updates and adds tests to validate prompt assembly ordering and plugin hook behavior. |
There was a problem hiding this comment.
Code Review
This pull request introduces a structured 'Prompt Assembly' mechanism to manage LLM prompt construction across three channels: system blocks, user message parts, and context history. It refactors the main agent logic to use this assembly system, enabling better ordering, tracing, and plugin-based mutation via the new on_prompt_assembly hook. The review feedback suggests improving the predictability of system prompt block ordering by assigning unique order values to individual blocks within the same category, such as skills and runtime notices.
astrbot/core/astr_main_agent.py: - 删除 _resolve_prompt_assembly 和 _finalize_prompt_assembly 两个辅助函数, 消除内部隐式的 assembly 创建和渲染逻辑,将职责上移到调用方 - 将所有内部函数的 assembly 参数从 Optional[PromptAssembly] 改为 PromptAssembly(必填), 使数据流更透明、类型更严格 - _process_quote_message 的 assembly 参数移至 config 之前,保持位置参数顺序一致 tests/unit/test_astr_main_agent.py: - 所有测试用例显式创建 PromptAssembly 实例并传入, 需要验证渲染结果的场景额外调用 _render_assembly tests/test_profile_aware_tools.py: - 适配 _apply_sandbox_tools 新签名,补传 PromptAssembly 参数
Motivation / 动机
原先
astr_main_agent.py中所有 prompt 组装逻辑(system prompt 拼接、context 前后插入、用户消息追加)全部内联在一个巨大的函数里,流程散乱且难以扩展。插件开发者想要在 prompt 中注入自定义内容时,只能通过修改ProviderRequest的原始字段,缺乏结构化支持。本 PR 将 prompt 组装抽象为三通道(system / user / context)结构化模型,提供注册 → 渲染 → 追踪的完整流水线,并为插件暴露了
on_prompt_assemblyhook。Fix #7028
#7028:部分修复。现在 core 的 prompt 组装已经收口到 PromptAssembly,并且有新的结构化扩展点 @on_prompt_assembly(),见 assembly.py 和 star_handler.py。但它没有彻底禁止旧式插件在 on_llm_request 里直接改最终 ProviderRequest,所以“插件相互干扰”并没有被制度上完全消灭,只是 core 自己先干净了很多。
Fix #4446:基本修了。现在插件有两个明确入口:
结构化插入,用 @on_prompt_assembly()
看最终渲染结果,用 @on_llm_request()
也就是说,“插件拿不到最终版 prompt”这个问题,在 internal main agent 这条链路上已经被解决了。触发点在 astr_main_agent.py 里 render 前和 render 后两段链路之间。
#6423:只修了后端一半。现在 backend 已经有 core_prompt_assembly 结构化 trace,能按 source/order 看 prompt 来源,见 tracing.py。但 WebUI 还没有专门把这个 trace 渲染成模块化视图,所以“前端看起来还是一整坨”这件事没有完全解决。
Modifications / 改动点
新增模块
astrbot/core/prompt/:models.py— 三通道数据模型(SystemBlock/UserAppendPart/ContextContribution)及排序常量assembly.py— 注册 API(add_system_block/add_user_text/add_context_prefix/add_context_suffix)及PromptMutationfacaderenderer.py— 将PromptAssembly渲染回ProviderRequest,渲染幂等tracing.py— 生成结构化 trace 快照,支持visible_in_trace控制敏感内容可见性__init__.py— 统一导出公共 API修改文件:
astrbot/core/astr_main_agent.py— 重构 prompt 组装流程,从内联拼接改为向PromptAssembly注册后统一渲染astrbot/core/star/star_handler.py— 新增EventType.OnPromptAssemblyEvent事件类型astrbot/core/star/register/__init__.py/star_handler.py— 新增register_on_prompt_assembly装饰器,支持插件注册 prompt assembly hookastrbot/api/event/filter/__init__.py— 导出新事件类型astrbot/dashboard/routes/plugin.py— 适配新事件类型新增测试:
tests/unit/test_prompt_renderer.py— 渲染器完整测试(244 行)tests/unit/test_internal_agent_trace.py— trace 快照测试(135 行)tests/unit/test_astr_main_agent.py— 补充 prompt assembly 相关用例文件变更统计:
改动类型:
重构/整理
新功能(插件 hook)
This is NOT a breaking change. / 这不是一个破坏性变更。
Implementation Details / 实现说明
三通道设计:
system_blocksreq.system_promptuser_append_partsreq.extra_user_content_partscontext_contributionsreq.contexts渲染流程:
PromptAssembly注册区块(按通道 + order 排序)render_prompt_assembly()一次性渲染回ProviderRequestrendered标志防重入)插件 hook (
on_prompt_assembly):PromptMutationfacade 向各通道追加内容Screenshots or Test Results / 运行截图或测试结果
# 运行测试 python -m pytest tests/unit/test_prompt_renderer.py tests/unit/test_internal_agent_trace.py tests/unit/test_astr_main_agent.py -v所有新增测试覆盖:
visible_in_trace过滤PromptMutationfacade 方法调用add_context_prefix/add_context_suffix正确拼接Checklist / 检查清单
😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
/ 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。
👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
/ 我的更改经过了良好的测试,并已在上方提供了"验证步骤"和"运行截图"。
🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in
requirements.txtandpyproject.toml./ 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到
requirements.txt和pyproject.toml文件相应位置。😮 My changes do not introduce malicious code.
/ 我的更改没有引入恶意代码。
Summary by Sourcery
Introduce a structured Prompt Assembly pipeline for composing system, user, and context prompts and integrate it into main agent request building and plugin infrastructure.
New Features:
Enhancements:
Tests: