diff --git a/docs/3_loops_and_systems/story_loop.md b/docs/3_loops_and_systems/story_loop.md index ab8452b7..e6ffc5cd 100644 --- a/docs/3_loops_and_systems/story_loop.md +++ b/docs/3_loops_and_systems/story_loop.md @@ -12,6 +12,26 @@ between them, and finalize when the arc completes. **Source:** `packages/agents/src/monitor_agents/loops/story_loop.py` (`build_story_graph`). Delegates each scene to the [Scene Loop](./scene_loop.md). +## Arc contract + +A story spans several scenes. Finishing a scene advances the story to the next +one; the story itself ends when its arc resolves. + +The arc runs through four phases, and the story's `scene_budget` (5 scenes by +default) decides how many scenes each phase takes: + +| Phase | Scenes | What happens | +|---|---|---| +| `rising_action` | until `budget - 3` | Tension climbs; the story opens its plot threads. | +| `climax` | `budget - 2` | The conflict is confronted and the plot threads are paid off. | +| `falling_action` | `budget - 1` | The aftermath of the climax. | +| `resolution` | `budget` | The arc resolves and the story is complete. | + +Because the phases are ordered, the arc always reaches its resolution — a +story cannot stall. A caller that wants a longer or shorter campaign sets +`scene_budget` on the `StoryState`; a caller that wants the story to end early +passes `story_complete` in the scene outcome. + ## Nodes - `init_story` — establish arc parameters + initial world state. @@ -22,5 +42,12 @@ between them, and finalize when the arc completes. - `transition` — update continuity + plot threads before the next scene. - `finalize` — wrap the arc; ensure final world state is canonized. +## State + +The story state is checkpointed after every scene, so a campaign that is +interrupted — or resumed in another process — keeps its scenes, its arc +position and its in-game time. `StoryLoop.start_or_resume()` returns the +checkpointed story when there is one instead of starting over. + ## See Also - [Scene Loop](./scene_loop.md) · [Loops Index](./_index.md) diff --git a/packages/agents/src/monitor_agents/indexer/agent.py b/packages/agents/src/monitor_agents/indexer/agent.py index ab66fc83..d38a111c 100644 --- a/packages/agents/src/monitor_agents/indexer/agent.py +++ b/packages/agents/src/monitor_agents/indexer/agent.py @@ -766,8 +766,13 @@ async def _chunk_pdf_with_structure( # worker has no shared state race with the chunking loop below. blocks_with_tags: list[tuple[Any, str]] = [] for block in section_blocks: - heading_text = " > ".join(block.heading_path) - tag_category = TagPool.classify_text(heading_text) or "general" + # A subsection gives a more precise cue than its parent chapter. + # Fall back through ancestors when a heading has no known tag. + tag_category = next( + (category for heading in reversed(block.heading_path) + if (category := TagPool.classify_text(heading)) and category != "general"), + "general", + ) blocks_with_tags.append((block, tag_category)) sem = asyncio.Semaphore(_SECTION_CATEGORIZATION_CONCURRENCY) @@ -848,6 +853,9 @@ async def _classify(block: Any, tag_category: str) -> str: "heading_path": block.heading_path, "section_path": section_path, "section_depth": block.depth, + # These are section bounds, not exact chunk citations. + "section_page_start": page_num, + "section_page_end": block.page_end + 1, }, ) for chunk in block_chunks: diff --git a/packages/agents/src/monitor_agents/ingestion/agent.py b/packages/agents/src/monitor_agents/ingestion/agent.py index 9d8386db..9cff99e6 100644 --- a/packages/agents/src/monitor_agents/ingestion/agent.py +++ b/packages/agents/src/monitor_agents/ingestion/agent.py @@ -69,7 +69,7 @@ import traceback from datetime import UTC, datetime from pathlib import Path -from typing import Any +from typing import Any, Callable from uuid import UUID from monitor_data.db.minio import get_minio_client @@ -241,6 +241,7 @@ async def ingest_file( content_type: str | None = None, reuse_source_id: UUID | None = None, reuse_doc_id: UUID | None = None, + on_job_created: Callable[[UUID], None] | None = None, ) -> IngestionJobResponse: """ Ingest a raw file through the full pipeline. @@ -264,6 +265,7 @@ async def ingest_file( multiverse_id: Optional multiverse for auto-applying the KnowledgePack. auto_apply: If True, immediately creates ProposedChanges from the pack. content_type: MIME type override (auto-detected from filename if None). + on_job_created: Called with the persisted job ID as soon as it is created. Returns: IngestionJobResponse tracking pipeline progress. @@ -342,6 +344,8 @@ async def ingest_file( processing_checklist=selected_layers, ), ) + if on_job_created is not None: + on_job_created(job.job_id) return await asyncio.to_thread( mongodb_update_ingestion_job, job.job_id, @@ -433,6 +437,8 @@ async def ingest_file( ), ) job_id = job.job_id + if on_job_created is not None: + on_job_created(job_id) await asyncio.to_thread( mongodb_update_ingestion_job, job_id, diff --git a/packages/agents/src/monitor_agents/loops/conversation_loop.py b/packages/agents/src/monitor_agents/loops/conversation_loop.py index 0657cbb0..24c9aed8 100644 --- a/packages/agents/src/monitor_agents/loops/conversation_loop.py +++ b/packages/agents/src/monitor_agents/loops/conversation_loop.py @@ -369,12 +369,15 @@ async def generate_npc_responses(state: ConversationState) -> dict[str, Any]: for npc_id in state.npc_ids: try: + # Each NPC hears the replies the others already gave this round, + # so a group answers as a group instead of in isolation. + conversation_history = [*prior_turns, *new_turns[len(state.turns) :]] if state.mode == ConversationMode.DIRECT: result = await agent.respond_direct( conversation_id=state.conversation_id, npc_id=npc_id, player_said=state.current_player_input, - conversation_history=prior_turns, + conversation_history=conversation_history, player_entity_id=state.player_entity_id, scene_id=state.scene_id, story_id=state.story_id, @@ -445,7 +448,7 @@ async def generate_npc_responses(state: ConversationState) -> dict[str, Any]: conversation_id=state.conversation_id, npc_id=npc_id, gm_question=state.current_player_input, - conversation_history=prior_turns, + conversation_history=conversation_history, source_profile=state.source_profile, npc_data=state.npc_contexts.get(str(npc_id)), ) diff --git a/packages/agents/src/monitor_agents/loops/scene_loop.py b/packages/agents/src/monitor_agents/loops/scene_loop.py index 1dd2a9fc..c2cb5946 100644 --- a/packages/agents/src/monitor_agents/loops/scene_loop.py +++ b/packages/agents/src/monitor_agents/loops/scene_loop.py @@ -221,6 +221,9 @@ class SceneState(BaseModel): # Flow control scene_complete: bool = False + # Set by ``SceneLoop.finalize`` so the graph closes the scene without + # resolving or narrating an empty turn first. + finalize_requested: bool = False turns_count: int = 0 max_turns: int = 50 # safety ceiling per scene @@ -1475,7 +1478,13 @@ def build_scene_graph() -> StateGraph: graph.set_entry_point("load_context") # Edges - graph.add_edge("load_context", "resolve") + # A finalization request closes the scene straight away — no resolve, no + # narration of an empty turn. + graph.add_conditional_edges( + "load_context", + lambda state: "complete_current_scene" if state.finalize_requested else "resolve", + {"complete_current_scene": "complete_current_scene", "resolve": "resolve"}, + ) graph.add_conditional_edges( "resolve", route_after_resolve, @@ -1771,6 +1780,7 @@ async def finalize(self) -> None: "scene_id": self.scene_id, "story_id": self.story_id, "scene_complete": True, + "finalize_requested": True, "user_input": None, "memories_to_persist": [], "max_turns": self.max_turns, diff --git a/packages/agents/src/monitor_agents/loops/story_loop.py b/packages/agents/src/monitor_agents/loops/story_loop.py index 74d0f9ec..2286f661 100644 --- a/packages/agents/src/monitor_agents/loops/story_loop.py +++ b/packages/agents/src/monitor_agents/loops/story_loop.py @@ -71,6 +71,13 @@ class StoryState(BaseModel): scenes: list[UUID] = Field(default_factory=list) current_scene_id: UUID | None = None scenes_completed: int = 0 + # How many scenes the arc runs before it resolves. The arc spends the + # last three on the climax, the falling action and the resolution. + scene_budget: int = Field( + default=5, + ge=3, + description="Number of scenes in the story arc before it resolves.", + ) # Flow control story_complete: bool = False @@ -356,6 +363,12 @@ async def start_or_resume(self) -> dict[str, Any]: thread_id = f"story-{self.story_id}" config = {"configurable": {"thread_id": thread_id}} + # A story already under way keeps the state it checkpointed — + # re-running the graph here would reset it. + snapshot = await compiled.aget_state(config) # type: ignore[arg-type] + if snapshot.values: + return dict(snapshot.values) + initial = StoryState( story_id=self.story_id, universe_id=self.universe_id, @@ -373,6 +386,8 @@ async def advance(self, story_complete: bool = False) -> dict[str, Any]: story_id=self.story_id, universe_id=self.universe_id, story_complete=story_complete, + # The scene this call follows has been played. + scenes_completed=1, ) if story_complete: await finalize_story(state) @@ -416,7 +431,9 @@ async def complete_current_scene( Returns: Updated StoryState dict. """ - force_complete = bool((outcome or {}).get("scene_complete", False)) + # A finished scene advances the story; the story itself ends when its + # arc resolves, or when the caller says so outright. + force_complete = bool((outcome or {}).get("story_complete", False)) duration = int((outcome or {}).get("total_minutes_passed") or 0) with MongoDBSaver.from_conn_string(settings.mongodb_uri) as checkpointer: @@ -424,37 +441,44 @@ async def complete_current_scene( thread_id = f"story-{self.story_id}" config = {"configurable": {"thread_id": thread_id}} - # Load current checkpoint state - try: - current = await compiled.aget_state(config) # type: ignore[arg-type] - state_values = dict(current.values) if current.values else {} - except Exception: # noqa: BLE001 # reason: falls back to `{}` on failure (in complete_current_scene) - state_values = {} + # Load current checkpoint state. A read failure must not pass for + # "no story yet" — that would silently restart the campaign and + # overwrite its checkpoint. + current = await compiled.aget_state(config) # type: ignore[arg-type] + state_values = dict(current.values) if current.values else {} - # Build state with scene completion context + # Build state with scene completion context, keeping every other + # checkpointed field (in-game time, agreements, player pitch). previous_scenes = list(state_values.get("scenes", [])) previous_history = list(state_values.get("scenes_history", [])) previous_threads = list(state_values.get("active_threads", [])) previous_completed = list(state_values.get("completed_threads", [])) + # A story that already resolved stays resolved. + story_complete = force_complete or bool(state_values.get("story_complete", False)) - run_state = StoryState( - story_id=self.story_id, - universe_id=self.universe_id, - current_scene_id=scene_id, - story_complete=force_complete, - last_scene_duration_minutes=duration, - scenes=previous_scenes, - scenes_completed=state_values.get("scenes_completed", 0) + 1, - scenes_history=previous_history, - world_tone=state_values.get("world_tone", "dramatic"), - arc_label=state_values.get("arc_label", "rising_action"), - tension_score=float(state_values.get("tension_score", 0.3)), - active_threads=previous_threads, - completed_threads=previous_completed, + run_state = StoryState.model_validate( + { + **state_values, + "story_id": self.story_id, + "universe_id": self.universe_id, + "current_scene_id": scene_id, + "story_complete": story_complete, + "last_scene_duration_minutes": duration, + "scenes": previous_scenes, + "scenes_completed": state_values.get("scenes_completed", 0) + 1, + "scenes_history": previous_history, + "active_threads": previous_threads, + "completed_threads": previous_completed, + } ) - if force_complete: + if story_complete: await finalize_story(run_state) + await compiled.aupdate_state( + config, # type: ignore[arg-type] + run_state.model_dump(), + as_node="init_story", + ) return {"story_complete": True, "completed_scene_id": scene_id} # Step 1: Run world simulation @@ -510,6 +534,11 @@ async def complete_current_scene( if arc_update.get("story_complete"): await finalize_story(run_state) + await compiled.aupdate_state( + config, # type: ignore[arg-type] + run_state.model_dump(), + as_node="init_story", + ) return { **run_state.model_dump(), **progression_update, @@ -518,9 +547,14 @@ async def complete_current_scene( # Step 3: Create next scene via transition transition_update = await transition_scene(run_state) + checkpoint_state = {**run_state.model_dump(), **transition_update} + await compiled.aupdate_state( + config, # type: ignore[arg-type] + checkpoint_state, + as_node="init_story", + ) return { - **run_state.model_dump(), - **transition_update, + **checkpoint_state, **progression_update, "agenda_moves": agenda_moves, "completed_scene_id": scene_id, diff --git a/packages/agents/src/monitor_agents/story/agent.py b/packages/agents/src/monitor_agents/story/agent.py index 17d0ae87..98e9007e 100644 --- a/packages/agents/src/monitor_agents/story/agent.py +++ b/packages/agents/src/monitor_agents/story/agent.py @@ -69,34 +69,32 @@ async def evaluate_arc(self, state: StoryState) -> dict[str, Any]: """ Evaluate the narrative arc after each scene. - Uses heuristics based on scenes_completed, tension trends, and - thread counts to classify the arc phase. A DSPy-powered version - can override this with LLM-driven evaluation in the future. + The arc contract: a story spans several scenes and ends when its arc + resolves. The phases run in order — rising action, the climax that + confronts the conflict, the falling action after it, then the + resolution that closes the story — and the story's scene budget + decides how long each phase lasts. A story therefore always reaches + its resolution instead of stalling, and ``scene_budget`` is what a + caller sets to change how long a campaign runs. + + A DSPy-powered version can override this with LLM-driven evaluation + in the future. """ scenes_done = state.scenes_completed - active = len(state.active_threads) - completed_threads = len(state.completed_threads) tension = state.tension_score - - # Heuristic arc evaluation - if scenes_done <= 2: - label = "rising_action" - tension = min(0.5, tension + 0.15 * scenes_done) - elif tension >= 0.7 and active >= 2: - label = "climax" - tension = min(1.0, tension + 0.1) - elif completed_threads > active and tension > 0.5: - label = "falling_action" - tension = max(0.3, tension - 0.15) - elif active == 0 and completed_threads > 0: - label = "resolution" - tension = max(0.1, tension - 0.2) - elif active == 0 and scenes_done >= 3: - label = "new_thread" - tension = 0.3 - else: - label = "rising_action" - tension = min(0.6, tension + 0.1) + budget = state.scene_budget + + # The arc runs through four phases. The budget decides how many + # scenes the rising action takes; the last three are the climax, the + # falling action and the resolution. + phase = min(max(scenes_done - (budget - 3), 0), 3) + label = ("rising_action", "climax", "falling_action", "resolution")[phase] + tension = ( + min(0.9, tension + 0.15 * scenes_done), + min(1.0, tension + 0.1), + max(0.3, tension - 0.25), + max(0.1, tension - 0.2), + )[phase] # Suggest next scene type based on arc phase next_type = { @@ -107,14 +105,20 @@ async def evaluate_arc(self, state: StoryState) -> dict[str, Any]: "new_thread": "revelation", }.get(label, "action") - # Auto-complete story after resolution if no new threads - story_complete = label == "resolution" and active == 0 - - # Seed threads if none exist yet (world-building phase) - if not state.active_threads and scenes_done <= 1: + # Threads open with the story and are paid off at the climax; they + # stay paid off after it, so the resolution arrives with nothing + # left open. + active_threads = list(state.active_threads) + completed_threads = list(state.completed_threads) + if not active_threads and scenes_done <= 1: active_threads = ["central_conflict", "character_motivation"] - else: - active_threads = state.active_threads + if phase >= 1: + completed_threads = completed_threads + active_threads + active_threads = [] + + # The story is over once the arc reaches its resolution and no thread + # is left open. + story_complete = phase >= 3 and not active_threads # P-6: Automated Story Outlining - Call DSPy planner plan = await self._plan_next_scene(state) @@ -147,6 +151,7 @@ async def evaluate_arc(self, state: StoryState) -> dict[str, Any]: "arc_label": label, "tension_score": tension, "active_threads": active_threads, + "completed_threads": completed_threads, "next_scene_type": next_type, "scene_hook": scene_hook, "story_complete": story_complete, @@ -203,7 +208,9 @@ async def transition_scene(self, state: StoryState) -> dict[str, Any]: return { "scenes": state.scenes + [scene_id], "current_scene_id": scene_id, - "scenes_completed": state.scenes_completed + 1, + # The caller counts the scene it just finished; the transition + # only opens the next one. + "scenes_completed": state.scenes_completed, "next_scene_type": None, # consumed } diff --git a/packages/agents/tests/test_indexer.py b/packages/agents/tests/test_indexer.py index 9ae7746a..d3d43186 100644 --- a/packages/agents/tests/test_indexer.py +++ b/packages/agents/tests/test_indexer.py @@ -495,6 +495,29 @@ def _make_section_block( class TestChunkPdfWithStructure: """Tests for the structure-aware PDF chunking pre-pass.""" + @pytest.mark.asyncio + async def test_specific_bookmark_takes_precedence_and_preserves_page_range(self): + agent = _make_indexer() + block = _make_section_block( + heading_path=["Combat Rules", "Magic"], depth=1, + text="A spell effect.", page_start=9, page_end=14, + ) + with ( + patch("monitor_agents.indexer.agent.extract_pdf_structure", return_value=[block]), + patch("monitor_agents.indexer.agent.TagPool") as tags, + patch("monitor_agents.indexer.agent.SectionCategorizationModule"), + patch("monitor_agents.indexer.agent.chunk_text", return_value=[]) as chunker, + ): + tags.classify_text.side_effect = lambda heading: { + "Combat Rules": "combat_rules", "Magic": "power_system" + }.get(heading) + await agent._chunk_pdf_with_structure(b"%PDF", "rulebook") + + metadata = chunker.call_args.kwargs["metadata"] + assert metadata["semantic_category"] == "power_system" + assert metadata["section_path"] == "Combat Rules > Magic" + assert (metadata["section_page_start"], metadata["section_page_end"]) == (10, 15) + @pytest.mark.asyncio async def test_heading_path_appears_in_chunk_metadata(self): """heading_path and section_depth must be set on every chunk produced.""" diff --git a/packages/agents/tests/test_ingestion_pipeline.py b/packages/agents/tests/test_ingestion_pipeline.py index 7583433c..36465ff7 100644 --- a/packages/agents/tests/test_ingestion_pipeline.py +++ b/packages/agents/tests/test_ingestion_pipeline.py @@ -38,6 +38,7 @@ async def test_ingest_failure_preserves_source_and_uploaded_object(monkeypatch): source_id = uuid4() doc_id = uuid4() job_id = uuid4() + created_jobs = [] monkeypatch.setattr("monitor_agents.ingestion.agent.get_minio_client", lambda: fake_minio) monkeypatch.setattr( @@ -73,12 +74,45 @@ async def _boom(**_kwargs): universe_id=uuid4(), pack_name="Broken Source", pack_type=KnowledgePackType.SETTING_SUPPLEMENT, + on_job_created=created_jobs.append, ) + assert created_jobs == [job_id] assert deleted_sources == [] assert fake_minio.deleted_keys == [] +@pytest.mark.asyncio +async def test_duplicate_ingest_reports_created_job(monkeypatch): + pipeline = IngestionPipeline(agent_id="test-duplicate-callback") + job_id = uuid4() + existing = SimpleNamespace( + source_id=uuid4(), doc_id=uuid4(), extraction_status=ExtractionStatus.COMPLETED + ) + monkeypatch.setattr( + "monitor_agents.ingestion.agent.mongodb_find_document_by_filename", + lambda *_args, **_kwargs: existing, + ) + monkeypatch.setattr( + "monitor_agents.ingestion.agent.mongodb_create_ingestion_job", + lambda *_args, **_kwargs: SimpleNamespace(job_id=job_id), + ) + monkeypatch.setattr( + "monitor_agents.ingestion.agent.mongodb_update_ingestion_job", + lambda *_args, **_kwargs: SimpleNamespace(job_id=job_id), + ) + created_jobs = [] + await pipeline.ingest_file( + file_bytes=b"duplicate", + filename="duplicate.pdf", + source_title="Duplicate", + universe_id=uuid4(), + pack_name="Duplicate", + on_job_created=created_jobs.append, + ) + assert created_jobs == [job_id] + + @pytest.mark.asyncio async def test_ingest_reuses_existing_doc_for_rescan(monkeypatch): pipeline = IngestionPipeline(agent_id="test-reingest") diff --git a/packages/agents/tests/test_story_arc_contract.py b/packages/agents/tests/test_story_arc_contract.py new file mode 100644 index 00000000..3ac5034d --- /dev/null +++ b/packages/agents/tests/test_story_arc_contract.py @@ -0,0 +1,375 @@ +""" +Tests for the story arc contract: a story spans several scenes and ends +when its arc resolves. + +The arc runs rising action → climax → falling action → resolution, and the +story's scene budget decides how long it lasts. These tests pin that the +arc always reaches its resolution, that finishing a scene advances the story +instead of ending it, and that the story's progress survives a restart. + +Run: + cd /path/to/monitor_dm_system && pytest packages/agents/tests/test_story_arc_contract.py -v +""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from langgraph.checkpoint.memory import MemorySaver + +from monitor_agents.loops import scene_loop, story_loop + + +@contextmanager +def _saver_context(saver): + yield saver + + +@pytest.fixture +def arc_planner(monkeypatch): + """Stub the DSPy scene planner so the arc heuristic runs on its own.""" + from monitor_agents.story.agent import StoryAgent + + monkeypatch.setattr( + StoryAgent, + "_plan_next_scene", + AsyncMock(return_value={"next_scene_type": "exploration", "plot_hook": "a hook"}), + ) + + +async def _arc_labels(budget: int = 5) -> list[tuple[str, bool]]: + """Run the real arc evaluation over a whole story and collect its labels. + + Each round stands for one finished scene: the story counts the scene + before the arc looks at it, exactly as ``complete_current_scene`` does. + """ + from monitor_agents.story.agent import StoryAgent + + agent = StoryAgent(agent_id="arc-contract") + state = story_loop.StoryState(story_id=uuid4(), universe_id=uuid4(), scene_budget=budget) + labels: list[tuple[str, bool]] = [] + for _ in range(budget): + state = state.model_copy(update={"scenes_completed": state.scenes_completed + 1}) + update = await agent.evaluate_arc(state) + state = state.model_copy(update=update) + labels.append((update["arc_label"], update["story_complete"])) + return labels + + +class TestArcContract: + @pytest.mark.asyncio + async def test_arc_reaches_resolution_and_ends_the_story(self, arc_planner): + labels = await _arc_labels() + + assert [label for label, _ in labels] == [ + "rising_action", + "rising_action", + "climax", + "falling_action", + "resolution", + ] + assert [complete for _, complete in labels] == [False, False, False, False, True] + + @pytest.mark.asyncio + async def test_arc_length_follows_the_scene_budget(self, arc_planner): + labels = await _arc_labels(budget=7) + + assert [label for label, _ in labels][:6] == [ + "rising_action", + "rising_action", + "rising_action", + "rising_action", + "climax", + "falling_action", + ] + assert labels[6] == ("resolution", True) + + @pytest.mark.asyncio + async def test_the_shortest_budget_still_resolves(self, arc_planner): + labels = await _arc_labels(budget=3) + + assert [label for label, _ in labels] == ["climax", "falling_action", "resolution"] + assert labels[-1][1] is True + + @pytest.mark.asyncio + async def test_an_even_budget_puts_each_phase_where_it_belongs(self, arc_planner): + labels = await _arc_labels(budget=6) + + assert [label for label, _ in labels] == [ + "rising_action", + "rising_action", + "rising_action", + "climax", + "falling_action", + "resolution", + ] + assert labels[-1][1] is True + + @pytest.mark.asyncio + async def test_the_arc_curve_pins_tension_and_phase_boundaries(self, arc_planner): + """The whole curve: phase, tension and threads for every scene.""" + from monitor_agents.story.agent import StoryAgent + + agent = StoryAgent(agent_id="arc-curve") + state = story_loop.StoryState(story_id=uuid4(), universe_id=uuid4(), scene_budget=5, tension_score=0.3) + + seen = [] + for _ in range(5): + state = state.model_copy(update={"scenes_completed": state.scenes_completed + 1}) + update = await agent.evaluate_arc(state) + state = state.model_copy(update=update) + seen.append( + ( + update["arc_label"], + update["tension_score"], + list(update["active_threads"]), + list(update["completed_threads"]), + ) + ) + + assert [(label, threads, done) for label, _, threads, done in seen] == [ + ("rising_action", ["central_conflict", "character_motivation"], []), + ("rising_action", ["central_conflict", "character_motivation"], []), + ("climax", [], ["central_conflict", "character_motivation"]), + ("falling_action", [], ["central_conflict", "character_motivation"]), + ("resolution", [], ["central_conflict", "character_motivation"]), + ] + assert [tension for _, tension, _, _ in seen] == pytest.approx([0.45, 0.75, 0.85, 0.6, 0.4]) + + @pytest.mark.asyncio + async def test_tension_never_falls_below_the_floor(self, arc_planner): + """A quiet story still keeps its tension inside the arc's range.""" + from monitor_agents.story.agent import StoryAgent + + agent = StoryAgent(agent_id="arc-floor") + state = story_loop.StoryState( + story_id=uuid4(), + universe_id=uuid4(), + scene_budget=3, + tension_score=0.1, + ) + + tensions = [] + for _ in range(3): + state = state.model_copy(update={"scenes_completed": state.scenes_completed + 1}) + update = await agent.evaluate_arc(state) + state = state.model_copy(update=update) + tensions.append(update["tension_score"]) + + assert tensions == pytest.approx([0.2, 0.3, 0.1]) + + @pytest.mark.asyncio + async def test_the_story_opens_its_threads_on_its_first_scene(self, arc_planner): + """The threads open with the story, not only when no scene has run.""" + from monitor_agents.story.agent import StoryAgent + + agent = StoryAgent(agent_id="arc-first-scene") + state = story_loop.StoryState(story_id=uuid4(), universe_id=uuid4(), scenes_completed=1) + + update = await agent.evaluate_arc(state) + + assert update["active_threads"] == ["central_conflict", "character_motivation"] + assert update["completed_threads"] == [] + + @pytest.mark.asyncio + async def test_climax_pays_off_the_plot_threads(self, arc_planner): + from monitor_agents.story.agent import StoryAgent + + agent = StoryAgent(agent_id="arc-threads") + state = story_loop.StoryState(story_id=uuid4(), universe_id=uuid4()) + + first = await agent.evaluate_arc(state) + assert first["active_threads"] == ["central_conflict", "character_motivation"] + + state = state.model_copy(update=first) + for _ in range(2): + state = state.model_copy(update={"scenes_completed": state.scenes_completed + 1}) + state = state.model_copy(update=await agent.evaluate_arc(state)) + + state = state.model_copy(update={"scenes_completed": state.scenes_completed + 1}) + climax = await agent.evaluate_arc(state) + assert climax["arc_label"] == "climax" + assert climax["active_threads"] == [] + assert climax["completed_threads"] == ["central_conflict", "character_motivation"] + + +class TestStoryProgression: + @pytest.mark.asyncio + async def test_finishing_a_scene_advances_the_story_without_ending_it(self, monkeypatch, arc_planner): + saver = MemorySaver() + monkeypatch.setattr(story_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + monkeypatch.setattr(story_loop, "simulate_world_advance", AsyncMock(return_value={})) + next_scene = uuid4() + monkeypatch.setattr(story_loop, "transition_scene", AsyncMock(return_value={"current_scene_id": next_scene})) + finalize = AsyncMock() + monkeypatch.setattr(story_loop, "finalize_story", finalize) + + loop = story_loop.StoryLoop(uuid4(), uuid4()) + loop.call_tool = AsyncMock(return_value=[]) + + result = await loop.complete_current_scene(uuid4(), {"scene_complete": True}) + + assert result["story_complete"] is False + assert result["current_scene_id"] == next_scene + finalize.assert_not_awaited() + + @pytest.mark.asyncio + async def test_story_ends_when_its_arc_resolves(self, monkeypatch, arc_planner): + saver = MemorySaver() + monkeypatch.setattr(story_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + monkeypatch.setattr(story_loop, "simulate_world_advance", AsyncMock(return_value={})) + monkeypatch.setattr(story_loop, "transition_scene", AsyncMock(return_value={"current_scene_id": uuid4()})) + finalize = AsyncMock() + monkeypatch.setattr(story_loop, "finalize_story", finalize) + + loop = story_loop.StoryLoop(uuid4(), uuid4()) + loop.call_tool = AsyncMock(return_value=[]) + + results = [] + for _ in range(5): + results.append(await loop.complete_current_scene(uuid4(), {"scene_complete": True})) + + assert results[-1]["story_complete"] is True + assert finalize.await_count == 1 + # One scene counted per finished scene — not two. + assert results[-1]["scenes_completed"] == 5 + + @pytest.mark.asyncio + async def test_story_progress_survives_a_restart(self, monkeypatch, arc_planner): + saver = MemorySaver() + monkeypatch.setattr(story_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + monkeypatch.setattr(story_loop, "simulate_world_advance", AsyncMock(return_value={})) + monkeypatch.setattr(story_loop, "transition_scene", AsyncMock(return_value={"current_scene_id": uuid4()})) + + loop = story_loop.StoryLoop(uuid4(), uuid4()) + loop.call_tool = AsyncMock(return_value=[]) + + first = await loop.complete_current_scene(uuid4(), {"total_minutes_passed": 10}) + second = await loop.complete_current_scene(first["current_scene_id"], {"total_minutes_passed": 20}) + + resumed = await loop.start_or_resume() + + assert second["scenes_completed"] == 2 + assert resumed["scenes_completed"] == 2 + assert resumed["current_scene_id"] == second["current_scene_id"] + assert resumed["arc_label"] == second["arc_label"] + + @pytest.mark.asyncio + async def test_a_story_asked_to_end_ends_at_once(self, monkeypatch, arc_planner): + saver = MemorySaver() + monkeypatch.setattr(story_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + finalize = AsyncMock() + monkeypatch.setattr(story_loop, "finalize_story", finalize) + + loop = story_loop.StoryLoop(uuid4(), uuid4()) + loop.call_tool = AsyncMock(return_value=[]) + + result = await loop.complete_current_scene(uuid4(), {"story_complete": True}) + + assert result["story_complete"] is True + finalize.assert_awaited_once() + + @pytest.mark.asyncio + async def test_a_resolved_story_does_not_open_another_scene(self, monkeypatch, arc_planner): + saver = MemorySaver() + monkeypatch.setattr(story_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + monkeypatch.setattr(story_loop, "simulate_world_advance", AsyncMock(return_value={})) + transition = AsyncMock(return_value={"current_scene_id": uuid4()}) + monkeypatch.setattr(story_loop, "transition_scene", transition) + monkeypatch.setattr(story_loop, "finalize_story", AsyncMock()) + + loop = story_loop.StoryLoop(uuid4(), uuid4()) + loop.call_tool = AsyncMock(return_value=[]) + + for _ in range(5): + await loop.complete_current_scene(uuid4(), {"scene_complete": True}) + + transition.reset_mock() + result = await loop.complete_current_scene(uuid4(), {"scene_complete": True}) + + assert result["story_complete"] is True + transition.assert_not_awaited() + + +class TestSceneFinalize: + @pytest.mark.asyncio + async def test_finalize_closes_the_scene_without_narrating_a_turn(self, monkeypatch): + saver = MemorySaver() + monkeypatch.setattr(scene_loop.MongoDBSaver, "from_conn_string", lambda _: _saver_context(saver)) + monkeypatch.setattr(scene_loop, "_SCENE_GRAPH_TEMPLATE", None) + calls: list[str] = [] + + async def load(state): + return {} + + async def resolve(state): + calls.append("resolve") + return {"resolution": {"resolution_type": "narrative"}} + + async def narrate(state): + calls.append("narrate") + return {"turns_count": state.turns_count + 1} + + async def complete(state): + calls.append("complete") + return {"scene_complete": True} + + async def empty(state): + return {} + + monkeypatch.setattr(scene_loop, "load_context", load) + monkeypatch.setattr(scene_loop, "resolve_action", resolve) + monkeypatch.setattr(scene_loop, "narrate", narrate) + monkeypatch.setattr(scene_loop, "complete_current_scene", complete) + for name in ( + "extract_new_entities", + "extract_memories", + "extract_facts", + "persist_memories_node", + "check_consistency", + "check_events", + "check_foreshadowing", + "persist_turn_artifacts", + "canonize_checkpoint", + ): + monkeypatch.setattr(scene_loop, name, empty) + + loop = scene_loop.SceneLoop(scene_id=uuid4(), story_id=uuid4()) + await loop.finalize() + + assert calls == ["complete"] + + +class TestGroupConversation: + @pytest.mark.asyncio + async def test_each_npc_hears_the_replies_already_given(self, monkeypatch): + from monitor_agents.loops import conversation_loop + from monitor_agents.npc_voice.agent import NPCVoice + from monitor_data.schemas.conversations import ConversationMode + + first, second = uuid4(), uuid4() + histories: list[list[dict]] = [] + + async def respond(self, **kwargs): + histories.append(kwargs["conversation_history"]) + return {"npc_response": str(kwargs["npc_id"]), "emotional_state_after": {}} + + monkeypatch.setattr(NPCVoice, "respond_direct", respond) + + state = conversation_loop.ConversationState( + conversation_id=uuid4(), + universe_id=uuid4(), + mode=ConversationMode.DIRECT, + npc_ids=[first, second], + current_player_input="hello", + turns=[{"speaker_role": "player", "text": "hello"}], + ) + + result = await conversation_loop.generate_npc_responses(state) + + assert len(result["current_npc_responses"]) == 2 + # The second NPC hears the first NPC's reply. + assert histories[1][-1]["entity_id"] == str(first) diff --git a/packages/agents/tests/test_story_loop.py b/packages/agents/tests/test_story_loop.py index 8058803e..b95d1695 100644 --- a/packages/agents/tests/test_story_loop.py +++ b/packages/agents/tests/test_story_loop.py @@ -263,7 +263,7 @@ async def test_run_scene_returns_scene_count(self): class TestTransitionSceneNode: @pytest.mark.asyncio async def test_transition_creates_next_scene(self): - """transition_scene creates a new scene and increments counter.""" + """transition_scene creates a new scene and leaves the counter alone.""" (StoryState, *_) = _import_story_loop() scene0 = uuid4() state = StoryState( @@ -285,7 +285,9 @@ async def test_transition_creates_next_scene(self): assert result["current_scene_id"] == new_scene_id assert new_scene_id in result["scenes"] - assert result["scenes_completed"] == 2 + # The caller counts the scene it just finished; the transition only + # opens the next one. + assert result["scenes_completed"] == 1 @pytest.mark.asyncio async def test_transition_noop_when_story_complete(self): diff --git a/packages/agents/tests/test_story_loop_procedural.py b/packages/agents/tests/test_story_loop_procedural.py index 2317b531..d405cd02 100644 --- a/packages/agents/tests/test_story_loop_procedural.py +++ b/packages/agents/tests/test_story_loop_procedural.py @@ -112,7 +112,7 @@ async def test_transition_scene_triggers_procedural_population(): res = await transition_scene(state) assert res["current_scene_id"] == new_scene_id - assert res["scenes_completed"] == 2 + assert res["scenes_completed"] == 1 # Verify procedural population was called mock_populate.assert_called_once_with(universe_id=u_id, scene_id=new_scene_id, location_type="dungeon") @@ -144,7 +144,7 @@ async def test_transition_scene_with_town_location(): res = await transition_scene(state) - assert res["scenes_completed"] == 1 + assert res["scenes_completed"] == 0 mock_populate.assert_called_once_with(universe_id=u_id, scene_id=new_scene_id, location_type="town") diff --git a/packages/data-layer/src/monitor_data/tools/ingest_tools/pdf_processing.py b/packages/data-layer/src/monitor_data/tools/ingest_tools/pdf_processing.py index 572e098b..18ac3ac1 100644 --- a/packages/data-layer/src/monitor_data/tools/ingest_tools/pdf_processing.py +++ b/packages/data-layer/src/monitor_data/tools/ingest_tools/pdf_processing.py @@ -233,6 +233,22 @@ def _extract_pdf_structure_uncached(pdf_bytes: bytes) -> list[SectionBlock]: if toc: sections = [] + # Parent and child bookmarks frequently span the same pages. Extract + # each page once per document rather than re-parsing it for each span. + page_text: dict[int, str] = {} + first_bookmark_page = max(0, min(doc.page_count, toc[0][2] - 1)) + if first_bookmark_page: + for p in range(first_bookmark_page): + page_text[p] = doc[p].get_text() + sections.append( + SectionBlock( + heading_path=["(front matter)"], + depth=0, + page_start=0, + page_end=first_bookmark_page - 1, + text="\n".join(page_text[p] for p in range(first_bookmark_page)).strip(), + ) + ) for i, entry in enumerate(toc): level = entry[0] title = entry[1] @@ -258,7 +274,11 @@ def _extract_pdf_structure_uncached(pdf_bytes: bytes) -> list[SectionBlock]: path.append(title) # Collect text from page range - text_parts = [doc[p].get_text() for p in range(max(0, page_start), min(doc.page_count, page_end + 1))] + text_parts = [] + for p in range(max(0, page_start), min(doc.page_count, page_end + 1)): + if p not in page_text: + page_text[p] = doc[p].get_text() + text_parts.append(page_text[p]) text = "\n".join(text_parts).strip() sections.append( diff --git a/packages/data-layer/tests/test_db/test_ingest_tools.py b/packages/data-layer/tests/test_db/test_ingest_tools.py index a01b38e2..caba24eb 100644 --- a/packages/data-layer/tests/test_db/test_ingest_tools.py +++ b/packages/data-layer/tests/test_db/test_ingest_tools.py @@ -71,6 +71,115 @@ def test_extract_pdf_structure_heading_path_nesting(): assert combat_sections[0].depth >= 1 +def test_overlapping_bookmarks_reuse_page_text(monkeypatch): + from monitor_data.tools.ingest_tools import pdf_processing + + reads = [] + + class Page: + def __init__(self, number): + self.number = number + + def get_text(self): + reads.append(self.number) + return f"Page {self.number}" + + class Document: + page_count = 2 + + def get_toc(self, simple=False): + return [[1, "Chapter", 1], [2, "Section", 2]] + + def __getitem__(self, number): + return Page(number) + + def close(self): + pass + + monkeypatch.setattr(pdf_processing, "_open_pdf", lambda _: Document()) + sections = pdf_processing._extract_pdf_structure_uncached(b"pdf") + + assert reads == [0, 1] + assert sections[0].text == "Page 0\nPage 1" + assert sections[1].text == "Page 1" + + +def test_bookmarks_preserve_unbookmarked_front_matter(): + doc = fitz.open() + doc.new_page().insert_text((72, 72), "Unbookmarked introduction") + doc.new_page().insert_text((72, 72), "Chapter text") + doc.set_toc([[1, "Chapter", 2]]) + data = doc.tobytes() + doc.close() + + sections = extract_pdf_structure(data) + + assert sections[0].page_start == sections[0].page_end == 0 + assert "Unbookmarked introduction" in sections[0].text + assert sections[1].heading_path == ["Chapter"] + assert "Chapter text" in sections[1].text + + +def test_front_matter_and_section_bounds_are_exact(): + """The pages a section claims are exactly the pages it holds. + + Front matter runs up to the page before the first bookmark; each + bookmark section runs up to the page before the next one starts. A page + that belongs to no bookmark is still covered by the front matter. + """ + doc = fitz.open() + for number in range(4): + doc.new_page().insert_text((72, 72), f"Body of page {number}") + doc.set_toc([[1, "Chapter One", 3], [1, "Chapter Two", 4]]) + data = doc.tobytes() + doc.close() + + sections = extract_pdf_structure(data) + + front, one, two = sections[0], sections[1], sections[2] + + assert front.heading_path == ["(front matter)"] + assert front.depth == 0 + assert (front.page_start, front.page_end) == (0, 1) + assert "Body of page 0" in front.text + assert "Body of page 1" in front.text + assert "Body of page 2" not in front.text + + assert one.heading_path == ["Chapter One"] + assert (one.page_start, one.page_end) == (2, 2) + assert "Body of page 2" in one.text + assert "Body of page 3" not in one.text + + assert two.heading_path == ["Chapter Two"] + assert (two.page_start, two.page_end) == (3, 3) + assert "Body of page 3" in two.text + + +def test_front_matter_spans_every_page_before_the_first_bookmark(): + """Three pages of front matter are three pages, not one or two.""" + doc = fitz.open() + for number in range(5): + doc.new_page().insert_text((72, 72), f"Body of page {number}") + doc.set_toc([[1, "Chapter One", 4], [1, "Chapter Two", 5]]) + data = doc.tobytes() + doc.close() + + sections = extract_pdf_structure(data) + + front, one, two = sections[0], sections[1], sections[2] + + assert (front.page_start, front.page_end) == (0, 2) + for number in range(3): + assert f"Body of page {number}" in front.text + assert "Body of page 3" not in front.text + + assert (one.page_start, one.page_end) == (3, 3) + assert "Body of page 3" in one.text + + assert (two.page_start, two.page_end) == (4, 4) + assert "Body of page 4" in two.text + + # [G-8](d) — extract_pdf_structure wraps an in-memory SHA-256 LRU. # Tests pin (a) cache hit on identical bytes (b) mutation isolation # (c) cache miss on different bytes (d) LRU eviction at 5th distinct PDF. diff --git a/packages/ui/backend/src/monitor_ui/routers/ingest.py b/packages/ui/backend/src/monitor_ui/routers/ingest.py index b4e34e71..3161f20d 100644 --- a/packages/ui/backend/src/monitor_ui/routers/ingest.py +++ b/packages/ui/backend/src/monitor_ui/routers/ingest.py @@ -294,8 +294,7 @@ def _run_ingest_in_thread( """Run the ingest pipeline in a dedicated thread with its own event loop. This prevents the long-running sync+async pipeline calls from blocking - uvicorn's main event loop. The _ingest_lock ensures only one pipeline - runs at a time. An overall timeout of _JOB_TIMEOUT_SECS is enforced via + uvicorn's main event loop. The executor limits concurrent pipelines. An overall timeout of _JOB_TIMEOUT_SECS is enforced via asyncio.wait_for so a stalled LLM call cannot block the executor forever. """ global _ingest_active_title, _ingest_active_job_id @@ -308,73 +307,54 @@ async def _run() -> None: pipeline = IngestionPipeline() async def _ingest_with_capture() -> None: - # We wrap ingest_file and capture the job_id as soon as MongoDB - # creates it by patching the symbol that `IngestionPipeline` actually uses. - import monitor_agents.ingestion.agent as _pipeline_module + def _capture_job(job_id: UUID) -> None: + _captured_job_id[0] = str(job_id) + _attach_job_id_to_request(queue_token, str(job_id)) + logger.info("Created ingestion job %s for '%s'", job_id, source_title) - # reason: capturing the original mongodb_create_ingestion_job reference for runtime monkey-patch restoration; module attribute access narrows to a typed Callable that the local doesn't preserve at the restoration site - _orig_create_job = _pipeline_module.mongodb_create_ingestion_job # type: ignore - - def _patched_create_job(params): # type: ignore[no-untyped-def] - result = _orig_create_job(params) - _captured_job_id[0] = str(result.job_id) - _attach_job_id_to_request(queue_token, _captured_job_id[0]) - logger.info( - "Created ingestion job %s for '%s'", - _captured_job_id[0], - source_title, - ) - return result + # 1. Standard ingestion pipeline + await pipeline.ingest_file( + file_bytes=file_bytes, + filename=filename, + source_title=source_title, + # reason: universe_uid: UUID | None passed to IngestionPipeline.ingest_file(universe_id: UUID), which expects a non-None UUID + universe_id=universe_uid, # type: ignore + pack_name=source_title, + pack_type=pack_type_enum, + analysis_layers=selected_layers, + auto_apply=False, + content_type=content_type, + reuse_source_id=reuse_source_id, + reuse_doc_id=reuse_doc_id, + on_job_created=_capture_job, + ) - # reason: installing the untyped _patched_create_job (no parameter annotation) over the typed module attribute — mypy can't bridge the runtime swap from an untyped function to the typed signature - _pipeline_module.mongodb_create_ingestion_job = _patched_create_job # type: ignore - try: - # 1. Standard ingestion pipeline - await pipeline.ingest_file( - file_bytes=file_bytes, - filename=filename, - source_title=source_title, - # reason: universe_uid: UUID | None passed to IngestionPipeline.ingest_file(universe_id: UUID), which expects a non-None UUID - universe_id=universe_uid, # type: ignore - pack_name=source_title, - pack_type=pack_type_enum, - analysis_layers=selected_layers, - auto_apply=False, - content_type=content_type, - reuse_source_id=reuse_source_id, - reuse_doc_id=reuse_doc_id, - ) + # 2. Record in World Library for the Architect + if universe_uid: + try: + from monitor_agents.utils.world_library import WorldLibrary + from monitor_data.tools.neo4j_tools.core import ( + neo4j_get_universe, + ) - # 2. Record in World Library for the Architect - if universe_uid: - try: - from monitor_agents.utils.world_library import WorldLibrary - from monitor_data.tools.neo4j_tools.core import ( - neo4j_get_universe, + library = WorldLibrary() + universe = neo4j_get_universe(universe_uid) + + if universe: + await library.add_source( + # reason: UniverseResponse.multiverse_id: UUID | None passed to WorldLibrary.add_source(multiverse_id: UUID), which expects a non-None UUID + multiverse_id=universe.multiverse_id, # type: ignore + source_type="file", + uri=filename, + # reason: 'system_name' is not a defined attribute on UniverseResponse (only 'default_system_name' exists); access would raise AttributeError at runtime — flagged as a latent bug, not fixed here + system_name=universe.system_name or "generic", # type: ignore + metadata={ + "title": source_title, + "universe_id": str(universe_uid), + }, ) - - library = WorldLibrary() - universe = neo4j_get_universe(universe_uid) - - if universe: - await library.add_source( - # reason: UniverseResponse.multiverse_id: UUID | None passed to WorldLibrary.add_source(multiverse_id: UUID), which expects a non-None UUID - multiverse_id=universe.multiverse_id, # type: ignore - source_type="file", - uri=filename, - # reason: 'system_name' is not a defined attribute on UniverseResponse (only 'default_system_name' exists); access would raise AttributeError at runtime — flagged as a latent bug, not fixed here - system_name=universe.system_name or "generic", # type: ignore - metadata={ - "title": source_title, - "universe_id": str(universe_uid), - }, - ) - except Exception as exc: # noqa: BLE001 # reason: boundary logs and degrades gracefully (in _patched_create_job) - logger.debug("Failed to record file in WorldLibrary: %s", exc) - - finally: - # reason: restoring the original function reference — symmetric to the capture/install pattern at lines 311/324; mypy sees the typed callable being written back without the runtime context of the prior untyped swap - _pipeline_module.mongodb_create_ingestion_job = _orig_create_job # type: ignore + except Exception as exc: # noqa: BLE001 # reason: boundary logs and degrades gracefully (in _patched_create_job) + logger.debug("Failed to record file in WorldLibrary: %s", exc) try: await asyncio.wait_for(_ingest_with_capture(), timeout=_JOB_TIMEOUT_SECS) diff --git a/packages/ui/backend/tests/test_ingest.py b/packages/ui/backend/tests/test_ingest.py index 78cb7839..d35a265d 100644 --- a/packages/ui/backend/tests/test_ingest.py +++ b/packages/ui/backend/tests/test_ingest.py @@ -287,6 +287,33 @@ def test_run_ingest_in_thread(mock_all_dbs): content_type="application/pdf" ) + +def test_ingest_worker_attaches_job_to_its_own_request(mock_all_dbs): + import monitor_ui.routers.ingest as ingest_mod + from monitor_data.schemas.knowledge_packs import KnowledgePackType + + job_id = uuid.uuid4() + attached = [] + + async def fake_ingest(**kwargs): + kwargs["on_job_created"](job_id) + + with patch.object(ingest_mod, "IngestionPipeline") as pipeline: + pipeline.return_value.ingest_file = fake_ingest + with patch.object(ingest_mod, "_attach_job_id_to_request", side_effect=lambda *args: attached.append(args)): + ingest_mod._run_ingest_in_thread( + queue_token="own-request", + file_bytes=b"test", + filename="test.pdf", + source_title="test", + universe_uid=None, + pack_type_enum=KnowledgePackType.SETTING_SUPPLEMENT, + selected_layers=["entities"], + content_type="application/pdf", + ) + + assert attached == [("own-request", str(job_id))] + def test_upload_exceptions(mock_all_dbs): import os os.environ["MONITOR_MAX_INGEST_FILE_BYTES"] = "5" diff --git a/scripts/assess_ingestion_corpus.py b/scripts/assess_ingestion_corpus.py new file mode 100644 index 00000000..e1287e27 --- /dev/null +++ b/scripts/assess_ingestion_corpus.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Measure PDF text and bookmark coverage using local files and no LLM calls. + +Run: uv run python scripts/assess_ingestion_corpus.py examples/*.pdf +The JSON output contains counts and timings, never extracted source text. +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import fitz +from monitor_data.tools.ingest_tools.pdf_processing import extract_pdf_structure + + +def assess(path: Path) -> dict: + data = path.read_bytes() + with fitz.open(stream=data, filetype="pdf") as doc: + page_count = len(doc) + bookmarks = len(doc.get_toc()) + chars = [len(page.get_text().strip()) for page in doc] + + started = time.perf_counter() + sections = extract_pdf_structure(data) + elapsed = time.perf_counter() - started + covered = { + page + for section in sections + for page in range(max(0, section.page_start), min(page_count, section.page_end + 1)) + } + return { + "file": path.name, + "bytes": len(data), + "pages": page_count, + "bookmarks": bookmarks, + "pages_without_text": sum(count == 0 for count in chars), + "pages_under_100_characters": sum(count < 100 for count in chars), + "text_characters": sum(chars), + "sections": len(sections), + "sections_without_text": sum(not section.text.strip() for section in sections), + "pages_outside_sections": page_count - len(covered), + "section_extraction_seconds": round(elapsed, 3), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("pdf", nargs="+", type=Path) + args = parser.parse_args() + results = [] + failed = False + for path in args.pdf: + try: + results.append(assess(path)) + except (OSError, ValueError, RuntimeError) as exc: + failed = True + results.append({"file": path.name, "error": str(exc)}) + print(json.dumps(results, indent=2)) + return int(failed) + + +if __name__ == "__main__": + raise SystemExit(main())