diff --git a/config/plugins.thought-aligner.example.json b/config/plugins.thought-aligner.example.json new file mode 100644 index 0000000..c311f38 --- /dev/null +++ b/config/plugins.thought-aligner.example.json @@ -0,0 +1,42 @@ +{ + "phases": { + "llm_before": { + "client": [], + "server": [] + }, + "llm_after": { + "client": [], + "server": [ + { + "name": "thought_aligner", + "env": { + "base_url": "$THOUGHT_ALIGNER_BASE_URL", + "api_key": "$THOUGHT_ALIGNER_API_KEY", + "model": "$THOUGHT_ALIGNER_MODEL" + }, + "kwargs": { + "timeout_s": 30, + "failure_mode": "deny", + "max_history_items": 8, + "max_instruction_chars": 12000, + "max_thought_chars": 8000, + "max_observation_chars": 12000 + } + } + ] + }, + "tool_before": { + "client": [], + "server": [ + { + "name": "rule_based_plugin", + "env": {} + } + ] + }, + "tool_after": { + "client": [], + "server": [] + } + } +} diff --git a/docs/en/SUMMARY.md b/docs/en/SUMMARY.md index 00a2270..543ae9e 100644 --- a/docs/en/SUMMARY.md +++ b/docs/en/SUMMARY.md @@ -18,6 +18,7 @@ * [AgentGuard Plugins](plugins.md) * [Builtin Plugins](plugins/builtin_plugins.md) * [rule_based_plugin](plugins/rule_based_plugin.md) + * [Thought-Aligner](plugins/thought_aligner.md) * [Visual Policy Configuration](policies/quick_config.md) * [Policy DSL Structure](policies/dsl_basic_structure.md) * [jailbreak_check](plugins/jailbreak_check.md) diff --git a/docs/en/plugins/builtin_plugins.md b/docs/en/plugins/builtin_plugins.md index 3296844..0cde44a 100644 --- a/docs/en/plugins/builtin_plugins.md +++ b/docs/en/plugins/builtin_plugins.md @@ -10,5 +10,6 @@ AgentGuard includes built-in plugins for common runtime protection needs. This p - [rule_based_plugin](rule_based_plugin.md): a server-only rule engine for tool-call access control and policy evaluation that can either return fixed `ALLOW` / `DENY` decisions or escalate matched cases to `HUMAN_CHECK` / `LLM_CHECK`. - [jailbreak_check](jailbreak_check.md): the same LLM-input detector can also run on the AgentGuard server when you want centralized governance and auditing. +- [Thought-Aligner](thought_aligner.md): an opt-in `llm_after` intervention that rewrites exposed reasoning on the server and makes a compatible Python client regenerate its action before execution. Other utility plugins also exist in the codebase, but this page focuses on the built-in plugins that are documented for direct use in the public docs. diff --git a/docs/en/plugins/custom_server_plugin.md b/docs/en/plugins/custom_server_plugin.md index 2ab7c26..3eb2fb5 100644 --- a/docs/en/plugins/custom_server_plugin.md +++ b/docs/en/plugins/custom_server_plugin.md @@ -143,7 +143,7 @@ Server plugin specs are read from the `server` list in `config/plugins.json` or - `name`: registered plugin name. - `class` or `plugin`: optional import-path alternatives to `name`. - The current server runtime resolves plugin classes by `name` or import path. -- Extra fields may remain in stored config, but the current server plugin manager does not inject `env` or `kwargs` into server plugin constructors. +- `kwargs` values are passed to the plugin constructor. `env` maps constructor attribute names to environment-variable references such as `$MY_PLUGIN_API_KEY`; references are resolved in the server process. ## Output diff --git a/docs/en/plugins/thought_aligner.md b/docs/en/plugins/thought_aligner.md new file mode 100644 index 0000000..aa025b1 --- /dev/null +++ b/docs/en/plugins/thought_aligner.md @@ -0,0 +1,78 @@ +# Thought-Aligner + +`thought_aligner` is an opt-in `llm_after` server plugin. It holds a Python agent's first model response, sends the exposed reasoning to a server-hosted Thought-Aligner endpoint, and instructs the same agent model to regenerate `Action` and `Action Input` from the aligned thought before the first action can reach the framework's parser or tool executor. + +The implementation follows the [Thought-Aligner-7B model-card prompt](https://huggingface.co/WhitzardAgent/Thought-Aligner-7B). Review that model's CC BY-NC 4.0 license before commercial use. + +## Execution flow + +1. The patched Python client calls the original agent model but keeps its response inside the wrapper. +2. The client normalizes the LLM output and tells the server whether that concrete request shape can be safely regenerated. +3. The server builds: + - the user instruction, preferring explicit metadata and otherwise reading user/human messages or rendered prompts; + - the current thought, preferring explicit normalized fields, then common reasoning aliases, thought tags, and ReAct `Thought:` text; + - completed earlier thought/tool-result pairs, formatted as ` ... ` and ` ... `. +4. If Thought-Aligner changes the thought, the server returns `align_thought` with the aligned thought. +5. The client copies the original request, injects the aligned thought plus an action-only instruction, calls the original model once more, and returns only the regenerated result. The retry is marked and cannot trigger another alignment loop. + +The Thought-Aligner model call and its credentials stay on the AgentGuard server. The original agent model remains on the client side because only the client owns the native model callable and framework-specific request/response objects. + +## Configuration + +Set dedicated server environment variables. Do not put a key in JSON or commit it to the repository. + +```bash +export THOUGHT_ALIGNER_BASE_URL="https://your-thought-aligner-host/v1" +export THOUGHT_ALIGNER_API_KEY="replace-with-a-secret" +export THOUGHT_ALIGNER_MODEL="thought-aligner-7b" +export AGENTGUARD_SERVER_PLUGIN_CONFIG="./config/plugins.thought-aligner.example.json" +``` + +Then start the AgentGuard server normally. The example file enables the plugin only in the server-side `llm_after` phase; the regular `config/plugins.json` remains unchanged, so Thought-Aligner is disabled by default. + +The Python client needs a server URL and a decision timeout longer than the server plugin's model timeout: + +```python +from agentguard import AgentGuard + +guard = AgentGuard( + "agent-session", + server_url="http://127.0.0.1:8000", + remote_timeout_s=45, + remote_retries=0, +) +guard.attach_langchain(agent) +``` + +If a framework passes an opaque prompt and the server cannot reliably recover the original user task, provide it explicitly before the guarded turn: + +```python +guard.context.metadata["instruction"] = user_instruction +``` + +Relevant plugin options: + +- `timeout_s`: Thought-Aligner endpoint timeout; default `30` seconds. +- `failure_mode`: `deny` by default, which withholds the first action if an attempted alignment fails. `allow` preserves availability but releases the original response after a model failure. +- `max_history_items`: maximum completed thought/observation pairs; default `8`. +- `max_instruction_chars`, `max_thought_chars`, `max_observation_chars`: per-field bounds before the external model call. + +## Supported input and output forms + +Thought extraction is independent of framework classes and recognizes: + +- normalized `thought`; +- nested `reasoning_content`, `reasoning`, `thinking`, `plan`, and `analysis` fields; +- ``, ``, ``, ``, and `` blocks; +- ReAct-style `Thought:` content before `Action`, `Action Input`, `Observation`, or `Final Answer`. + +Python regeneration currently supports non-streaming, concrete patched LLM calls whose request is a string, a dict/list/tuple of chat messages, an `input`/`prompt`/`messages` argument, or a LangChain-style `agent_scratchpad`. Common string, dictionary, and Pydantic message responses retain their native shape where possible. + +## Safety and compatibility boundaries + +- A provider that does not expose reasoning gives AgentGuard no truthful Thought to align. In that case the plugin is a no-op; it never invents hidden chain-of-thought. +- A thought-bearing call from an old or unsupported client is denied before the Thought-Aligner call because that client cannot prove it will regenerate the action. +- Streaming agents that emit action-bearing chunks before `llm_after`, including current streaming-specific integrations, are not covered by this interception path. Buffer the complete turn or add a framework-native pre-action hook before enabling the plugin. +- The implementation performs at most one regeneration. A second `align_thought` directive is blocked. +- Instruction, exposed thought, and selected observations leave the AgentGuard server for the configured model endpoint. Apply your data-retention, redaction, and regional-processing requirements to that endpoint. The server decision metadata contains the aligned thought but does not copy the original instruction or unsafe thought into the decision. +- Thought alignment complements, but does not replace, AgentGuard tool policies. The regenerated action still passes through the existing `tool_before` policy path. diff --git a/docs/zh/SUMMARY.md b/docs/zh/SUMMARY.md index 365bf99..80a853c 100644 --- a/docs/zh/SUMMARY.md +++ b/docs/zh/SUMMARY.md @@ -18,6 +18,7 @@ * [AgentGuard插件](plugins.md) * [内置插件](plugins/builtin_plugins.md) * [rule_based_plugin](plugins/rule_based_plugin.md) + * [Thought-Aligner](plugins/thought_aligner.md) * [可视化策略配置](policies/quick_config.md) * [策略 DSL 基本结构](policies/dsl_basic_structure.md) * [jailbreak_check](plugins/jailbreak_check.md) diff --git a/docs/zh/plugins/builtin_plugins.md b/docs/zh/plugins/builtin_plugins.md index f666f00..f58378d 100644 --- a/docs/zh/plugins/builtin_plugins.md +++ b/docs/zh/plugins/builtin_plugins.md @@ -10,5 +10,6 @@ AgentGuard 提供了一组面向常见运行时防护需求的内置 plugin。 - [rule_based_plugin](rule_based_plugin.md)：一个仅运行在 server 侧的规则引擎，用于工具调用访问控制和策略评估；它既可以直接返回固定的 `ALLOW` / `DENY`，也可以把命中的情况转入 `HUMAN_CHECK` / `LLM_CHECK`。 - [jailbreak_check](jailbreak_check.md)：同一个 LLM 输入检测器也可以部署在 AgentGuard Server 侧，用于集中式治理和审计。 +- [Thought-Aligner](thought_aligner.md)：一个按需启用的 `llm_after` 防御，在 server 侧改写可获得的推理，并让兼容的 Python client 在执行前重新生成 Action。 代码库里还有其他工具型 plugin，但这个页面当前聚焦于已经在公开文档中单独展开说明的内置 plugin。 diff --git a/docs/zh/plugins/custom_server_plugin.md b/docs/zh/plugins/custom_server_plugin.md index b0c38b0..de02731 100644 --- a/docs/zh/plugins/custom_server_plugin.md +++ b/docs/zh/plugins/custom_server_plugin.md @@ -143,7 +143,7 @@ Server plugin spec 从 `config/plugins.json` 或运行时 plugin config 的 `ser - `name`：注册后的 plugin 名称。 - `class` 或 `plugin`：也可以作为 `name` 的替代形式，用来写导入路径。 - 当前 server runtime 会按 `name` 或导入路径解析 plugin 类。 -- 额外字段会保留在配置中，但当前 server plugin manager 不会把 `env` 或 `kwargs` 注入 server plugin 构造函数。 +- `kwargs` 的值会传给 plugin 构造函数；`env` 可以把构造参数名映射为 `$MY_PLUGIN_API_KEY` 这类环境变量引用，并由 server 进程解析。 ## 输出 diff --git a/docs/zh/plugins/thought_aligner.md b/docs/zh/plugins/thought_aligner.md new file mode 100644 index 0000000..4dee082 --- /dev/null +++ b/docs/zh/plugins/thought_aligner.md @@ -0,0 +1,78 @@ +# Thought-Aligner + +`thought_aligner` 是一个默认关闭、按需启用的 server 侧 `llm_after` plugin。它会先扣住 Python agent 第一次生成的模型响应，把其中可获得的推理发送到 server 上配置的 Thought-Aligner 端点，然后让同一个 agent 模型基于安全 Thought 重新生成 `Action` 和 `Action Input`。第一次生成的 Action 不会先交给框架解析器或工具执行器。 + +实现遵循 [Thought-Aligner-7B 模型卡](https://huggingface.co/WhitzardAgent/Thought-Aligner-7B)给出的 prompt。商业使用前请确认该模型的 CC BY-NC 4.0 许可证是否适用。 + +## 执行流程 + +1. Python client 调用原 agent 模型，但 wrapper 暂不向 agent 返回第一次响应。 +2. Client 标准化 LLM 输出，并告诉 server 当前这次请求的数据形态是否能安全回跳。 +3. Server 构造： + - 用户指令：优先使用显式 metadata，否则从 user/human 消息、嵌套输入或渲染后的 prompt 中提取； + - 当前 Thought：依次尝试标准字段、常见 reasoning 别名、Thought 标签和 ReAct `Thought:` 文本； + - 之前已经完成的 Thought/工具结果对，并用 ` ... `、` ... ` 标记。 +4. Thought-Aligner 改写了 Thought 时，server 返回带安全 Thought 的 `align_thought` 决策。 +5. Client 复制原请求，注入安全 Thought 和“只生成 Action/Final Answer”的指令，再调用一次原模型，只把重新生成的结果返回给 agent。重试事件会被标记，不会再次进入对齐循环。 + +Thought-Aligner 调用及其凭证始终在 AgentGuard server 侧。原 agent 模型仍由 client 调用，因为只有 client 持有原生模型 callable 以及具体框架的请求、响应对象。 + +## 配置 + +在 server 进程中设置独立环境变量。不要把密钥直接写入 JSON，也不要提交到仓库。 + +```bash +export THOUGHT_ALIGNER_BASE_URL="https://your-thought-aligner-host/v1" +export THOUGHT_ALIGNER_API_KEY="replace-with-a-secret" +export THOUGHT_ALIGNER_MODEL="thought-aligner-7b" +export AGENTGUARD_SERVER_PLUGIN_CONFIG="./config/plugins.thought-aligner.example.json" +``` + +然后正常启动 AgentGuard server。示例配置只在 server 的 `llm_after` 阶段启用该 plugin；现有 `config/plugins.json` 不变，所以默认行为不会开启 Thought-Aligner。 + +Python client 需要配置 server 地址，并确保远程决策超时大于 server plugin 的模型超时： + +```python +from agentguard import AgentGuard + +guard = AgentGuard( + "agent-session", + server_url="http://127.0.0.1:8000", + remote_timeout_s=45, + remote_retries=0, +) +guard.attach_langchain(agent) +``` + +如果某个框架传入的是不透明 prompt，server 无法稳定还原最初的用户任务，可以在受保护轮次开始前显式设置： + +```python +guard.context.metadata["instruction"] = user_instruction +``` + +主要 plugin 参数： + +- `timeout_s`：Thought-Aligner 端点超时，默认 `30` 秒。 +- `failure_mode`：默认 `deny`；已经进入对齐的请求如果模型调用失败，会扣住第一次 Action。设置为 `allow` 可优先保证可用性，但模型失败后会释放原响应。 +- `max_history_items`：最多传入的已完成 Thought/Observation 对，默认 `8`。 +- `max_instruction_chars`、`max_thought_chars`、`max_observation_chars`：发送给外部模型前各字段的长度上限。 + +## 支持的输入与输出形式 + +Thought 提取不依赖具体框架类，目前识别： + +- 标准化后的 `thought`； +- 嵌套的 `reasoning_content`、`reasoning`、`thinking`、`plan`、`analysis` 字段； +- ``、``、``、``、`` 块； +- ReAct 文本中位于 `Action`、`Action Input`、`Observation` 或 `Final Answer` 之前的 `Thought:`。 + +Python 回跳当前支持非流式、已经被具体 patch 的 LLM 调用，请求形态可以是字符串、消息 dict/list/tuple、`input`/`prompt`/`messages` 参数，或 LangChain 风格 `agent_scratchpad`。对于常见字符串、字典和 Pydantic 消息响应，会尽量保留原生返回类型。 + +## 安全与兼容边界 + +- 如果模型供应商完全不暴露推理，AgentGuard 就没有真实 Thought 可供对齐。此时 plugin 不执行改写，也不会猜测或伪造隐藏思维链。 +- 如果旧 client 或不支持回跳的 client 发送了包含 Thought 的事件，server 会在调用 Thought-Aligner 前拒绝该轮，因为无法证明 client 会重新生成 Action。 +- 如果流式 agent 在 `llm_after` 之前就把包含 Action 的 chunk 发出，当前拦截路径无法提供保护。启用前需要缓冲完整一轮，或增加框架原生的 Action 前 hook。 +- 每轮最多回跳一次；第二个 `align_thought` 决策会被阻断。 +- 用户指令、可见 Thought 和选中的 Observation 会离开 AgentGuard server，发送到配置的模型端点。需要对该端点落实数据保留、脱敏和地域处理要求。server 的最终 decision metadata 只携带安全 Thought，不会再次复制原指令或不安全 Thought。 +- Thought 对齐不能替代工具策略。重新生成的 Action 仍会经过 AgentGuard 原有的 `tool_before` 策略链。 diff --git a/src/client/python/agentguard/adapters/agent/langchain.py b/src/client/python/agentguard/adapters/agent/langchain.py index 6c6ae8d..d079275 100644 --- a/src/client/python/agentguard/adapters/agent/langchain.py +++ b/src/client/python/agentguard/adapters/agent/langchain.py @@ -371,6 +371,15 @@ def _normalize_langchain_request( def _normalize_langchain_llm_output(value: Any) -> Any: normalized = _normalize_langchain_value(value) + if isinstance(normalized, str): + parsed = _parse_tagged_llm_output(normalized) + if parsed.thought is None: + return normalized + return { + "output": normalized, + "thought": parsed.thought, + "final_output": parsed.final_output, + } return _extract_langchain_llm_output_fields(normalized) @@ -419,7 +428,7 @@ def _first_non_empty_text(value: dict[str, Any], *keys: str) -> str | None: @dataclass(frozen=True) class _ParsedLLMOutput: thought: str | None - final_output: str + final_output: str | None _THOUGHT_TAG_RE = re.compile( @@ -430,12 +439,28 @@ class _ParsedLLMOutput: r"<(?Panswer|final|final_output)\b[^>]*>(?P.*?)", flags=re.IGNORECASE | re.DOTALL, ) +_REACT_THOUGHT_RE = re.compile( + r"(?:^|\n)\s*(?:Thought|Reasoning|Analysis|思考)\s*:\s*(?P.*?)" + r"(?=\n\s*(?:Action(?:\s+Input)?|Observation|Final\s+Answer|Answer|" + r"行动|观察|最终答案)\s*:|\Z)", + flags=re.IGNORECASE | re.DOTALL, +) def _parse_tagged_llm_output(output: str) -> _ParsedLLMOutput: thought_matches = list(_THOUGHT_TAG_RE.finditer(output)) if not thought_matches: - return _ParsedLLMOutput(thought=None, final_output=output) + react_match = _REACT_THOUGHT_RE.search(output) + if react_match is None: + return _ParsedLLMOutput(thought=None, final_output=output) + thought = react_match.group("body").strip() or None + remainder = f"{output[:react_match.start()]}{output[react_match.end():]}".strip() + final_output = ( + None + if re.match(r"^\s*Action\s*:", remainder, flags=re.IGNORECASE) + else remainder + ) + return _ParsedLLMOutput(thought=thought, final_output=final_output) thought_parts = [match.group("body").strip() for match in thought_matches] thought = "\n\n".join(part for part in thought_parts if part) or None diff --git a/src/client/python/agentguard/adapters/agent/patching.py b/src/client/python/agentguard/adapters/agent/patching.py index da79fe6..e2768fc 100644 --- a/src/client/python/agentguard/adapters/agent/patching.py +++ b/src/client/python/agentguard/adapters/agent/patching.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import Any +from agentguard.adapters.agent import thought_alignment from agentguard.adapters.agent.normalization import ( DEFAULT_AGENT_EVENT_NORMALIZER, AgentEventNormalizer, @@ -127,6 +128,8 @@ def guard_llm_after( normalizer: AgentEventNormalizer | None = None, fn: Callable[..., Any] | None = None, owner: Any = None, + extra_metadata: dict[str, Any] | None = None, + thought_override: str | None = None, ) -> GuardDecision: normalized = _resolve_normalizer(normalizer).normalize_llm_output( label=label, @@ -134,8 +137,19 @@ def guard_llm_after( fn=fn, owner=owner, ) + payload = normalized.payload + metadata = dict(normalized.metadata) + metadata.setdefault("output_type", type(output).__name__) + if extra_metadata: + metadata.update(extra_metadata) + if thought_override: + if isinstance(payload, dict): + payload = dict(payload) + payload["thought"] = thought_override + else: + payload = {"output": payload, "thought": thought_override} return guard.runtime.guard( - ev.llm_output(guard.context, normalized.payload, **dict(normalized.metadata)), + ev.llm_output(guard.context, payload, **metadata), phase="after", ).decision @@ -302,6 +316,9 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: if before_blocked is not None: return before_blocked raw = await fn(*args, **kwargs) + regeneration_supported = thought_alignment.supports_thought_regeneration( + args, kwargs + ) decision = guard_llm_after( guard, raw, @@ -309,7 +326,23 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: normalizer=normalizer, fn=fn, owner=owner, + extra_metadata={ + "thought_regeneration_supported": regeneration_supported, + "thought_alignment_attempt": 0, + }, ) + if decision.decision_type == DecisionType.ALIGN_THOUGHT: + return await _regenerate_async( + guard, + fn, + args=args, + kwargs=kwargs, + raw=raw, + decision=decision, + label=label, + normalizer=normalizer, + owner=owner, + ) blocked = _blocked_llm_value(decision) return blocked if blocked is not None else raw except Exception: @@ -336,6 +369,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: if before_blocked is not None: return before_blocked raw = fn(*args, **kwargs) + regeneration_supported = thought_alignment.supports_thought_regeneration(args, kwargs) decision = guard_llm_after( guard, raw, @@ -343,7 +377,23 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: normalizer=normalizer, fn=fn, owner=owner, + extra_metadata={ + "thought_regeneration_supported": regeneration_supported, + "thought_alignment_attempt": 0, + }, ) + if decision.decision_type == DecisionType.ALIGN_THOUGHT: + return _regenerate_sync( + guard, + fn, + args=args, + kwargs=kwargs, + raw=raw, + decision=decision, + label=label, + normalizer=normalizer, + owner=owner, + ) blocked = _blocked_llm_value(decision) return blocked if blocked is not None else raw except Exception: @@ -453,9 +503,126 @@ def _blocked_llm_value(decision: GuardDecision) -> Any | None: "reason": decision.reason, "decision": decision.decision_type.value, } + if decision.decision_type in { + DecisionType.ALIGN_THOUGHT, + DecisionType.LOOP_BACK_TO_LLM, + }: + return { + "agentguard": "blocked", + "reason": decision.reason, + "decision": decision.decision_type.value, + } return None +def _regenerate_sync( + guard: Any, + fn: Callable[..., Any], + *, + args: tuple[Any, ...], + kwargs: dict[str, Any], + raw: Any, + decision: GuardDecision, + label: str, + normalizer: AgentEventNormalizer | None, + owner: Any, +) -> Any: + aligned_thought = thought_alignment.aligned_thought(decision) + prepared = ( + thought_alignment.prepare_thought_regeneration(args, kwargs, aligned_thought) + if aligned_thought + else None + ) + if prepared is None: + return _blocked_llm_value(decision) + retry_args, retry_kwargs = prepared + before_decision = guard_llm_before( + guard, + label=label, + args=retry_args, + kwargs=retry_kwargs, + normalizer=normalizer, + fn=fn, + owner=owner, + ) + before_blocked = _blocked_llm_value(before_decision) + if before_blocked is not None: + return before_blocked + regenerated = fn(*retry_args, **retry_kwargs) + retry_decision = guard_llm_after( + guard, + regenerated, + label=label, + normalizer=normalizer, + fn=fn, + owner=owner, + extra_metadata={ + "thought_regeneration_supported": True, + "thought_alignment_attempt": 1, + "thought_alignment_protocol": "thought_alignment_v1", + }, + thought_override=aligned_thought, + ) + blocked = _blocked_llm_value(retry_decision) + if blocked is not None: + return blocked + return thought_alignment.merge_aligned_thought(regenerated, raw, aligned_thought) + + +async def _regenerate_async( + guard: Any, + fn: Callable[..., Any], + *, + args: tuple[Any, ...], + kwargs: dict[str, Any], + raw: Any, + decision: GuardDecision, + label: str, + normalizer: AgentEventNormalizer | None, + owner: Any, +) -> Any: + aligned_thought = thought_alignment.aligned_thought(decision) + prepared = ( + thought_alignment.prepare_thought_regeneration(args, kwargs, aligned_thought) + if aligned_thought + else None + ) + if prepared is None: + return _blocked_llm_value(decision) + retry_args, retry_kwargs = prepared + before_decision = guard_llm_before( + guard, + label=label, + args=retry_args, + kwargs=retry_kwargs, + normalizer=normalizer, + fn=fn, + owner=owner, + ) + before_blocked = _blocked_llm_value(before_decision) + if before_blocked is not None: + return before_blocked + regenerated = await fn(*retry_args, **retry_kwargs) + retry_decision = guard_llm_after( + guard, + regenerated, + label=label, + normalizer=normalizer, + fn=fn, + owner=owner, + extra_metadata={ + "thought_regeneration_supported": True, + "thought_alignment_attempt": 1, + "thought_alignment_protocol": "thought_alignment_v1", + }, + thought_override=aligned_thought, + ) + blocked = _blocked_llm_value(retry_decision) + if blocked is not None: + return blocked + return thought_alignment.merge_aligned_thought(regenerated, raw, aligned_thought) + + def _sync_local_cache_now(guard: Any, *, reason: str) -> None: rt = getattr(guard, "runtime", None) sync = getattr(rt, "sync_local_cache_now", None) diff --git a/src/client/python/agentguard/adapters/agent/thought_alignment.py b/src/client/python/agentguard/adapters/agent/thought_alignment.py new file mode 100644 index 0000000..1995da0 --- /dev/null +++ b/src/client/python/agentguard/adapters/agent/thought_alignment.py @@ -0,0 +1,175 @@ +"""Framework-neutral helpers for one bounded Thought-Aligner regeneration.""" +from __future__ import annotations + +import copy +import re +from typing import Any + +from agentguard.schemas.decisions import GuardDecision + + +def aligned_thought(decision: GuardDecision) -> str | None: + if decision.metadata.get("protocol") != "thought_alignment_v1": + return None + value = decision.metadata.get("aligned_thought") + if not isinstance(value, str) or not value.strip(): + return None + return value.strip() + + +def supports_thought_regeneration( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> bool: + return prepare_thought_regeneration(args, kwargs, "probe") is not None + + +def prepare_thought_regeneration( + args: tuple[Any, ...], + kwargs: dict[str, Any], + safe_thought: str, +) -> tuple[tuple[Any, ...], dict[str, Any]] | None: + retry_args = tuple(args) + retry_kwargs = dict(kwargs) + + if isinstance(retry_kwargs.get("agent_scratchpad"), str): + retry_kwargs["agent_scratchpad"] = _append_aligned_thought( + retry_kwargs["agent_scratchpad"], safe_thought + ) + return retry_args, retry_kwargs + + for key in ("messages", "input", "prompt"): + if key not in retry_kwargs: + continue + injected = _inject_thought(retry_kwargs[key], safe_thought) + if injected is not None: + retry_kwargs[key] = injected + return retry_args, retry_kwargs + + if not retry_args: + return None + injected = _inject_thought(retry_args[0], safe_thought) + if injected is None: + return None + return (injected, *retry_args[1:]), retry_kwargs + + +def merge_aligned_thought(regenerated: Any, original: Any, safe_thought: str) -> Any: + """Preserve common native response shapes while recording the effective thought.""" + if isinstance(regenerated, str): + if not _has_react_action(original) and not _has_react_action(regenerated): + return regenerated + match = _REACT_CONTINUATION_RE.search(regenerated) + continuation = match.group("body").strip() if match else regenerated.strip() + return f"Thought: {safe_thought}\n{continuation}" + + if isinstance(regenerated, dict): + merged = copy.deepcopy(regenerated) + merged["thought"] = safe_thought + additional = merged.get("additional_kwargs") + if isinstance(additional, dict): + additional["reasoning_content"] = safe_thought + return merged + + model_copy = getattr(regenerated, "model_copy", None) + if callable(model_copy): + additional = getattr(regenerated, "additional_kwargs", None) + if isinstance(additional, dict): + updated = dict(additional) + updated["reasoning_content"] = safe_thought + try: + return model_copy(update={"additional_kwargs": updated}) + except Exception: + pass + return regenerated + + +def _inject_thought(value: Any, safe_thought: str) -> Any | None: + if isinstance(value, str): + return _append_aligned_thought(value, safe_thought) + + if isinstance(value, dict): + cloned = dict(value) + if isinstance(cloned.get("agent_scratchpad"), str): + cloned["agent_scratchpad"] = _append_aligned_thought( + cloned["agent_scratchpad"], safe_thought + ) + return cloned + for key in ("messages", "input", "prompt"): + if key not in cloned: + continue + injected = _inject_thought(cloned[key], safe_thought) + if injected is not None: + cloned[key] = injected + return cloned + return None + + if isinstance(value, (list, tuple)): + messages = list(value) + additions = _thought_messages(messages, safe_thought) + if additions is None: + return None + combined = [*messages, *additions] + return tuple(combined) if isinstance(value, tuple) else combined + return None + + +def _thought_messages(messages: list[Any], safe_thought: str) -> list[Any] | None: + directive = _action_only_directive() + if not messages or all(isinstance(item, dict) for item in messages): + return [ + {"role": "assistant", "content": f"Thought: {safe_thought}"}, + {"role": "user", "content": directive}, + ] + if all(isinstance(item, str) for item in messages): + return [f"Thought: {safe_thought}\n{directive}"] + + try: + from langchain_core.messages import AIMessage, HumanMessage + except Exception: + return None + if not all( + hasattr(item, "content") and "langchain" in type(item).__module__.lower() + for item in messages + ): + return None + return [AIMessage(content=f"Thought: {safe_thought}"), HumanMessage(content=directive)] + + +def _append_aligned_thought(value: str, safe_thought: str) -> str: + prefix = value.rstrip() + if prefix: + prefix += "\n" + return f"{prefix}Thought: {safe_thought}\n\n{_action_only_directive()}" + + +def _action_only_directive() -> str: + return ( + "(**DO NOT generate 'Thought' again! Output 'Action' and 'Action Input', " + "or 'Final Answer', directly. If user confirmation is required, output " + "the 'Final Answer' directly and do not generate 'Action'.**)" + ) + + +_REACT_CONTINUATION_RE = re.compile( + r"(?:^|\n)\s*(?P(?:Action(?:\s+Input)?|Final\s+Answer)\s*:.*)\Z", + flags=re.IGNORECASE | re.DOTALL, +) + + +def _has_react_action(value: Any) -> bool: + if isinstance(value, str): + return bool(re.search(r"(?:^|\n)\s*Action\s*:", value, flags=re.IGNORECASE)) + if isinstance(value, dict): + for key in ("output", "text", "content", "message"): + if _has_react_action(value.get(key)): + return True + return False + + +__all__ = [ + "aligned_thought", + "merge_aligned_thought", + "prepare_thought_regeneration", + "supports_thought_regeneration", +] diff --git a/src/client/python/agentguard/schemas/events.py b/src/client/python/agentguard/schemas/events.py index 50fc814..12f9845 100644 --- a/src/client/python/agentguard/schemas/events.py +++ b/src/client/python/agentguard/schemas/events.py @@ -384,6 +384,8 @@ def _coerce_llm_output(value: Any) -> LLMOutput: return LLMOutput(output=_coerce_text(value)) thought = data.get("thought") + if thought is None: + thought = _nested_reasoning_value(data) final_output = data.get("final_output") output = data.get("output") if output is None: @@ -422,7 +424,18 @@ def _llm_output_fields(value: Any) -> dict[str, Any] | None: if data is None: attrs = { key: getattr(value, key) - for key in ("output", "text", "content", "message", "thought", "final_output") + for key in ( + "output", + "text", + "content", + "message", + "thought", + "reasoning_content", + "reasoning", + "thinking", + "analysis", + "final_output", + ) if getattr(value, key, None) is not None } data = attrs or None @@ -430,7 +443,64 @@ def _llm_output_fields(value: Any) -> dict[str, Any] | None: if not data: return None - recognized = ("output", "text", "content", "message", "thought", "final_output") + recognized = ( + "output", + "text", + "content", + "message", + "thought", + "reasoning_content", + "reasoning", + "thinking", + "analysis", + "final_output", + ) if not any(key in data for key in recognized): return None return data + + +_REASONING_KEYS = ( + "reasoning_content", + "reasoningContent", + "reasoning", + "thinking", + "plan", + "analysis", +) + + +def _nested_reasoning_value(value: Any, depth: int = 0) -> str | None: + if depth > 5: + return None + if isinstance(value, dict): + for key in _REASONING_KEYS: + text = _reasoning_text(value.get(key)) + if text: + return text + for item in value.values(): + nested = _nested_reasoning_value(item, depth + 1) + if nested: + return nested + elif isinstance(value, list): + for item in value: + nested = _nested_reasoning_value(item, depth + 1) + if nested: + return nested + return None + + +def _reasoning_text(value: Any) -> str | None: + if isinstance(value, str): + return value.strip() or None + if isinstance(value, list): + parts: list[str] = [] + for item in value: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict): + text = item.get("text") or item.get("content") or item.get("summary") + if isinstance(text, str) and text.strip(): + parts.append(text) + return "\n".join(parts).strip() or None + return None diff --git a/src/server/backend/runtime/plugins/llm_after/__init__.py b/src/server/backend/runtime/plugins/llm_after/__init__.py index 51729f7..417bca3 100644 --- a/src/server/backend/runtime/plugins/llm_after/__init__.py +++ b/src/server/backend/runtime/plugins/llm_after/__init__.py @@ -2,5 +2,6 @@ from __future__ import annotations from backend.runtime.plugins.llm_after.llm_output import LLMOutputPlugin +from backend.runtime.plugins.llm_after.thought_aligner import ThoughtAlignerPlugin -__all__ = ["LLMOutputPlugin"] +__all__ = ["LLMOutputPlugin", "ThoughtAlignerPlugin"] diff --git a/src/server/backend/runtime/plugins/llm_after/thought_aligner.py b/src/server/backend/runtime/plugins/llm_after/thought_aligner.py new file mode 100644 index 0000000..5a35e4a --- /dev/null +++ b/src/server/backend/runtime/plugins/llm_after/thought_aligner.py @@ -0,0 +1,146 @@ +"""Server-side Thought-Aligner intervention for LLM outputs.""" +from __future__ import annotations + +from typing import Any + +from backend.runtime.plugins.base import BasePlugin, CheckResult +from backend.runtime.plugins.registry import register +from backend.runtime.thought_alignment import ( + ThoughtAlignerClient, + ThoughtAlignmentError, + build_alignment_context, +) +from shared.schemas.context import RuntimeContext +from shared.schemas.decisions import DecisionType, GuardDecision +from shared.schemas.events import EventType, RuntimeEvent + + +@register( + name="thought_aligner", + description="Rewrite exposed agent reasoning before the client regenerates its action.", +) +class ThoughtAlignerPlugin(BasePlugin): + event_types = [EventType.LLM_OUTPUT] + + def __init__( + self, + *, + aligner: Any = None, + env: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + self._aligner_override = aligner + super().__init__(env=env, **kwargs) + + def check( + self, + event: RuntimeEvent, + context: RuntimeContext, + trajectory_window: list[RuntimeEvent] | None = None, + ) -> CheckResult: + if _as_int(event.metadata.get("thought_alignment_attempt"), 0) > 0: + return CheckResult(metadata={"thought_alignment": "retry_skipped"}) + + alignment = build_alignment_context( + event, + context, + trajectory_window, + max_instruction_chars=_as_int( + getattr(self, "max_instruction_chars", 12_000), 12_000 + ), + max_thought_chars=_as_int(getattr(self, "max_thought_chars", 8_000), 8_000), + max_observation_chars=_as_int( + getattr(self, "max_observation_chars", 12_000), 12_000 + ), + max_history_items=_as_int(getattr(self, "max_history_items", 8), 8), + ) + if alignment is None: + return CheckResult(metadata={"thought_alignment": "context_unavailable"}) + + if event.metadata.get("thought_regeneration_supported") is not True: + return CheckResult( + decision_candidate=GuardDecision.deny( + "Thought alignment requires client-side action regeneration support.", + metadata={"protocol": "thought_alignment_v1"}, + ), + risk_signals=["thought_alignment_unsupported_client"], + is_final=True, + ) + + try: + aligned = self._aligner().align( + alignment.formatted_instruction, + alignment.thought, + ) + thought_limit = max( + 1, + _as_int(getattr(self, "max_thought_chars", 8_000), 8_000), + ) + aligned = str(aligned).strip()[:thought_limit].strip() + if not aligned: + raise ThoughtAlignmentError("Thought-Aligner returned empty text") + except Exception: + return self._failure_result() + + if aligned == alignment.thought.strip(): + return CheckResult(metadata={"thought_alignment": "unchanged"}) + + decision = GuardDecision( + DecisionType.ALIGN_THOUGHT, + "Thought-Aligner rewrote the current reasoning; regenerate the action.", + risk_signals=["thought_alignment_applied"], + metadata={ + "aligned_thought": aligned, + "protocol": "thought_alignment_v1", + }, + ) + return CheckResult( + decision_candidate=decision, + risk_signals=["thought_alignment_applied"], + is_final=True, + metadata={"thought_alignment": "aligned"}, + ) + + def _aligner(self) -> Any: + if self._aligner_override is not None: + return self._aligner_override + return ThoughtAlignerClient( + base_url=getattr(self, "base_url", None), + api_key=getattr(self, "api_key", None), + model=getattr(self, "model", None), + timeout_s=_as_float(getattr(self, "timeout_s", 30.0), 30.0), + ) + + def _failure_result(self) -> CheckResult: + failure_mode = str(getattr(self, "failure_mode", "deny") or "deny").lower() + if failure_mode == "allow": + return CheckResult( + risk_signals=["thought_alignment_error"], + metadata={"thought_alignment": "error_allowed"}, + ) + return CheckResult( + decision_candidate=GuardDecision.deny( + "Thought alignment failed; the original action was not released.", + metadata={"protocol": "thought_alignment_v1"}, + ), + risk_signals=["thought_alignment_error"], + is_final=True, + metadata={"thought_alignment": "error_denied"}, + ) + + +def _as_int(value: Any, default: int) -> int: + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _as_float(value: Any, default: float) -> float: + try: + return float(value) + except (TypeError, ValueError): + return default + + +__all__ = ["ThoughtAlignerPlugin"] diff --git a/src/server/backend/runtime/thought_alignment.py b/src/server/backend/runtime/thought_alignment.py new file mode 100644 index 0000000..6ab14f4 --- /dev/null +++ b/src/server/backend/runtime/thought_alignment.py @@ -0,0 +1,467 @@ +"""Thought-Aligner context extraction and OpenAI-compatible transport.""" +from __future__ import annotations + +import json +import re +import urllib.request +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from shared.schemas.context import RuntimeContext +from shared.schemas.events import EventType, RuntimeEvent + +_THOUGHT_KEYS = ( + "thought", + "reasoning_content", + "reasoningContent", + "reasoning", + "thinking", + "plan", + "analysis", +) +_INSTRUCTION_KEYS = ( + "instruction", + "user_instruction", + "task_instruction", + "goal", + "task", +) +_THOUGHT_TAG_RE = re.compile( + r"<(?Pthink|thought|reason|reasoning|analysis)\b[^>]*>" + r"(?P.*?)", + flags=re.IGNORECASE | re.DOTALL, +) +_REACT_THOUGHT_RE = re.compile( + r"(?:^|\n)\s*(?:Thought|Reasoning|Analysis|思考)\s*:\s*(?P.*?)" + r"(?=\n\s*(?:Action(?:\s+Input)?|Observation|Final\s+Answer|Answer|" + r"行动|观察|最终答案)\s*:|\Z)", + flags=re.IGNORECASE | re.DOTALL, +) +_RENDERED_USER_RE = re.compile( + r"(?:^|\n)\s*(?:Human|User|用户)\s*:\s*(?P.*?)" + r"(?=\n\s*(?:AI|Assistant|System|Human|User|工具|观察)\s*:|\Z)", + flags=re.IGNORECASE | re.DOTALL, +) +_OFFICIAL_PROMPT = ( + "BEGINNING OF CONVERSATION: USER: Edit the following Instruction-Thought pair " + "to make it more helpful and harmless. : {instruction} | {thought} ASSISTANT:" +) + + +class ThoughtAlignmentError(RuntimeError): + """Raised when a configured Thought-Aligner endpoint cannot produce a thought.""" + + +@dataclass(frozen=True) +class ThoughtObservation: + thought: str + observation: str + + +@dataclass(frozen=True) +class AlignmentContext: + instruction: str + thought: str + history: tuple[ThoughtObservation, ...] = () + + @property + def formatted_instruction(self) -> str: + parts = [self.instruction] + for item in self.history: + parts.append(f" {_escape_marker_text(item.thought)} ") + parts.append( + f" {_escape_marker_text(item.observation)} " + ) + return "\n".join(part for part in parts if part).strip() + + +def build_alignment_context( + event: RuntimeEvent, + context: RuntimeContext, + trajectory_window: list[RuntimeEvent] | None, + *, + max_instruction_chars: int = 12_000, + max_thought_chars: int = 8_000, + max_observation_chars: int = 12_000, + max_history_items: int = 8, +) -> AlignmentContext | None: + """Build the model input without depending on framework-specific classes.""" + if event.event_type != EventType.LLM_OUTPUT: + return None + + thought = extract_thought(event.payload) + if not thought: + return None + + trace = list(trajectory_window or []) + instruction = _extract_instruction(event, context, trace) + if not instruction: + return None + + history = _extract_history(trace) + history_limit = max(0, int(max_history_items)) + selected_history = history[-history_limit:] if history_limit else [] + bounded_history = tuple( + ThoughtObservation( + thought=_clip(item.thought, max_thought_chars), + observation=_clip(item.observation, max_observation_chars), + ) + for item in selected_history + ) + return AlignmentContext( + instruction=_clip(instruction, max_instruction_chars), + thought=_clip(thought, max_thought_chars), + history=bounded_history, + ) + + +def extract_thought(value: Any) -> str | None: + """Extract exposed reasoning from normalized fields, tags, or ReAct text.""" + explicit = _find_reasoning_value(value, top_level_only=True) + if explicit: + return explicit + + nested = _find_reasoning_value(value, top_level_only=False) + if nested: + return nested + + for text in _candidate_output_texts(value): + tagged = _extract_tagged_thought(text) + if tagged: + return tagged + react = _extract_react_thought(text) + if react: + return react + return None + + +class ThoughtAlignerClient: + """Small dedicated client for an OpenAI-compatible Thought-Aligner endpoint.""" + + def __init__( + self, + *, + base_url: str | None, + api_key: str | None, + model: str | None, + timeout_s: float = 30.0, + opener: Callable[..., Any] | None = None, + ) -> None: + self.base_url = str(base_url or "").strip() + self.api_key = str(api_key or "").strip() + self.model = str(model or "").strip() + self.timeout_s = float(timeout_s) + self._opener = opener or urllib.request.urlopen + + def align(self, instruction: str, thought: str) -> str: + if not self.base_url or not self.api_key or not self.model: + raise ThoughtAlignmentError("Thought-Aligner configuration is incomplete") + + prompt = _OFFICIAL_PROMPT.format(instruction=instruction, thought=thought) + payload = { + "model": self.model, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": prompt}, + ], + } + request = urllib.request.Request( + _chat_completions_url(self.base_url), + data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with self._opener(request, timeout=self.timeout_s) as response: + decoded = json.loads(response.read().decode("utf-8")) + content = decoded["choices"][0]["message"]["content"] + except ThoughtAlignmentError: + raise + except Exception as exc: + raise ThoughtAlignmentError("Thought-Aligner request or response failed") from exc + + if not isinstance(content, str) or not content.strip(): + raise ThoughtAlignmentError("Thought-Aligner response did not contain text") + return _clean_aligned_thought(content) + + +def _extract_instruction( + event: RuntimeEvent, + context: RuntimeContext, + trace: list[RuntimeEvent], +) -> str | None: + for source in (event.metadata, context.metadata): + explicit = _first_text_for_keys(source, _INSTRUCTION_KEYS) + if explicit: + return explicit + + for candidate in reversed(trace): + if candidate.event_type != EventType.LLM_INPUT: + continue + instruction = _extract_user_instruction(candidate.payload) + if instruction: + return instruction + return None + + +def _extract_user_instruction(value: Any) -> str | None: + normalized = _mapping_or_attributes(value) + if normalized is not None: + messages = normalized.get("messages") + if messages is not None: + role_text = _last_role_text(messages) + if role_text: + return role_text + nested = _extract_user_instruction(messages) + if nested: + return nested + + role = str(normalized.get("role") or normalized.get("type") or "").lower() + content = _as_text(normalized.get("content")) + if role in {"user", "human"} and content: + return _rendered_user_or_text(content) + + for key in ("input", "prompt", "query", "request", "text", "content"): + if key not in normalized: + continue + nested = _extract_user_instruction(normalized[key]) + if nested: + return nested + return None + + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + role_text = _last_role_text(value) + if role_text: + return role_text + for item in reversed(value): + nested = _extract_user_instruction(item) + if nested: + return nested + return None + + text = _as_text(value) + return _rendered_user_or_text(text) if text else None + + +def _last_role_text(value: Any) -> str | None: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + return None + for item in reversed(value): + mapping = _mapping_or_attributes(item) + if mapping is None: + continue + role = str(mapping.get("role") or mapping.get("type") or "").lower() + if role not in {"user", "human"}: + continue + content = _as_text(mapping.get("content")) + if content: + return _rendered_user_or_text(content) + return None + + +def _rendered_user_or_text(text: str) -> str: + matches = list(_RENDERED_USER_RE.finditer(text)) + if matches: + return matches[-1].group("body").strip() + return text.strip() + + +def _extract_history(trace: list[RuntimeEvent]) -> list[ThoughtObservation]: + history: list[ThoughtObservation] = [] + pending_thought: str | None = None + observations: list[Any] = [] + + def flush() -> None: + nonlocal pending_thought, observations + if pending_thought and observations: + observation_value: Any = observations[0] if len(observations) == 1 else observations + history.append( + ThoughtObservation( + thought=pending_thought, + observation=_serialize_observation(observation_value), + ) + ) + pending_thought = None + observations = [] + + for item in trace: + if item.event_type == EventType.LLM_OUTPUT: + flush() + pending_thought = extract_thought(item.payload) + elif item.event_type == EventType.TOOL_RESULT and pending_thought: + observations.append(item.payload.get("result")) + flush() + return history + + +def _find_reasoning_value(value: Any, *, top_level_only: bool) -> str | None: + mapping = _mapping_or_attributes(value) + if mapping is None: + return None + + direct = _first_text_for_keys(mapping, _THOUGHT_KEYS) + if direct: + return direct + if top_level_only: + return None + + seen: set[int] = set() + + def walk(item: Any, depth: int) -> str | None: + if depth > 5 or id(item) in seen: + return None + if isinstance(item, (Mapping, list, tuple)): + seen.add(id(item)) + nested_mapping = _mapping_or_attributes(item) + if nested_mapping is not None: + found = _first_text_for_keys(nested_mapping, _THOUGHT_KEYS) + if found: + return found + for nested_value in nested_mapping.values(): + found = walk(nested_value, depth + 1) + if found: + return found + elif isinstance(item, Sequence) and not isinstance(item, (str, bytes, bytearray)): + for nested_value in item: + found = walk(nested_value, depth + 1) + if found: + return found + return None + + return walk(value, 0) + + +def _candidate_output_texts(value: Any) -> list[str]: + if isinstance(value, str): + return [value] + mapping = _mapping_or_attributes(value) + if mapping is None: + return [] + values: list[str] = [] + for key in ("output", "text", "content", "message", "final_output"): + text = _as_text(mapping.get(key)) + if text: + values.append(text) + return values + + +def _mapping_or_attributes(value: Any) -> dict[str, Any] | None: + if isinstance(value, Mapping): + return dict(value) + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + try: + dumped = to_dict() + except Exception: + dumped = None + if isinstance(dumped, Mapping): + return dict(dumped) + attributes = { + key: getattr(value, key) + for key in ( + *_THOUGHT_KEYS, + "output", + "text", + "content", + "message", + "final_output", + "messages", + "role", + "type", + ) + if getattr(value, key, None) is not None + } + return attributes or None + + +def _first_text_for_keys(value: Mapping[str, Any], keys: Sequence[str]) -> str | None: + for key in keys: + text = _as_text(value.get(key)) + if text: + return text + return None + + +def _as_text(value: Any) -> str | None: + if isinstance(value, str): + return value.strip() or None + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + parts: list[str] = [] + for item in value: + if isinstance(item, str): + parts.append(item) + continue + mapping = _mapping_or_attributes(item) + if mapping is None: + continue + text = _first_text_for_keys(mapping, ("text", "content", "summary")) + if text: + parts.append(text) + joined = "\n".join(part.strip() for part in parts if part.strip()) + return joined or None + return None + + +def _extract_tagged_thought(text: str) -> str | None: + parts = [match.group("body").strip() for match in _THOUGHT_TAG_RE.finditer(text)] + return "\n\n".join(part for part in parts if part) or None + + +def _extract_react_thought(text: str) -> str | None: + match = _REACT_THOUGHT_RE.search(text) + return match.group("body").strip() or None if match else None + + +def _serialize_observation(value: Any) -> str: + if isinstance(value, str): + return value.strip() + try: + return json.dumps(value, ensure_ascii=False, sort_keys=True) + except (TypeError, ValueError): + return str(value) + + +def _clip(value: str, limit: int) -> str: + cleaned = value.strip() + bounded = max(1, int(limit)) + if len(cleaned) <= bounded: + return cleaned + return cleaned[:bounded] + + +def _escape_marker_text(value: str) -> str: + return re.sub( + r"", + lambda match: match.group(0).replace("<", "<").replace(">", ">"), + value, + flags=re.IGNORECASE, + ) + + +def _chat_completions_url(base_url: str) -> str: + normalized = base_url.rstrip("/") + if normalized.endswith("/chat/completions"): + return normalized + return f"{normalized}/chat/completions" + + +def _clean_aligned_thought(value: str) -> str: + cleaned = value.strip() + tagged = _extract_tagged_thought(cleaned) + if tagged and _THOUGHT_TAG_RE.fullmatch(cleaned): + return tagged + return re.sub(r"^\s*Thought\s*:\s*", "", cleaned, count=1, flags=re.IGNORECASE).strip() + + +__all__ = [ + "AlignmentContext", + "ThoughtAlignerClient", + "ThoughtAlignmentError", + "ThoughtObservation", + "build_alignment_context", + "extract_thought", +] diff --git a/src/shared/schemas/events.py b/src/shared/schemas/events.py index f5cc83a..6b8b2b8 100644 --- a/src/shared/schemas/events.py +++ b/src/shared/schemas/events.py @@ -403,6 +403,8 @@ def _coerce_llm_output(value: Any) -> LLMOutput: return LLMOutput(output=_coerce_text(value)) thought = data.get("thought") + if thought is None: + thought = _nested_reasoning_value(data) final_output = data.get("final_output") output = data.get("output") if output is None: @@ -441,7 +443,18 @@ def _llm_output_fields(value: Any) -> dict[str, Any] | None: if data is None: attrs = { key: getattr(value, key) - for key in ("output", "text", "content", "message", "thought", "final_output") + for key in ( + "output", + "text", + "content", + "message", + "thought", + "reasoning_content", + "reasoning", + "thinking", + "analysis", + "final_output", + ) if getattr(value, key, None) is not None } data = attrs or None @@ -449,7 +462,64 @@ def _llm_output_fields(value: Any) -> dict[str, Any] | None: if not data: return None - recognized = ("output", "text", "content", "message", "thought", "final_output") + recognized = ( + "output", + "text", + "content", + "message", + "thought", + "reasoning_content", + "reasoning", + "thinking", + "analysis", + "final_output", + ) if not any(key in data for key in recognized): return None return data + + +_REASONING_KEYS = ( + "reasoning_content", + "reasoningContent", + "reasoning", + "thinking", + "plan", + "analysis", +) + + +def _nested_reasoning_value(value: Any, depth: int = 0) -> str | None: + if depth > 5: + return None + if isinstance(value, dict): + for key in _REASONING_KEYS: + text = _reasoning_text(value.get(key)) + if text: + return text + for item in value.values(): + nested = _nested_reasoning_value(item, depth + 1) + if nested: + return nested + elif isinstance(value, list): + for item in value: + nested = _nested_reasoning_value(item, depth + 1) + if nested: + return nested + return None + + +def _reasoning_text(value: Any) -> str | None: + if isinstance(value, str): + return value.strip() or None + if isinstance(value, list): + parts: list[str] = [] + for item in value: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict): + text = item.get("text") or item.get("content") or item.get("summary") + if isinstance(text, str) and text.strip(): + parts.append(text) + return "\n".join(parts).strip() or None + return None diff --git a/tests/test_attach_adapters.py b/tests/test_attach_adapters.py index 791e7a6..fd7a39a 100644 --- a/tests/test_attach_adapters.py +++ b/tests/test_attach_adapters.py @@ -10,8 +10,10 @@ from agentguard.adapters.agent import langgraph as langgraph_adapter from agentguard.adapters.agent import openai_agents as openai_agents_adapter from agentguard.adapters.agent.base import BaseAgentAdapter +from agentguard.adapters.agent.patching import make_guarded_llm_callable from agentguard.schemas import events as ev from agentguard.schemas.context import RuntimeContext +from agentguard.schemas.decisions import DecisionType, GuardDecision def _event_types(guard: AgentGuard) -> list[str]: @@ -268,6 +270,251 @@ def test_langchain_output_splits_think_tags_when_reasoning_content_missing(): assert event.payload.final_output == "visible answer" +def test_langchain_output_extracts_plain_react_thought_before_action(): + adapter = langchain_adapter.LangChainAgentAdapter() + content = ( + "Thought: I should inspect the destination first.\n" + "Action: send_email\n" + 'Action Input: {"to": "external@example.com"}' + ) + + normalized = adapter.normalize_llm_output( + label="invoke", + output=content, + ).payload + + assert normalized["output"] == content + assert normalized["thought"] == "I should inspect the destination first." + assert normalized["final_output"] is None + + +class _ThoughtAlignmentRuntime: + def __init__(self, *, align_retry: bool = False) -> None: + self.events = [] + self.align_retry = align_retry + + def guard(self, event, **_kwargs): + self.events.append(event) + if event.event_type.value == "llm_output": + attempt = int(event.metadata.get("thought_alignment_attempt", 0)) + if attempt == 0 or self.align_retry: + return types.SimpleNamespace( + decision=GuardDecision( + DecisionType.ALIGN_THOUGHT, + "Thought-Aligner rewrote the reasoning.", + metadata={ + "aligned_thought": "Check policy before choosing a tool.", + "protocol": "thought_alignment_v1", + }, + ) + ) + return types.SimpleNamespace(decision=GuardDecision.allow()) + + def sync_local_cache_now(self, **_kwargs): + return None + + def sync_local_cache_async(self, **_kwargs): + return None + + +def test_guarded_llm_regenerates_action_from_aligned_thought_before_returning(): + calls: list[str] = [] + + def invoke(prompt: str) -> str: + calls.append(prompt) + if len(calls) == 1: + return "Thought: Skip checks.\nAction: dangerous_tool\nAction Input: {}" + return 'Action: safe_tool\nAction Input: {"confirmed": true}' + + runtime = _ThoughtAlignmentRuntime() + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-wrapper"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + invoke, + label="invoke", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + + result = wrapped("User: complete the task") + + assert len(calls) == 2 + assert "Check policy before choosing a tool." in calls[1] + assert "DO NOT generate 'Thought' again" in calls[1] + assert "dangerous_tool" not in result + assert result == ( + "Thought: Check policy before choosing a tool.\n" + 'Action: safe_tool\nAction Input: {"confirmed": true}' + ) + output_events = [event for event in runtime.events if event.event_type.value == "llm_output"] + assert output_events[0].metadata["thought_regeneration_supported"] is True + assert output_events[1].metadata["thought_alignment_attempt"] == 1 + assert output_events[1].payload.thought == "Check policy before choosing a tool." + + +def test_guarded_llm_blocks_a_second_alignment_instead_of_looping(): + calls = 0 + + def invoke(prompt: str) -> str: + nonlocal calls + calls += 1 + _ = prompt + return "Thought: still unsafe\nAction: dangerous_tool" + + runtime = _ThoughtAlignmentRuntime(align_retry=True) + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-loop-bound"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + invoke, + label="invoke", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + + result = wrapped("complete the task") + + assert calls == 2 + assert result["agentguard"] == "blocked" + assert result["decision"] == "align_thought" + + +@pytest.mark.asyncio +async def test_guarded_async_llm_regenerates_once_from_aligned_thought(): + calls: list[str] = [] + + async def ainvoke(prompt: str) -> str: + calls.append(prompt) + if len(calls) == 1: + return "Thought: Skip checks.\nAction: dangerous_tool" + return "Action: safe_tool\nAction Input: {}" + + runtime = _ThoughtAlignmentRuntime() + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-async-wrapper"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + ainvoke, + label="ainvoke", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + + result = await wrapped("complete the task") + + assert len(calls) == 2 + assert result.startswith("Thought: Check policy before choosing a tool.\nAction: safe_tool") + + +def test_guarded_llm_regeneration_supports_agent_scratchpad_without_mutating_caller(): + calls: list[dict[str, str]] = [] + + def predict(*, input: str, agent_scratchpad: str) -> str: + calls.append({"input": input, "agent_scratchpad": agent_scratchpad}) + if len(calls) == 1: + return "Thought: Skip checks.\nAction: dangerous_tool" + return "Action: safe_tool\nAction Input: {}" + + runtime = _ThoughtAlignmentRuntime() + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-scratchpad"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + predict, + label="predict", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + original_scratchpad = "Thought: locate document\nObservation: found" + + result = wrapped(input="send report", agent_scratchpad=original_scratchpad) + + assert calls[0]["agent_scratchpad"] == original_scratchpad + assert calls[1]["agent_scratchpad"].startswith(original_scratchpad) + assert "Check policy before choosing a tool." in calls[1]["agent_scratchpad"] + assert result.startswith("Thought: Check policy before choosing a tool.\nAction: safe_tool") + + +def test_guarded_llm_regeneration_supports_message_dicts_and_structured_output(): + calls: list[list[dict[str, str]]] = [] + + def invoke(messages: list[dict[str, str]]) -> dict: + calls.append(messages) + if len(calls) == 1: + return { + "content": "", + "additional_kwargs": {"reasoning_content": "Skip checks."}, + "tool_calls": [{"name": "dangerous_tool", "args": {}}], + } + return { + "content": "", + "additional_kwargs": {}, + "tool_calls": [{"name": "safe_tool", "args": {}}], + } + + runtime = _ThoughtAlignmentRuntime() + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-messages"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + invoke, + label="invoke", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + original_messages = [{"role": "user", "content": "complete the task"}] + + result = wrapped(original_messages) + + assert original_messages == [{"role": "user", "content": "complete the task"}] + assert len(calls) == 2 + assert calls[1][-2] == { + "role": "assistant", + "content": "Thought: Check policy before choosing a tool.", + } + assert calls[1][-1]["role"] == "user" + assert result["tool_calls"] == [{"name": "safe_tool", "args": {}}] + assert result["thought"] == "Check policy before choosing a tool." + assert result["additional_kwargs"]["reasoning_content"] == ( + "Check policy before choosing a tool." + ) + + +def test_guarded_llm_blocks_alignment_when_request_shape_cannot_be_reinjected(): + calls = 0 + + def invoke(request: object) -> str: + nonlocal calls + calls += 1 + _ = request + return "Thought: unsafe\nAction: dangerous_tool" + + runtime = _ThoughtAlignmentRuntime() + guard = types.SimpleNamespace( + context=RuntimeContext(session_id="thought-aligner-unsupported-shape"), + runtime=runtime, + ) + wrapped = make_guarded_llm_callable( + guard, + invoke, + label="invoke", + normalizer=langchain_adapter.LangChainAgentAdapter(), + ) + + result = wrapped(object()) + + assert calls == 1 + assert result["agentguard"] == "blocked" + output = next(event for event in runtime.events if event.event_type.value == "llm_output") + assert output.metadata["thought_regeneration_supported"] is False + + def test_attach_langchain_patches_agent_executor_llm_chain_model(): class Tool: name = "lookup" diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 8b2fb87..7ed1ca3 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -158,7 +158,7 @@ def test_llm_output_does_not_parse_thought_only_output(): assert event.payload.final_output is None -def test_llm_output_ignores_nested_reasoning_fields(): +def test_llm_output_extracts_nested_reasoning_alias(): ctx = RuntimeContext(session_id="s") event = ev.llm_output( ctx, @@ -169,5 +169,19 @@ def test_llm_output_ignores_nested_reasoning_fields(): ) assert event.payload.output == "visible answer" - assert event.payload.thought is None + assert event.payload.thought == "hidden reasoning" assert event.payload.final_output is None + + +def test_llm_output_prefers_explicit_thought_over_nested_reasoning_alias(): + ctx = RuntimeContext(session_id="s") + event = ev.llm_output( + ctx, + { + "text": "visible answer", + "thought": "explicit thought", + "additional_kwargs": {"reasoning_content": "nested reasoning"}, + }, + ) + + assert event.payload.thought == "explicit thought" diff --git a/tests/test_thought_aligner.py b/tests/test_thought_aligner.py new file mode 100644 index 0000000..86008e1 --- /dev/null +++ b/tests/test_thought_aligner.py @@ -0,0 +1,385 @@ +from __future__ import annotations + +import json +from typing import Any + +import pytest +from backend.runtime.manager import RuntimeManager +from backend.runtime.plugins.llm_after.thought_aligner import ThoughtAlignerPlugin +from backend.runtime.thought_alignment import ( + ThoughtAlignerClient, + ThoughtAlignmentError, + build_alignment_context, + extract_thought, +) + +from shared.schemas import events as ev +from shared.schemas.context import RuntimeContext +from shared.schemas.decisions import DecisionType + + +def _ctx(**metadata: Any) -> RuntimeContext: + return RuntimeContext(session_id="thought-aligner-test", metadata=metadata) + + +def test_context_prefers_explicit_instruction_and_current_thought() -> None: + context = _ctx(instruction="context instruction") + event = ev.llm_output( + context, + {"output": "Action: delete", "thought": "current thought"}, + instruction="event instruction", + thought_regeneration_supported=True, + ) + + alignment = build_alignment_context(event, context, []) + + assert alignment is not None + assert alignment.instruction == "event instruction" + assert alignment.thought == "current thought" + assert alignment.history == () + + +def test_context_extracts_instruction_from_nested_langchain_input() -> None: + context = _ctx() + llm_input = ev.llm_input( + context, + { + "input": [ + {"role": "system", "content": "You are an agent."}, + {"role": "human", "content": "Send the weekly report."}, + ] + }, + ) + event = ev.llm_output( + context, + "Thought: I should email it.\nAction: send_email", + thought_regeneration_supported=True, + ) + + alignment = build_alignment_context(event, context, [llm_input]) + + assert alignment is not None + assert alignment.instruction == "Send the weekly report." + assert alignment.thought == "I should email it." + + +def test_context_extracts_last_human_section_from_rendered_prompt() -> None: + context = _ctx() + llm_input = ev.llm_input( + context, + "System: You are an agent.\nHuman: Read the invoice safely.\nAI:", + ) + event = ev.llm_output( + context, + "Inspect the file first.\nAction: read_file", + thought_regeneration_supported=True, + ) + + alignment = build_alignment_context(event, context, [llm_input]) + + assert alignment is not None + assert alignment.instruction == "Read the invoice safely." + assert alignment.thought == "Inspect the file first." + + +def test_context_formats_completed_thought_observation_history() -> None: + context = _ctx() + trace = [ + ev.user_input(context, "Find the document, then summarize it."), + ev.llm_output( + context, + {"thought": "Search for the document.", "output": "Action: search"}, + ), + ev.tool_invoke(context, "search", {"query": "document"}), + ev.tool_result(context, "search", {"id": 7, "title": "Quarterly report"}), + ] + event = ev.llm_output( + context, + {"thought": "Open result 7.", "output": "Action: open"}, + thought_regeneration_supported=True, + ) + + alignment = build_alignment_context(event, context, trace) + + assert alignment is not None + assert len(alignment.history) == 1 + assert alignment.history[0].thought == "Search for the document." + assert alignment.history[0].observation == '{"id": 7, "title": "Quarterly report"}' + assert alignment.formatted_instruction == ( + "Find the document, then summarize it.\n" + " Search for the document. \n" + ' {"id": 7, "title": "Quarterly report"} ' + ) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ({"thought": "explicit", "reasoning_content": "fallback"}, "explicit"), + ({"additional_kwargs": {"reasoning_content": "hidden"}}, "hidden"), + ({"metadata": {"thinking": "structured thinking"}}, "structured thinking"), + ("tagged thought\nAction: lookup", "tagged thought"), + ("Thought: use lookup\nAction: lookup\nAction Input: {}", "use lookup"), + ("Final Answer: done", None), + ], +) +def test_extract_thought_supports_explicit_structured_tagged_and_react_forms( + value: Any, + expected: str | None, +) -> None: + assert extract_thought(value) == expected + + +def test_context_returns_none_when_provider_exposes_no_thought() -> None: + context = _ctx() + event = ev.llm_output( + context, + {"output": "Action: lookup", "final_output": "Action: lookup"}, + thought_regeneration_supported=True, + ) + + assert build_alignment_context(event, context, []) is None + + +class _Response: + def __init__(self, payload: dict[str, Any]) -> None: + self._body = json.dumps(payload).encode("utf-8") + + def __enter__(self) -> _Response: + return self + + def __exit__(self, *_: Any) -> None: + return None + + def read(self) -> bytes: + return self._body + + +def test_thought_aligner_client_uses_dedicated_openai_compatible_endpoint() -> None: + seen: dict[str, Any] = {} + + def opener(request: Any, *, timeout: float) -> _Response: + seen["url"] = request.full_url + seen["authorization"] = request.get_header("Authorization") + seen["body"] = json.loads(request.data.decode("utf-8")) + seen["timeout"] = timeout + return _Response({"choices": [{"message": {"content": "safe thought"}}]}) + + client = ThoughtAlignerClient( + base_url="https://aligner.example/v1/", + api_key="test-only-key", + model="thought-aligner-7b", + timeout_s=17.0, + opener=opener, + ) + + result = client.align("instruction with history", "unsafe thought") + + assert result == "safe thought" + assert seen == { + "url": "https://aligner.example/v1/chat/completions", + "authorization": "Bearer test-only-key", + "timeout": 17.0, + "body": { + "model": "thought-aligner-7b", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + { + "role": "user", + "content": ( + "BEGINNING OF CONVERSATION: USER: Edit the following " + "Instruction-Thought pair to make it more helpful and harmless. : " + "instruction with history | unsafe thought ASSISTANT:" + ), + }, + ], + }, + } + + +@pytest.mark.parametrize( + "kwargs", + [ + {"base_url": "", "api_key": "key", "model": "model"}, + {"base_url": "https://aligner.example/v1", "api_key": "", "model": "model"}, + {"base_url": "https://aligner.example/v1", "api_key": "key", "model": ""}, + ], +) +def test_thought_aligner_client_rejects_missing_configuration(kwargs: dict[str, str]) -> None: + client = ThoughtAlignerClient(**kwargs) + + with pytest.raises(ThoughtAlignmentError, match="configuration"): + client.align("instruction", "thought") + + +def test_thought_aligner_client_rejects_malformed_response() -> None: + client = ThoughtAlignerClient( + base_url="https://aligner.example/v1", + api_key="test-only-key", + model="thought-aligner-7b", + opener=lambda *_args, **_kwargs: _Response({"choices": []}), + ) + + with pytest.raises(ThoughtAlignmentError, match="response"): + client.align("instruction", "thought") + + +class _FakeAligner: + def __init__(self, result: str = "safe thought", error: Exception | None = None) -> None: + self.result = result + self.error = error + self.calls: list[tuple[str, str]] = [] + + def align(self, instruction: str, thought: str) -> str: + self.calls.append((instruction, thought)) + if self.error is not None: + raise self.error + return self.result + + +def _thought_event( + context: RuntimeContext, + *, + supported: bool = True, + attempt: int = 0, +) -> Any: + return ev.llm_output( + context, + {"thought": "unsafe thought", "output": "Action: dangerous"}, + instruction="Complete the task safely.", + thought_regeneration_supported=supported, + thought_alignment_attempt=attempt, + ) + + +def test_thought_aligner_plugin_returns_alignment_directive_without_sensitive_prompt() -> None: + context = _ctx() + aligner = _FakeAligner("safe thought") + plugin = ThoughtAlignerPlugin(aligner=aligner) + + result = plugin.check(_thought_event(context), context, []) + + assert result.decision_candidate is not None + assert result.decision_candidate.decision_type == DecisionType.ALIGN_THOUGHT + assert result.decision_candidate.metadata == { + "aligned_thought": "safe thought", + "protocol": "thought_alignment_v1", + } + assert result.risk_signals == ["thought_alignment_applied"] + assert aligner.calls == [("Complete the task safely.", "unsafe thought")] + + +def test_thought_aligner_plugin_is_noop_for_missing_thought_or_retry() -> None: + context = _ctx() + aligner = _FakeAligner() + plugin = ThoughtAlignerPlugin(aligner=aligner) + no_thought = ev.llm_output( + context, + {"output": "Final Answer: done"}, + thought_regeneration_supported=True, + ) + + missing_result = plugin.check(no_thought, context, []) + retry_result = plugin.check(_thought_event(context, attempt=1), context, []) + + assert missing_result.decision_candidate is None + assert retry_result.decision_candidate is None + assert aligner.calls == [] + + +def test_thought_aligner_plugin_denies_old_or_unsupported_client_before_model_call() -> None: + context = _ctx() + aligner = _FakeAligner() + plugin = ThoughtAlignerPlugin(aligner=aligner) + + result = plugin.check(_thought_event(context, supported=False), context, []) + + assert result.decision_candidate is not None + assert result.decision_candidate.decision_type == DecisionType.DENY + assert "regeneration" in result.decision_candidate.reason.lower() + assert aligner.calls == [] + + +def test_thought_aligner_plugin_is_noop_when_thought_is_unchanged() -> None: + context = _ctx() + aligner = _FakeAligner(" unsafe thought \n") + plugin = ThoughtAlignerPlugin(aligner=aligner) + + result = plugin.check(_thought_event(context), context, []) + + assert result.decision_candidate is None + assert result.metadata["thought_alignment"] == "unchanged" + + +@pytest.mark.parametrize( + ("failure_mode", "expected"), + [("allow", None), ("deny", DecisionType.DENY)], +) +def test_thought_aligner_plugin_has_configurable_model_failure_mode( + failure_mode: str, + expected: DecisionType | None, +) -> None: + context = _ctx() + plugin = ThoughtAlignerPlugin( + aligner=_FakeAligner(error=ThoughtAlignmentError("endpoint unavailable")), + failure_mode=failure_mode, + ) + + result = plugin.check(_thought_event(context), context, []) + + if expected is None: + assert result.decision_candidate is None + else: + assert result.decision_candidate is not None + assert result.decision_candidate.decision_type == expected + assert result.risk_signals == ["thought_alignment_error"] + + +def test_runtime_manager_returns_serialized_alignment_decision_before_action() -> None: + aligner = _FakeAligner("Verify recipient authorization before sending.") + manager = RuntimeManager(enable_session_health_monitor=False) + manager.plugins.add(ThoughtAlignerPlugin(aligner=aligner), phase="llm_after") + context = _ctx() + previous_input = ev.user_input(context, "Send the report to the authorized recipient.") + current = ev.llm_output( + context, + "Thought: Send it immediately.\nAction: send_email\nAction Input: {}", + thought_regeneration_supported=True, + thought_alignment_attempt=0, + ) + + response = manager.decide( + { + "request_id": "thought-aligner-runtime", + "context": context.to_dict(), + "current_event": current.to_dict(), + "trajectory_window": [previous_input.to_dict()], + "local_signals": [], + } + ) + + assert response["decision"]["decision_type"] == "align_thought" + assert response["decision"]["metadata"]["aligned_thought"] == ( + "Verify recipient authorization before sending." + ) + assert "Send it immediately" not in json.dumps(response["decision"]) + assert aligner.calls == [ + ("Send the report to the authorized recipient.", "Send it immediately.") + ] + + +def test_example_config_loads_credentials_from_server_environment(monkeypatch) -> None: + monkeypatch.setenv("THOUGHT_ALIGNER_BASE_URL", "https://aligner.example/v1") + monkeypatch.setenv("THOUGHT_ALIGNER_API_KEY", "test-only-key") + monkeypatch.setenv("THOUGHT_ALIGNER_MODEL", "thought-aligner-test") + + manager = RuntimeManager( + plugin_config="config/plugins.thought-aligner.example.json", + enable_session_health_monitor=False, + ) + plugin = next(item for item in manager.plugins.plugins if item.name == "thought_aligner") + + assert plugin.base_url == "https://aligner.example/v1" + assert plugin.api_key == "test-only-key" + assert plugin.model == "thought-aligner-test" + assert plugin.failure_mode == "deny"