Merged
[fix]修复AgentContext.template_vars未渲染到主Agent的System Prompt中问题 #44
yuliangbinAdd on May 27
[fix]修复AgentContext.template_vars未渲染到主Agent的System Prompt中问题 #44
Merged
Merged fromyu_liangbin/msagentfix_templateintoAscend/msagentmaster
Total: 2 files changed+204-1
| @@ -21,6 +21,7 @@ | |||
| 21 | from __future__ import annotations | 21 | from __future__ import annotations |
| 22 | 22 | ||
| 23 | import logging | 23 | import logging |
| 24 | +import string | ||
| 24 | import tempfile | 25 | import tempfile |
| 25 | from fnmatch import fnmatch | 26 | from fnmatch import fnmatch |
| 26 | from importlib import import_module | 27 | from importlib import import_module |
| @@ -75,6 +76,79 @@ class _ToolPatternFilterMiddleware(AgentMiddleware[Any, Any, Any]): | |||
| 75 | return await handler(request) | 76 | return await handler(request) |
| 76 | 77 | ||
| 77 | 78 | ||
| 79 | +class _SystemMessageMiddleware(AgentMiddleware[Any, Any, Any]): | ||
| 80 | + """Populate system message placeholders from runtime AgentContext.template_vars.""" | ||
| 81 | + | ||
| 82 | + class _SafeTemplateFormatter(string.Formatter): | ||
| 83 | + """String formatter that leaves unknown placeholders unchanged.""" | ||
| 84 | + | ||
| 85 | + def __init__(self, context: dict[str, Any]) -> None: | ||
| 86 | + super().__init__() | ||
| 87 | + self._context = context | ||
| 88 | + | ||
| 89 | + def get_value(self, key: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: | ||
| 90 | + if isinstance(key, str): | ||
| 91 | + if key in self._context: | ||
| 92 | + return self._context[key] | ||
| 93 | + return "{" + key + "}" | ||
| 94 | + return super().get_value(key, args, kwargs) | ||
| 95 | + | ||
| 96 | + @classmethod | ||
| 97 | + def _safe_render_templates(cls, data: Any, context: dict[str, Any] | None) -> Any: | ||
| 98 | + formatter = cls._SafeTemplateFormatter(context or {}) | ||
| 99 | + if isinstance(data, str): | ||
| 100 | + return cls._safe_render_text(data, formatter) | ||
| 101 | + if isinstance(data, list): | ||
| 102 | + return [cls._render_text_block(item, formatter) for item in data] | ||
| 103 | + return data | ||
| 104 | + | ||
| 105 | + @staticmethod | ||
| 106 | + def _safe_render_text(text: str, formatter: string.Formatter) -> str: | ||
| 107 | + try: | ||
| 108 | + return formatter.vformat(text, (), {}) | ||
| 109 | + except ValueError: | ||
| 110 | + return text | ||
| 111 | + | ||
| 112 | + @classmethod | ||
| 113 | + def _render_text_block(cls, item: Any, formatter: string.Formatter) -> Any: | ||
| 114 | + if not isinstance(item, dict) or item.get("type") != "text" or not isinstance(item.get("text"), str): | ||
| 115 | + return item | ||
| 116 | + | ||
| 117 | + rendered_text = cls._safe_render_text(item["text"], formatter) | ||
| 118 | + if rendered_text == item["text"]: | ||
| 119 | + return item | ||
| 120 | + | ||
| 121 | + rendered_item = dict(item) | ||
| 122 | + rendered_item["text"] = rendered_text | ||
| 123 | + return rendered_item | ||
| 124 | + | ||
| 125 | + @staticmethod | ||
| 126 | + def _render_request_system_message(request): | ||
| 127 | + system_message = getattr(request, "system_message", None) | ||
| 128 | + runtime = getattr(request, "runtime", None) | ||
| 129 | + context = getattr(runtime, "context", None) if runtime is not None else None | ||
| 130 | + template_vars = getattr(context, "template_vars", None) if context is not None else None | ||
| 131 | + if system_message is None or not template_vars: | ||
| 132 | + return request | ||
| 133 | + | ||
| 134 | + rendered_content = _SystemMessageMiddleware._safe_render_templates( | ||
| 135 | + system_message.content, | ||
| 136 | + template_vars, | ||
| 137 | + ) | ||
| 138 | + if rendered_content == system_message.content: | ||
| 139 | + return request | ||
| 140 | + | ||
| 141 | + return request.override(system_message=system_message.model_copy(update={"content": rendered_content})) | ||
| 142 | + | ||
| 143 | + def wrap_model_call(self, request, handler): | ||
| 144 | + request = self._render_request_system_message(request) | ||
| 145 | + return handler(request) | ||
| 146 | + | ||
| 147 | + async def awrap_model_call(self, request, handler): | ||
| 148 | + request = self._render_request_system_message(request) | ||
| 149 | + return await handler(request) | ||
| 150 | + | ||
| 151 | + | ||
| 78 | class AgentFactory: | 152 | class AgentFactory: |
| 79 | """Factory for creating deepagents-based graphs.""" | 153 | """Factory for creating deepagents-based graphs.""" |
| 80 | 154 | ||
| @@ -220,6 +294,7 @@ class AgentFactory: | |||
| 220 | tool_token_limit_before_evict=tool_output_max_tokens, | 294 | tool_token_limit_before_evict=tool_output_max_tokens, |
| 221 | ) | 295 | ) |
| 222 | ) | 296 | ) |
| 297 | + middleware.append(_SystemMessageMiddleware()) | ||
| 223 | 298 | ||
| 224 | raw_system_prompt = config.prompt | 299 | raw_system_prompt = config.prompt |
| 225 | if isinstance(raw_system_prompt, list): | 300 | if isinstance(raw_system_prompt, list): |
| @@ -21,9 +21,10 @@ from __future__ import annotations | |||
| 21 | from types import SimpleNamespace | 21 | from types import SimpleNamespace |
| 22 | 22 | ||
| 23 | import pytest | 23 | import pytest |
| 24 | +from langchain_core.messages import SystemMessage | ||
| 24 | 25 | ||
| 25 | import msagent.agents.factory as factory_module | 26 | import msagent.agents.factory as factory_module |
| 26 | -from msagent.agents.factory import AgentFactory | 27 | +from msagent.agents.factory import AgentFactory, _SystemMessageMiddleware |
| 27 | 28 | ||
| 28 | 29 | ||
| 29 | class _DummyLLMFactory: | 30 | class _DummyLLMFactory: |
| @@ -46,6 +47,28 @@ class _DummyGraph: | |||
| 46 | pass | 47 | pass |
| 47 | 48 | ||
| 48 | 49 | ||
| 50 | +def _make_model_request( | ||
| 51 | + *, | ||
| 52 | + system_message: SystemMessage, | ||
| 53 | + template_vars: dict[str, object] | None, | ||
| 54 | +) -> SimpleNamespace: | ||
| 55 | + request = SimpleNamespace( | ||
| 56 | + system_message=system_message, | ||
| 57 | + runtime=SimpleNamespace(context=SimpleNamespace(template_vars=template_vars or {})), | ||
| 58 | + ) | ||
| 59 | + | ||
| 60 | + def _override(**overrides): | ||
| 61 | + data = { | ||
| 62 | + "system_message": request.system_message, | ||
| 63 | + "runtime": request.runtime, | ||
| 64 | + } | ||
| 65 | + data.update(overrides) | ||
| 66 | + return SimpleNamespace(**data, override=_override) | ||
| 67 | + | ||
| 68 | + request.override = _override | ||
| 69 | + return request | ||
| 70 | + | ||
| 71 | + | ||
| 49 | def _patch_deepagent_entrypoints(monkeypatch: pytest.MonkeyPatch) -> None: | 72 | def _patch_deepagent_entrypoints(monkeypatch: pytest.MonkeyPatch) -> None: |
| 50 | class _DummyBackend: | 73 | class _DummyBackend: |
| 51 | def __init__(self, root_dir: str, **kwargs): | 74 | def __init__(self, root_dir: str, **kwargs): |
| @@ -241,6 +264,111 @@ def test_agent_factory_filters_deepagents_default_tools_from_model_request() -> | |||
| 241 | assert {tool.name for tool in filtered} == {"get_skill", "msprof-mcp_ping"} | 264 | assert {tool.name for tool in filtered} == {"get_skill", "msprof-mcp_ping"} |
| 242 | 265 | ||
| 243 | 266 | ||
| 267 | +def test_system_message_middleware_renders_known_vars_and_preserves_unknown_placeholders() -> None: | ||
| 268 | + request = _make_model_request( | ||
| 269 | + system_message=SystemMessage( | ||
| 270 | + content="cwd={working_dir}; local={local_environment_context}; worker={worker}; rank={Rank_ID}" | ||
| 271 | + ), | ||
| 272 | + template_vars={ | ||
| 273 | + "working_dir": "/tmp/project", | ||
| 274 | + "local_environment_context": "GPU=Ascend", | ||
| 275 | + }, | ||
| 276 | + ) | ||
| 277 | + | ||
| 278 | + updated = _SystemMessageMiddleware._render_request_system_message(request) | ||
| 279 | + | ||
| 280 | + assert updated is not request | ||
| 281 | + assert str(updated.system_message.content) == ( | ||
| 282 | + "cwd=/tmp/project; local=GPU=Ascend; worker={worker}; rank={Rank_ID}" | ||
| 283 | + ) | ||
| 284 | + | ||
| 285 | + | ||
| 286 | +def test_system_message_middleware_leaves_request_unchanged_without_template_vars() -> None: | ||
| 287 | + request = _make_model_request( | ||
| 288 | + system_message=SystemMessage(content="cwd={working_dir}; worker={worker}"), | ||
| 289 | + template_vars={}, | ||
| 290 | + ) | ||
| 291 | + | ||
| 292 | + updated = _SystemMessageMiddleware._render_request_system_message(request) | ||
| 293 | + | ||
| 294 | + assert updated is request | ||
| 295 | + assert str(updated.system_message.content) == "cwd={working_dir}; worker={worker}" | ||
| 296 | + | ||
| 297 | + | ||
| 298 | +def test_system_message_middleware_renders_text_blocks_only() -> None: | ||
| 299 | + content = [ | ||
| 300 | + { | ||
| 301 | + "type": "text", | ||
| 302 | + "text": "cwd={working_dir}", | ||
| 303 | + "metadata": {"template": "{working_dir}"}, | ||
| 304 | + }, | ||
| 305 | + { | ||
| 306 | + "type": "image_url", | ||
| 307 | + "image_url": {"url": "https://example.com/{working_dir}.png"}, | ||
| 308 | + }, | ||
| 309 | + ] | ||
| 310 | + request = _make_model_request( | ||
| 311 | + system_message=SystemMessage(content=content), | ||
| 312 | + template_vars={"working_dir": "/workspace"}, | ||
| 313 | + ) | ||
| 314 | + | ||
| 315 | + updated = _SystemMessageMiddleware._render_request_system_message(request) | ||
| 316 | + | ||
| 317 | + assert updated is not request | ||
| 318 | + assert updated.system_message.content == [ | ||
| 319 | + { | ||
| 320 | + "type": "text", | ||
| 321 | + "text": "cwd=/workspace", | ||
| 322 | + "metadata": {"template": "{working_dir}"}, | ||
| 323 | + }, | ||
| 324 | + { | ||
| 325 | + "type": "image_url", | ||
| 326 | + "image_url": {"url": "https://example.com/{working_dir}.png"}, | ||
| 327 | + }, | ||
| 328 | + ] | ||
| 329 | + | ||
| 330 | + | ||
| 331 | +def test_system_message_middleware_preserves_system_message_metadata() -> None: | ||
| 332 | + request = _make_model_request( | ||
| 333 | + system_message=SystemMessage( | ||
| 334 | + content="cwd={working_dir}", | ||
| 335 | + additional_kwargs={"source": "runtime"}, | ||
| 336 | + response_metadata={"trace_id": "abc"}, | ||
| 337 | + name="system-name", | ||
| 338 | + id="system-id", | ||
| 339 | + ), | ||
| 340 | + template_vars={"working_dir": "/workspace"}, | ||
| 341 | + ) | ||
| 342 | + | ||
| 343 | + updated = _SystemMessageMiddleware._render_request_system_message(request) | ||
| 344 | + | ||
| 345 | + assert str(updated.system_message.content) == "cwd=/workspace" | ||
| 346 | + assert updated.system_message.additional_kwargs == {"source": "runtime"} | ||
| 347 | + assert updated.system_message.response_metadata == {"trace_id": "abc"} | ||
| 348 | + assert updated.system_message.name == "system-name" | ||
| 349 | + assert updated.system_message.id == "system-id" | ||
| 350 | + | ||
| 351 | + | ||
| 352 | +@pytest.mark.asyncio | ||
| 353 | +async def test_system_message_middleware_awrap_model_call_applies_rendering() -> None: | ||
| 354 | + middleware = _SystemMessageMiddleware() | ||
| 355 | + request = _make_model_request( | ||
| 356 | + system_message=SystemMessage(content="local={local_environment_context}; rank={Rank_ID}"), | ||
| 357 | + template_vars={"local_environment_context": "NPU=910B"}, | ||
| 358 | + ) | ||
| 359 | + | ||
| 360 | + captured = {} | ||
| 361 | + | ||
| 362 | + async def _handler(updated_request): | ||
| 363 | + captured["request"] = updated_request | ||
| 364 | + return "ok" | ||
| 365 | + | ||
| 366 | + result = await middleware.awrap_model_call(request, _handler) | ||
| 367 | + | ||
| 368 | + assert result == "ok" | ||
| 369 | + assert str(captured["request"].system_message.content) == "local=NPU=910B; rank={Rank_ID}" | ||
| 370 | + | ||
| 371 | + | ||
| 244 | @pytest.mark.asyncio | 372 | @pytest.mark.asyncio |
| 245 | async def test_agent_factory_maps_retry_config_to_deepagents_primitives( | 373 | async def test_agent_factory_maps_retry_config_to_deepagents_primitives( |
| 246 | monkeypatch, | 374 | monkeypatch, |