返回 ViMax
test_agent_prompt_builder.py
根目录 / tests / test_agent_prompt_builder.py
1 import tempfile
2 import unittest
3 from pathlib import Path
4
5 from agent_runtime.prompts import PromptBuilder
6 from agent_runtime.session_index import SessionIndex
7 from agent_runtime.tools import build_builtin_registry
8
9
10 class PromptBuilderTests(unittest.TestCase):
11 def test_prompt_injects_context_and_tool_manifest(self):
12 with tempfile.TemporaryDirectory() as tmp:
13 root = Path(tmp)
14 (root / "prompts").mkdir()
15 (root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
16 (root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
17 index = SessionIndex(root)
18 index.create(idea="cat")
19 registry = build_builtin_registry(root, index)
20 builder = PromptBuilder(root / "prompts", index, registry)
21 messages = builder.build_messages("start")
22 self.assertIn("Available tools", messages[0]["content"])
23 self.assertIn("当前 working_dir 尚未完成结构化文本文件", messages[0]["content"])
24 self.assertIn("read_file", messages[0]["content"])
25 trace = builder.trace(builder.build_parts("start"))
26 self.assertGreater(trace["total_estimated_tokens"], 0)
27
28
29 def test_prompt_injects_compacted_session_summary_as_reference(self):
30 with tempfile.TemporaryDirectory() as tmp:
31 root = Path(tmp)
32 (root / "prompts").mkdir()
33 (root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
34 (root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
35 index = SessionIndex(root)
36 record = index.create(idea="cat")
37 index.update_compaction(record["session_id"], {"summary": "## Reference Context Only\n- user wants moon cat", "compacted_message_count": 4, "preserved_message_count": 2})
38 registry = build_builtin_registry(root, index)
39 builder = PromptBuilder(root / "prompts", index, registry)
40 message = builder.build_messages("continue")[0]["content"]
41 self.assertIn("Session context summary", message)
42 self.assertIn("reference context only", message)
43 self.assertIn("user wants moon cat", message)
44 trace = builder.trace(builder.build_parts("continue"))
45 self.assertGreater(trace["totals"]["dynamic_tokens"], 0)
46
47 def test_prompt_treats_novel_text_artifacts_as_text_stage_complete(self):
48 with tempfile.TemporaryDirectory() as tmp:
49 root = Path(tmp)
50 (root / "prompts").mkdir()
51 (root / "prompts" / "agent.md").write_text("agent rules", encoding="utf-8")
52 (root / "prompts" / "workflow.md").write_text("workflow rules", encoding="utf-8")
53 index = SessionIndex(root)
54 record = index.create(idea="novel")
55 session_root = root / record["working_dir"] / "novel2video"
56 (session_root / "novel").mkdir(parents=True)
57 (session_root / "novel" / "novel_compressed.txt").write_text("compressed", encoding="utf-8")
58 (session_root / "events").mkdir()
59 (session_root / "events" / "event_0.json").write_text("{}", encoding="utf-8")
60 (session_root / "relevant_chunks" / "event_0").mkdir(parents=True)
61 (session_root / "relevant_chunks" / "event_0" / "chunk.txt").write_text("chunk", encoding="utf-8")
62 (session_root / "scenes" / "event_0").mkdir(parents=True)
63 (session_root / "scenes" / "event_0" / "scene_0.json").write_text("{}", encoding="utf-8")
64 (session_root / "global_information" / "characters" / "novel_level").mkdir(parents=True)
65 (session_root / "global_information" / "characters" / "novel_level" / "novel_characters_after_event_0.json").write_text("[]", encoding="utf-8")
66 registry = build_builtin_registry(root, index)
67 builder = PromptBuilder(root / "prompts", index, registry)
68 messages = builder.build_messages("continue")
69 self.assertIn("文本规划阶段已完成", messages[0]["content"])
70 self.assertIn("novel2video/events/event_*.json: present", messages[0]["content"])
71
71 lines PYTHON