| 1 | import contextlib |
| 2 | import io |
| 3 | import json |
| 4 | import logging |
| 5 | import tempfile |
| 6 | import unittest |
| 7 | from pathlib import Path |
| 8 | from types import SimpleNamespace |
| 9 | from unittest.mock import patch |
| 10 | |
| 11 | from interfaces import CharacterInEvent, CharacterInNovel, CharacterInScene, Event, Scene |
| 12 | from interfaces.environment import EnvironmentInScene |
| 13 | from agent_runtime.session_index import SessionIndex |
| 14 | from agent_runtime.tools import ToolRuntimeContext |
| 15 | from agent_runtime.vimax_adapters import ViMaxAdapters, _run_planning_step |
| 16 | from agents.global_information_planner import GlobalInformationPlanner, MergeCharactersAcrossScenesInEventResponse |
| 17 | from pipelines.novel2movie_pipeline import Novel2MoviePipeline |
| 18 | |
| 19 | |
| 20 | class FakeCompressor: |
| 21 | def split(self, novel_text): |
| 22 | return [novel_text] |
| 23 | |
| 24 | async def compress_single_novel_chunk(self, semaphore, index, novel_chunk): |
| 25 | return index, f"compressed {novel_chunk}" |
| 26 | |
| 27 | def aggregate(self, chunks): |
| 28 | return "\n".join(chunks) |
| 29 | |
| 30 | |
| 31 | class FakeEventExtractor: |
| 32 | def extract_next_event(self, novel_text, extracted_events): |
| 33 | return Event(index=len(extracted_events), is_last=True, description="Hero leaves home", process_chain=["Hero opens the door"]) |
| 34 | |
| 35 | |
| 36 | class FakeKnowledgeBase: |
| 37 | def similarity_search(self, process, k=10): |
| 38 | return [SimpleNamespace(page_content="Hero opens the old wooden door.")] |
| 39 | |
| 40 | |
| 41 | class FakeReranker: |
| 42 | async def __call__(self, documents, query, top_n): |
| 43 | return [(documents[0], 0.95)] |
| 44 | |
| 45 | |
| 46 | class FakeSceneExtractor: |
| 47 | async def get_next_scene(self, relevant_chunks, event, previous_scenes): |
| 48 | return Scene( |
| 49 | idx=len(previous_scenes), |
| 50 | is_last=True, |
| 51 | environment=EnvironmentInScene(slugline="INT. HOUSE - DAY", description="A quiet room."), |
| 52 | characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], |
| 53 | script="<Hero> opens the door.", |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | class FakeGlobalPlanner: |
| 58 | async def merge_characters_across_scenes_in_event(self, event_idx, scenes): |
| 59 | return [CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")] |
| 60 | |
| 61 | def merge_characters_to_existing_characters_in_novel(self, event_idx, existing_characters_in_novel, characters_in_event): |
| 62 | return [CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={event_idx: "Hero"}, static_features="adult")] |
| 63 | |
| 64 | |
| 65 | class FakeNovelPipeline: |
| 66 | def __init__(self, working_dir: Path): |
| 67 | self.working_dir = working_dir |
| 68 | |
| 69 | async def plan_text_artifacts(self, novel_text, user_requirement="", style="", progress=None, quiet=False): |
| 70 | if progress: |
| 71 | for stage in ["save_novel", "compress_novel", "extract_events", "retrieve_chunks", "extract_scenes", "merge_characters", "completed"]: |
| 72 | progress(stage, stage, {}) |
| 73 | novel = self.working_dir / "novel" |
| 74 | novel.mkdir(parents=True, exist_ok=True) |
| 75 | (novel / "novel.txt").write_text(novel_text, encoding="utf-8") |
| 76 | (novel / "novel_compressed.txt").write_text("compressed", encoding="utf-8") |
| 77 | events = self.working_dir / "events" |
| 78 | events.mkdir(parents=True, exist_ok=True) |
| 79 | (events / "event_0.json").write_text(json.dumps(Event(index=0, is_last=True, description="d", process_chain=["p"]).model_dump()), encoding="utf-8") |
| 80 | chunks = self.working_dir / "relevant_chunks" / "event_0" |
| 81 | chunks.mkdir(parents=True, exist_ok=True) |
| 82 | (chunks / "chunk_0-score_0.95.txt").write_text("chunk", encoding="utf-8") |
| 83 | scenes = self.working_dir / "scenes" / "event_0" |
| 84 | scenes.mkdir(parents=True, exist_ok=True) |
| 85 | scene = Scene(idx=0, is_last=True, environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"), characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], script="<Hero> walks.") |
| 86 | (scenes / "scene_0.json").write_text(json.dumps(scene.model_dump()), encoding="utf-8") |
| 87 | event_level = self.working_dir / "global_information" / "characters" / "event_level" |
| 88 | event_level.mkdir(parents=True, exist_ok=True) |
| 89 | event_char = CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult") |
| 90 | (event_level / "event_0_characters.json").write_text(json.dumps([event_char.model_dump()]), encoding="utf-8") |
| 91 | novel_level = self.working_dir / "global_information" / "characters" / "novel_level" |
| 92 | novel_level.mkdir(parents=True, exist_ok=True) |
| 93 | novel_char = CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={0: "Hero"}, static_features="adult") |
| 94 | (novel_level / "novel_characters_after_event_0.json").write_text(json.dumps([novel_char.model_dump()]), encoding="utf-8") |
| 95 | return {} |
| 96 | |
| 97 | |
| 98 | class FakeMergeChain: |
| 99 | async def ainvoke(self, messages): |
| 100 | return MergeCharactersAcrossScenesInEventResponse( |
| 101 | characters=[CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult")] |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | class FakeMergeChatModel: |
| 106 | def __or__(self, parser): |
| 107 | return FakeMergeChain() |
| 108 | |
| 109 | |
| 110 | class GlobalInformationPlannerCompatibilityTests(unittest.IsolatedAsyncioTestCase): |
| 111 | async def test_merge_event_characters_uses_scene_character_idx(self): |
| 112 | planner = GlobalInformationPlanner.__new__(GlobalInformationPlanner) |
| 113 | planner.chat_model = FakeMergeChatModel() |
| 114 | scene = Scene( |
| 115 | idx=0, |
| 116 | is_last=True, |
| 117 | environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"), |
| 118 | characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], |
| 119 | script="<Hero> walks.", |
| 120 | ) |
| 121 | characters = await planner.merge_characters_across_scenes_in_event(event_idx=0, scenes=[scene]) |
| 122 | self.assertEqual(characters[0].identifier_in_event, "Hero") |
| 123 | |
| 124 | |
| 125 | class PlanningStepOutputSuppressionTests(unittest.IsolatedAsyncioTestCase): |
| 126 | async def test_run_planning_step_suppresses_stdout_stderr_and_warnings(self): |
| 127 | async def noisy_step(): |
| 128 | print("NOISE_STDOUT") |
| 129 | logging.warning("NOISE_WARNING") |
| 130 | return "ok" |
| 131 | |
| 132 | stdout = io.StringIO() |
| 133 | stderr = io.StringIO() |
| 134 | with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): |
| 135 | result = await _run_planning_step("message", "stage", noisy_step(), runtime=None) |
| 136 | self.assertEqual(result, "ok") |
| 137 | self.assertNotIn("NOISE_STDOUT", stdout.getvalue()) |
| 138 | self.assertNotIn("NOISE_WARNING", stderr.getvalue()) |
| 139 | |
| 140 | |
| 141 | class Novel2MoviePlanningTests(unittest.IsolatedAsyncioTestCase): |
| 142 | async def test_plan_text_artifacts_writes_structured_text_and_progress(self): |
| 143 | with tempfile.TemporaryDirectory() as tmp: |
| 144 | pipeline = Novel2MoviePipeline( |
| 145 | novel_compressor=FakeCompressor(), |
| 146 | event_extractor=FakeEventExtractor(), |
| 147 | embeddings=SimpleNamespace(model="fake-embedding"), |
| 148 | rerank_model=FakeReranker(), |
| 149 | scene_extractor=FakeSceneExtractor(), |
| 150 | global_information_planner=FakeGlobalPlanner(), |
| 151 | image_generator=object(), |
| 152 | rewriter=object(), |
| 153 | script2video_pipeline=object(), |
| 154 | working_dir=tmp, |
| 155 | ) |
| 156 | events = [] |
| 157 | with patch("pipelines.novel2movie_pipeline.CacheBackedEmbeddings.from_bytes_store", return_value=object()), \ |
| 158 | patch("pipelines.novel2movie_pipeline.FAISS.from_texts", return_value=FakeKnowledgeBase()): |
| 159 | result = await pipeline.plan_text_artifacts("Hero opens a door.", progress=lambda stage, message, metadata=None: events.append(stage), quiet=True) |
| 160 | self.assertEqual(events, ["save_novel", "compress_novel", "extract_events", "retrieve_chunks", "extract_scenes", "merge_characters", "completed"]) |
| 161 | root = Path(tmp) |
| 162 | self.assertTrue((root / "novel" / "novel_compressed.txt").exists()) |
| 163 | self.assertTrue((root / "events" / "event_0.json").exists()) |
| 164 | self.assertTrue((root / "relevant_chunks" / "event_0" / "chunk_0-score_0.95.txt").exists()) |
| 165 | self.assertTrue((root / "scenes" / "event_0" / "scene_0.json").exists()) |
| 166 | self.assertTrue((root / "global_information" / "characters" / "novel_level" / "novel_characters_after_event_0.json").exists()) |
| 167 | self.assertFalse((root / "character_portraits").exists()) |
| 168 | self.assertFalse((root / "videos").exists()) |
| 169 | self.assertEqual(len(result["events"]), 1) |
| 170 | |
| 171 | |
| 172 | class FakeNovelRenderPipeline: |
| 173 | def __init__(self, working_dir: Path): |
| 174 | self.working_dir = Path(working_dir) |
| 175 | |
| 176 | async def render_video_artifacts(self, style, user_requirement="", progress=None, quiet=False): |
| 177 | if progress: |
| 178 | progress("novel_portraits_start", "portraits", {}) |
| 179 | progress("novel_scene_render_start", "scene", {"event_idx": 0, "scene_idx": 0}) |
| 180 | progress("novel_render_completed", "done", {"scene_count": 1}) |
| 181 | scene_dir = self.working_dir / "videos" / "event_0" / "scene_0" |
| 182 | scene_dir.mkdir(parents=True, exist_ok=True) |
| 183 | (scene_dir / "final_video.mp4").write_text("video", encoding="utf-8") |
| 184 | return { |
| 185 | "character_portraits_dir": str(self.working_dir / "character_portraits"), |
| 186 | "scene_videos_dir": str(self.working_dir / "videos"), |
| 187 | "scene_video_dirs": [str(scene_dir)], |
| 188 | "scene_count": 1, |
| 189 | } |
| 190 | |
| 191 | |
| 192 | def write_minimal_novel_artifacts(root: Path): |
| 193 | novel = root / "novel2video" |
| 194 | (novel / "novel").mkdir(parents=True, exist_ok=True) |
| 195 | (novel / "novel" / "novel.txt").write_text("novel", encoding="utf-8") |
| 196 | (novel / "novel" / "novel_compressed.txt").write_text("compressed", encoding="utf-8") |
| 197 | events = novel / "events" |
| 198 | events.mkdir(parents=True, exist_ok=True) |
| 199 | (events / "event_0.json").write_text(json.dumps(Event(index=0, is_last=True, description="d", process_chain=["p"]).model_dump()), encoding="utf-8") |
| 200 | chunks = novel / "relevant_chunks" / "event_0" |
| 201 | chunks.mkdir(parents=True, exist_ok=True) |
| 202 | (chunks / "chunk_0-score_0.95.txt").write_text("chunk", encoding="utf-8") |
| 203 | scenes = novel / "scenes" / "event_0" |
| 204 | scenes.mkdir(parents=True, exist_ok=True) |
| 205 | scene = Scene(idx=0, is_last=True, environment=EnvironmentInScene(slugline="INT. ROOM - DAY", description="room"), characters=[CharacterInScene(idx=0, identifier_in_scene="Hero", is_visible=True, static_features="adult", dynamic_features="coat")], script="<Hero> walks.") |
| 206 | (scenes / "scene_0.json").write_text(json.dumps(scene.model_dump()), encoding="utf-8") |
| 207 | event_level = novel / "global_information" / "characters" / "event_level" |
| 208 | event_level.mkdir(parents=True, exist_ok=True) |
| 209 | event_char = CharacterInEvent(index=0, identifier_in_event="Hero", active_scenes={0: "Hero"}, static_features="adult") |
| 210 | (event_level / "event_0_characters.json").write_text(json.dumps([event_char.model_dump()]), encoding="utf-8") |
| 211 | novel_level = novel / "global_information" / "characters" / "novel_level" |
| 212 | novel_level.mkdir(parents=True, exist_ok=True) |
| 213 | novel_char = CharacterInNovel(index=0, identifier_in_novel="Hero", active_events={0: "Hero"}, static_features="adult") |
| 214 | (novel_level / "novel_characters_after_event_0.json").write_text(json.dumps([novel_char.model_dump()]), encoding="utf-8") |
| 215 | |
| 216 | |
| 217 | class NovelAdapterTests(unittest.IsolatedAsyncioTestCase): |
| 218 | async def test_novel_initializes_named_empty_active_session(self): |
| 219 | with tempfile.TemporaryDirectory() as tmp: |
| 220 | index = SessionIndex(tmp) |
| 221 | empty = index.create(project_name="Novel draft") |
| 222 | adapter = ViMaxAdapters(Path(tmp), index) |
| 223 | with patch("agent_runtime.vimax_adapters._build_novel_pipeline", side_effect=lambda working_dir: FakeNovelPipeline(Path(working_dir))): |
| 224 | result = await adapter.vimax_novel_planning({"novel_text": "Hero opens a door."}) |
| 225 | self.assertTrue(result.ok) |
| 226 | payload = json.loads(result.content) |
| 227 | self.assertEqual(payload["session_id"], empty["session_id"]) |
| 228 | self.assertEqual(index.active()["project_name"], "Novel draft") |
| 229 | self.assertEqual(len(index.load()["sessions"]), 1) |
| 230 | |
| 231 | async def test_missing_rag_config_returns_tool_error_and_marks_session_error(self): |
| 232 | with tempfile.TemporaryDirectory() as tmp: |
| 233 | index = SessionIndex(tmp) |
| 234 | adapter = ViMaxAdapters(Path(tmp), index) |
| 235 | with patch.dict("os.environ", {"VIMAX_LLM_API_KEY": "llm-key", "VIMAX_LLM_BASE_URL": "https://llm.test/v1"}, clear=True), \ |
| 236 | patch("agent_runtime.vimax_adapters._build_embedding_model", side_effect=RuntimeError("embedding config missing")): |
| 237 | result = await adapter.vimax_novel_planning({"novel_text": "Hero opens a door."}) |
| 238 | self.assertFalse(result.ok) |
| 239 | self.assertIn("embedding", result.content.lower()) |
| 240 | self.assertEqual(index.active()["stage"], "error") |
| 241 | |
| 242 | async def test_success_writes_novel2video_artifacts_and_marks_novel_planned(self): |
| 243 | with tempfile.TemporaryDirectory() as tmp: |
| 244 | index = SessionIndex(tmp) |
| 245 | adapter = ViMaxAdapters(Path(tmp), index) |
| 246 | progress_events = [] |
| 247 | runtime = ToolRuntimeContext("vimax_novel_planning", "vimax_novel_planning", turn_id="turn-test", progress_callback=progress_events.append) |
| 248 | with patch("agent_runtime.vimax_adapters._build_novel_pipeline", side_effect=lambda working_dir: FakeNovelPipeline(Path(working_dir))): |
| 249 | result = await adapter.vimax_novel_planning({"novel_text": "Hero opens a door.", "style": "noir"}, runtime) |
| 250 | self.assertTrue(result.ok) |
| 251 | payload = json.loads(result.content) |
| 252 | root = Path(tmp) / payload["working_dir"] |
| 253 | self.assertTrue((root / "novel2video" / "novel" / "novel_compressed.txt").exists()) |
| 254 | self.assertTrue((root / "novel2video" / "events" / "event_0.json").exists()) |
| 255 | self.assertIn("novel2video/scenes/event_*/scene_*.json", payload["generated"]) |
| 256 | self.assertFalse(payload["ready_for_scene_render"]) |
| 257 | self.assertEqual(index.active()["stage"], "novel_planned") |
| 258 | stages = [event["progress"]["stage"] for event in progress_events if event.get("type") == "tool_progress"] |
| 259 | self.assertIn("novel_plan_text_artifacts", stages) |
| 260 | self.assertIn("merge_characters", stages) |
| 261 | |
| 262 | async def test_render_video_routes_novel2video_with_mock_pipeline(self): |
| 263 | with tempfile.TemporaryDirectory() as tmp: |
| 264 | index = SessionIndex(tmp) |
| 265 | record = index.create(idea="novel", user_requirement="scene render", style="noir") |
| 266 | root = Path(tmp) / record["working_dir"] |
| 267 | write_minimal_novel_artifacts(root) |
| 268 | adapter = ViMaxAdapters(Path(tmp), index) |
| 269 | progress_events = [] |
| 270 | runtime = ToolRuntimeContext("vimax_render_video", "vimax_render_video", turn_id="turn-test", progress_callback=progress_events.append) |
| 271 | with patch("agent_runtime.vimax_adapters._build_chat_model", return_value=object()), \ |
| 272 | patch("agent_runtime.vimax_adapters._build_image_generator", return_value=object()), \ |
| 273 | patch("agent_runtime.vimax_adapters._build_video_generator", return_value=object()), \ |
| 274 | patch("agent_runtime.vimax_adapters._build_novel_render_pipeline", side_effect=lambda working_dir, chat_model, image_generator, video_generator: FakeNovelRenderPipeline(Path(working_dir))): |
| 275 | result = await adapter.vimax_render_video({}, runtime) |
| 276 | self.assertTrue(result.ok) |
| 277 | payload = json.loads(result.content) |
| 278 | self.assertEqual(payload["render_mode"], "novel2video") |
| 279 | self.assertTrue(payload["scene_render_completed"]) |
| 280 | self.assertIsNone(payload["final_video_path"]) |
| 281 | self.assertEqual(payload["scene_count"], 1) |
| 282 | self.assertEqual(index.get(record["session_id"])["stage"], "novel_scene_rendered") |
| 283 | self.assertTrue((root / "novel2video" / "videos" / "event_0" / "scene_0" / "final_video.mp4").exists()) |
| 284 | stages = [event["progress"]["stage"] for event in progress_events if event.get("type") == "tool_progress"] |
| 285 | self.assertIn("novel_scene_render_start", stages) |
| 286 | self.assertIn("novel_render_completed", stages) |
| 287 | |
| 288 | |
| 289 | if __name__ == "__main__": |
| 290 | unittest.main() |
| 291 |