| 1 | from types import SimpleNamespace |
| 2 | |
| 3 | import pytest |
| 4 | from langchain_core.messages import AIMessage, HumanMessage, SystemMessage |
| 5 | |
| 6 | from agent import Agent, LoopData |
| 7 | from helpers import extension, history |
| 8 | from extensions.python.system_prompt import _14_project_prompt as project_prompt |
| 9 | |
| 10 | |
| 11 | class _DummyLog: |
| 12 | def set_progress(self, _message: str) -> None: |
| 13 | return None |
| 14 | |
| 15 | |
| 16 | @pytest.mark.asyncio |
| 17 | async def test_prepare_prompt_places_protocol_before_history_and_extras(monkeypatch): |
| 18 | async def fake_call_extensions(extension_point: str, agent=None, **kwargs): |
| 19 | if extension_point == "message_loop_prompts_after": |
| 20 | loop_data = kwargs["loop_data"] |
| 21 | loop_data.protocol_persistent["project_instructions"] = "Project rule." |
| 22 | loop_data.extras_temporary["current_datetime"] = "Today." |
| 23 | |
| 24 | monkeypatch.setattr(extension, "call_extensions_async", fake_call_extensions) |
| 25 | monkeypatch.setattr(history.History, "_get_max_embeds", lambda self: 0) |
| 26 | |
| 27 | agent = object.__new__(Agent) |
| 28 | loop_data = LoopData() |
| 29 | agent.loop_data = loop_data |
| 30 | agent.context = SimpleNamespace(log=_DummyLog()) |
| 31 | agent.history = history.History(agent) |
| 32 | agent.data = {} |
| 33 | |
| 34 | agent.history.add_message(False, "User asks.") |
| 35 | agent.history.add_message(True, "Assistant answers.") |
| 36 | |
| 37 | async def get_system_prompt(_loop_data): |
| 38 | return ["System root."] |
| 39 | |
| 40 | def read_prompt(prompt_file: str, **kwargs) -> str: |
| 41 | if prompt_file == "agent.context.protocol.md": |
| 42 | return "[PROTOCOL]\n" + kwargs["protocol"] |
| 43 | if prompt_file == "agent.context.extras.md": |
| 44 | return "[EXTRAS]\n" + kwargs["extras"] |
| 45 | raise AssertionError(f"Unexpected prompt file: {prompt_file}") |
| 46 | |
| 47 | agent.get_system_prompt = get_system_prompt |
| 48 | agent.read_prompt = read_prompt |
| 49 | agent.set_data = lambda key, value: agent.data.__setitem__(key, value) |
| 50 | |
| 51 | prompt = await Agent.prepare_prompt(agent, loop_data) |
| 52 | |
| 53 | assert isinstance(prompt[0], SystemMessage) |
| 54 | assert prompt[0].content == "System root." |
| 55 | assert isinstance(prompt[1], HumanMessage) |
| 56 | assert str(prompt[1].content).startswith("[PROTOCOL]") |
| 57 | assert str(prompt[1].content).index("Project rule.") < str(prompt[1].content).index( |
| 58 | "User asks." |
| 59 | ) |
| 60 | assert isinstance(prompt[2], AIMessage) |
| 61 | assert prompt[2].content == "Assistant answers." |
| 62 | assert isinstance(prompt[3], HumanMessage) |
| 63 | assert str(prompt[3].content).startswith("[EXTRAS]") |
| 64 | assert "Today." in str(prompt[3].content) |
| 65 | assert '{"current_datetime":"Today."' in str(prompt[3].content) |
| 66 | |
| 67 | serialized_history = agent.history.serialize() |
| 68 | assert "Project rule." not in serialized_history |
| 69 | assert "Today." not in serialized_history |
| 70 | assert "protocol" not in serialized_history.lower() |
| 71 | assert loop_data.protocol_temporary == {} |
| 72 | assert loop_data.extras_temporary == {} |
| 73 | |
| 74 | class FakeResponsesModel: |
| 75 | def _convert_messages(self, messages): |
| 76 | role_by_type = {"system": "system", "human": "user", "ai": "assistant"} |
| 77 | return [ |
| 78 | {"role": role_by_type[message.type], "content": message.content} |
| 79 | for message in messages |
| 80 | ] |
| 81 | |
| 82 | input_items = Agent._responses_prompt_input_items( |
| 83 | agent, |
| 84 | FakeResponsesModel(), |
| 85 | prompt, |
| 86 | ) |
| 87 | assert input_items[1]["role"] == "user" |
| 88 | assert "[PROTOCOL]" in input_items[1]["content"] |
| 89 | assert "[EXTRAS]" in input_items[-1]["content"] |
| 90 | |
| 91 | |
| 92 | @pytest.mark.asyncio |
| 93 | async def test_project_prompt_moves_project_instructions_to_protocol(monkeypatch): |
| 94 | project_vars = { |
| 95 | "project_name": "Demo", |
| 96 | "project_description": "", |
| 97 | "project_instructions": "Project rule.", |
| 98 | "include_agents_md": True, |
| 99 | "project_path": "/a0/usr/projects/demo", |
| 100 | "project_git_url": "", |
| 101 | } |
| 102 | loop_data = LoopData() |
| 103 | |
| 104 | class FakeContext: |
| 105 | def get_data(self, key): |
| 106 | assert key == project_prompt.projects.CONTEXT_DATA_KEY_PROJECT |
| 107 | return "demo" |
| 108 | |
| 109 | class FakeAgent: |
| 110 | context = FakeContext() |
| 111 | |
| 112 | def read_prompt(self, prompt_file: str, **kwargs) -> str: |
| 113 | if prompt_file == "agent.system.projects.main.md": |
| 114 | return "project context may be active" |
| 115 | if prompt_file == "agent.system.projects.active.md": |
| 116 | return f"active project: {kwargs['project_path']}" |
| 117 | if prompt_file == "agent.protocol.projects.instructions.md": |
| 118 | return "protocol project instructions:\n" + kwargs["project_instructions"] |
| 119 | raise AssertionError(f"Unexpected prompt file: {prompt_file}") |
| 120 | |
| 121 | monkeypatch.setattr( |
| 122 | project_prompt.projects, |
| 123 | "build_system_prompt_vars", |
| 124 | lambda _name: project_vars, |
| 125 | ) |
| 126 | monkeypatch.setattr( |
| 127 | project_prompt.projects, |
| 128 | "build_agents_md_protocol", |
| 129 | lambda _name: "AGENTS path rule.", |
| 130 | ) |
| 131 | |
| 132 | prompt = await project_prompt.build_prompt.__wrapped__( # type: ignore[attr-defined] |
| 133 | FakeAgent(), |
| 134 | loop_data=loop_data, |
| 135 | ) |
| 136 | |
| 137 | assert "Project rule." not in prompt |
| 138 | assert "AGENTS path rule." not in prompt |
| 139 | assert list(loop_data.protocol_persistent) == [ |
| 140 | "agents_md_instructions", |
| 141 | "project_instructions", |
| 142 | ] |
| 143 | assert "AGENTS path rule." in loop_data.protocol_persistent["agents_md_instructions"] |
| 144 | assert "Project rule." in loop_data.protocol_persistent["project_instructions"] |
| 145 | |
| 146 | |
| 147 | @pytest.mark.asyncio |
| 148 | async def test_project_prompt_does_not_load_agents_md_without_project(monkeypatch): |
| 149 | loop_data = LoopData() |
| 150 | |
| 151 | class FakeContext: |
| 152 | def get_data(self, key): |
| 153 | assert key == project_prompt.projects.CONTEXT_DATA_KEY_PROJECT |
| 154 | return None |
| 155 | |
| 156 | class FakeAgent: |
| 157 | context = FakeContext() |
| 158 | |
| 159 | def read_prompt(self, prompt_file: str, **kwargs) -> str: |
| 160 | if prompt_file == "agent.system.projects.main.md": |
| 161 | return "project context may be active" |
| 162 | if prompt_file == "agent.system.projects.inactive.md": |
| 163 | return "no active project" |
| 164 | raise AssertionError(f"Unexpected prompt file: {prompt_file}") |
| 165 | |
| 166 | monkeypatch.setattr( |
| 167 | project_prompt.projects, |
| 168 | "build_agents_md_protocol", |
| 169 | lambda _name: (_ for _ in ()).throw(AssertionError("unexpected AGENTS load")), |
| 170 | ) |
| 171 | |
| 172 | await project_prompt.build_prompt.__wrapped__( # type: ignore[attr-defined] |
| 173 | FakeAgent(), |
| 174 | loop_data=loop_data, |
| 175 | ) |
| 176 | |
| 177 | assert "agents_md_instructions" not in loop_data.protocol_persistent |
| 178 | assert "project_instructions" not in loop_data.protocol_persistent |