diff --git a/narratives/AGENTS.md b/narratives/AGENTS.md index f349d45c..9e2bce58 100644 --- a/narratives/AGENTS.md +++ b/narratives/AGENTS.md @@ -74,8 +74,7 @@ process and PR expectations: [`CONTRIBUTING.md`](../CONTRIBUTING.md). the default chart floor (`chart_defaults.py`), the synthesis prompt builder (`synthesis.py`), the streaming block splitter (`blocks.py`), - chart configuration (`chart_config.py`), follow-up generation - (`follow_up.py`), + follow-up generation (`follow_up.py`), user-visible messages (`messages.py`), the typed SSE frames the agent loop sends (`sse_frames.py`), the turn's state slots (`state_slots.py`), and signed transcripts (`transcript.py`). Of the diff --git a/narratives/README.md b/narratives/README.md index 40bb41bf..d1d18a89 100644 --- a/narratives/README.md +++ b/narratives/README.md @@ -10,7 +10,7 @@ Context Protocol (MCP). - **Grounded statistical Q&A:** Answers natural-language questions using Data Commons MCP tools. -- **Interactive visualizations:** Renders Data Commons charts alongside +- **Interactive visualizations:** Renders Data Commons charts inline within generated responses. - **Source attribution:** Attaches inline citations and dataset provenance to statistical claims. diff --git a/narratives/agent/src/narratives_agent/config_model.py b/narratives/agent/src/narratives_agent/config_model.py index 385bd4f9..e1f86cef 100644 --- a/narratives/agent/src/narratives_agent/config_model.py +++ b/narratives/agent/src/narratives_agent/config_model.py @@ -60,8 +60,6 @@ class GeminiRole(Enum): CLASSIFY = "classify" DISCOVERY = "discovery" SYNTHESIS = "synthesis" - CHART_CONFIG = "chart_config" - CHART_VALIDATION = "chart_validation" FOLLOW_UP = "follow_up" COMPACTION = "compaction" @@ -85,8 +83,6 @@ class RoleConfig: GeminiRole.CLASSIFY: 0.0, GeminiRole.DISCOVERY: 0.2, GeminiRole.SYNTHESIS: 0.3, - GeminiRole.CHART_CONFIG: 0.2, - GeminiRole.CHART_VALIDATION: 0.0, # A higher temperature varies the phrasing of the follow-up questions. GeminiRole.FOLLOW_UP: 0.8, GeminiRole.COMPACTION: 0.1, @@ -100,8 +96,6 @@ class RoleConfig: # accept `MINIMAL`. _FIXED_THINKING_LEVELS: Mapping[GeminiRole, types.ThinkingLevel] = { GeminiRole.CLASSIFY: types.ThinkingLevel.MINIMAL, - GeminiRole.CHART_CONFIG: types.ThinkingLevel.LOW, - GeminiRole.CHART_VALIDATION: types.ThinkingLevel.LOW, GeminiRole.FOLLOW_UP: types.ThinkingLevel.LOW, GeminiRole.COMPACTION: types.ThinkingLevel.LOW, } @@ -179,8 +173,6 @@ class GeminiConfig: classify_model: str = DEFAULT_GEMINI_CLASSIFY_MODEL discovery_model: str = DEFAULT_GEMINI_MODEL synthesis_model: str = DEFAULT_GEMINI_MODEL - chart_config_model: str = DEFAULT_GEMINI_MODEL - chart_validation_model: str = DEFAULT_GEMINI_MODEL follow_up_model: str = DEFAULT_GEMINI_MODEL compaction_model: str = DEFAULT_GEMINI_MODEL api_key: str = field(default="", repr=False) @@ -445,14 +437,6 @@ def parse_agent_config(document: Mapping[str, Any]) -> AgentConfig: synthesis_model=_read_model_name( gemini, "synthesis_model", GeminiConfig.synthesis_model ), - chart_config_model=_read_model_name( - gemini, "chart_config_model", GeminiConfig.chart_config_model - ), - chart_validation_model=_read_model_name( - gemini, - "chart_validation_model", - GeminiConfig.chart_validation_model, - ), follow_up_model=_read_model_name( gemini, "follow_up_model", GeminiConfig.follow_up_model ), diff --git a/narratives/agent/src/narratives_agent/config_model_test.py b/narratives/agent/src/narratives_agent/config_model_test.py index 677e7e24..0ccb8066 100644 --- a/narratives/agent/src/narratives_agent/config_model_test.py +++ b/narratives/agent/src/narratives_agent/config_model_test.py @@ -77,8 +77,6 @@ def test_an_empty_document_gives_every_default() -> None: classify_model=DEFAULT_GEMINI_CLASSIFY_MODEL, discovery_model=DEFAULT_GEMINI_MODEL, synthesis_model=DEFAULT_GEMINI_MODEL, - chart_config_model=DEFAULT_GEMINI_MODEL, - chart_validation_model=DEFAULT_GEMINI_MODEL, follow_up_model=DEFAULT_GEMINI_MODEL, compaction_model=DEFAULT_GEMINI_MODEL, ) @@ -104,8 +102,6 @@ def test_configured_values_are_carried_through() -> None: "classify_model": " gemini-3.8-flash ", "discovery_model": " gemini-3-pro ", "synthesis_model": " gemini-3-pro-synthesis ", - "chart_config_model": " gemini-3-flash-charts ", - "chart_validation_model": " gemini-3-flash-val ", "follow_up_model": " gemini-3-flash-followup ", "compaction_model": " gemini-3-flash-compact ", "api_key": " test-key ", @@ -130,8 +126,6 @@ def test_configured_values_are_carried_through() -> None: classify_model="gemini-3.8-flash", discovery_model="gemini-3-pro", synthesis_model="gemini-3-pro-synthesis", - chart_config_model="gemini-3-flash-charts", - chart_validation_model="gemini-3-flash-val", follow_up_model="gemini-3-flash-followup", compaction_model="gemini-3-flash-compact", api_key="test-key", @@ -361,16 +355,6 @@ def test_an_unrecognized_thinking_level_gives_the_default( RoleConfig(DEFAULT_GEMINI_MODEL, 0.3, types.ThinkingLevel.MEDIUM), id="synthesis", ), - pytest.param( - GeminiRole.CHART_CONFIG, - RoleConfig(DEFAULT_GEMINI_MODEL, 0.2, types.ThinkingLevel.LOW), - id="chart_config", - ), - pytest.param( - GeminiRole.CHART_VALIDATION, - RoleConfig(DEFAULT_GEMINI_MODEL, 0.0, types.ThinkingLevel.LOW), - id="chart_validation", - ), pytest.param( GeminiRole.FOLLOW_UP, RoleConfig(DEFAULT_GEMINI_MODEL, 0.8, types.ThinkingLevel.LOW), @@ -446,8 +430,6 @@ def test_configured_values_apply_to_the_roles_they_govern() -> None: classify_model="gemini-3.6-flash", discovery_model="gemini-discovery-custom", synthesis_model="gemini-synthesis-custom", - chart_config_model="gemini-chart-config-custom", - chart_validation_model="gemini-chart-validation-custom", follow_up_model="gemini-follow-up-custom", compaction_model="gemini-compaction-custom", ), @@ -466,14 +448,6 @@ def test_configured_values_apply_to_the_roles_they_govern() -> None: assert config.resolve_role_config(GeminiRole.SYNTHESIS) == RoleConfig( "gemini-synthesis-custom", 0.3, types.ThinkingLevel.LOW ) - assert config.resolve_role_config(GeminiRole.CHART_CONFIG) == RoleConfig( - "gemini-chart-config-custom", 0.2, types.ThinkingLevel.LOW - ) - assert config.resolve_role_config( - GeminiRole.CHART_VALIDATION - ) == RoleConfig( - "gemini-chart-validation-custom", 0.0, types.ThinkingLevel.LOW - ) assert config.resolve_role_config(GeminiRole.FOLLOW_UP) == RoleConfig( "gemini-follow-up-custom", 0.8, types.ThinkingLevel.LOW ) diff --git a/narratives/agent/src/narratives_agent/gemini/client_test.py b/narratives/agent/src/narratives_agent/gemini/client_test.py index 2b125fb8..bef5bbaa 100644 --- a/narratives/agent/src/narratives_agent/gemini/client_test.py +++ b/narratives/agent/src/narratives_agent/gemini/client_test.py @@ -41,8 +41,6 @@ StreamChunk, ) from narratives_agent.gemini.schemas import ( - ChartConfigResponse, - DataValidationResponse, FollowUpResponse, QueryClassification, ) @@ -440,37 +438,37 @@ async def test_a_structured_output_request_sends_the_schema_and_config( gemini: _FakeGemini, ) -> None: # Test: Request configuration for a structured-output call. - # Situation: A caller passes ChartConfigResponse as response_schema with a + # Situation: A caller passes FollowUpResponse as response_schema with a # system instruction, temperature 0.2, and `ThinkingLevel.LOW`. # Expectation: The request sets responseMimeType to "application/json" # with the converted responseSchema, includes the rendered system # instruction and generation settings, and returns a `GeminiResponse` # whose text is the structured reply. - answer = json.dumps({"should_render": True, "charts": []}) + answer = json.dumps({"questions": ["And Spain?"]}) gemini.respond( _answer(200, json=_response_json([{"text": answer}], _USAGE)) ) result = await _call( "async_request", - system_instruction="Extract charts.", + system_instruction="Suggest questions.", temperature=0.2, thinking_level=types.ThinkingLevel.LOW, - response_schema=ChartConfigResponse, + response_schema=FollowUpResponse, ) config = gemini.body()["generationConfig"] assert config["responseMimeType"] == "application/json" - assert config["responseSchema"]["required"] == ["should_render"] + assert config["responseSchema"]["required"] == ["questions"] assert config["temperature"] == 0.2 # The SDK serializes ThinkingConfig fields in snake_case, which the # Gemini API's protobuf JSON parser accepts alongside camelCase. assert config["thinkingConfig"] == {"thinking_level": "LOW"} assert gemini.body()["systemInstruction"]["parts"] == [ - {"text": "rendered:Extract charts."} + {"text": "rendered:Suggest questions."} ] assert isinstance(result, GeminiResponse) - assert json.loads(result.text) == {"should_render": True, "charts": []} + assert json.loads(result.text) == {"questions": ["And Spain?"]} @pytest.mark.asyncio @@ -602,8 +600,6 @@ async def test_an_api_key_crossing_the_error_length_limit_is_redacted( ("schema", "required"), [ (QueryClassification, ["serviceable"]), - (ChartConfigResponse, ["should_render"]), - (DataValidationResponse, ["data_found"]), (FollowUpResponse, ["questions"]), ], ) @@ -611,7 +607,7 @@ async def test_each_schema_is_sent_to_gemini_as_expected( gemini: _FakeGemini, schema: type[BaseModel], required: list[str] ) -> None: # Test: The response schema each workflow sends on the wire. - # Situation: A structured-output request is made with each of the four + # Situation: A structured-output request is made with each of the two # workflow schemas. # Expectation: The request carries the required fields and typed # properties, and omits additional_properties (which the Gemini @@ -629,42 +625,6 @@ async def test_each_schema_is_sent_to_gemini_as_expected( assert set(sent["properties"]) == set(schema.model_fields) -@pytest.mark.asyncio -async def test_the_chart_schema_constrains_each_chart_on_the_wire( - gemini: _FakeGemini, -) -> None: - # Test: The nested chart item in the chart config schema. - # Situation: A request is made with the chart config schema. - # Expectation: Each chart item requires a title, enumerates only the - # supported chart types, marks optional fields nullable, and omits - # additional_properties. - gemini.respond(_answer(200, json=_response_json([{"text": "{}"}]))) - - await client.async_gemini_request( - _MESSAGES, "", _MODEL, response_schema=ChartConfigResponse - ) - - charts = gemini.body()["generationConfig"]["responseSchema"]["properties"][ - "charts" - ] - item = charts["items"] - assert charts["nullable"] is True - assert item["required"] == ["title"] - assert "additional_properties" not in item - assert item["properties"]["viz_type"]["enum"] == [ - "line", - "bar", - "ranking", - "pie", - "highlight", - "gauge", - "scatter", - "slider", - ] - assert item["properties"]["viz_type"]["nullable"] is True - assert item["properties"]["date"]["nullable"] is True - - @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["async_request", "async_thought_streaming"]) @pytest.mark.parametrize( diff --git a/narratives/agent/src/narratives_agent/gemini/schemas.py b/narratives/agent/src/narratives_agent/gemini/schemas.py index ab366b60..a32b5068 100644 --- a/narratives/agent/src/narratives_agent/gemini/schemas.py +++ b/narratives/agent/src/narratives_agent/gemini/schemas.py @@ -22,17 +22,6 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator -ChartVizType = Literal[ - "line", - "bar", - "ranking", - "pie", - "highlight", - "gauge", - "scatter", - "slider", -] - QueryPattern = Literal[ "single_point", "top_n_ranking", @@ -137,58 +126,6 @@ def _check_refusal_consistency(self) -> Self: return self -class ChartItem(BaseModel): - """One chart configuration within a chart config response.""" - - model_config = _SCHEMA_CONFIG - - viz_type: ChartVizType | None = None - # Required to prevent blank UI headers: forces the model to always - # generate a descriptive title rather than rendering an empty header strip. - # Asking for a title in the prompt alone was not enough; the model dropped - # it often enough to notice. - title: str = Field(description="Descriptive chart title") - variable_dcids: list[str] | None = None - place_dcids: list[str] | None = None - parent_place: str | None = None - child_place_type: str | None = None - date: str | None = Field( - default=None, - description="Single comparison date in YYYY, YYYY-MM, or YYYY-MM-DD", - ) - - -class ChartConfigResponse(BaseModel): - """Chart configurations, one per group of compatible variables.""" - - model_config = _SCHEMA_CONFIG - - should_render: bool = Field( - description="True if at least one chart should be rendered" - ) - charts: list[ChartItem] | None = Field( - default=None, - description=( - "Array of chart configurations (max 3). Group compatible " - "variables together." - ), - ) - - -class DataValidationResponse(BaseModel): - """Verdict on whether a synthesized answer contains actual data.""" - - model_config = _SCHEMA_CONFIG - - data_found: bool = Field( - description=( - "True if the response contains actual data/statistics that " - "answer the query. False if data is unavailable, not found, " - "or the response says data doesn't exist." - ) - ) - - class FollowUpResponse(BaseModel): """Suggested follow-up questions for a completed answer.""" @@ -199,7 +136,5 @@ class FollowUpResponse(BaseModel): ) -# The workflows pass these names as `response_schema`. -CHART_CONFIG_SCHEMA = ChartConfigResponse -DATA_VALIDATION_SCHEMA = DataValidationResponse +# Passed by the follow-up workflow as `response_schema`. FOLLOW_UP_SCHEMA = FollowUpResponse diff --git a/narratives/agent/src/narratives_agent/gemini/schemas_test.py b/narratives/agent/src/narratives_agent/gemini/schemas_test.py index 22284826..f108331b 100644 --- a/narratives/agent/src/narratives_agent/gemini/schemas_test.py +++ b/narratives/agent/src/narratives_agent/gemini/schemas_test.py @@ -19,10 +19,6 @@ from pydantic import BaseModel, ValidationError from narratives_agent.gemini.schemas import ( - ChartConfigResponse, - ChartItem, - ChartVizType, - DataValidationResponse, FollowUpResponse, GeographicShape, QueryClassification, @@ -32,60 +28,20 @@ _MODELS: list[type[BaseModel]] = [ QueryClassification, - ChartItem, - ChartConfigResponse, - DataValidationResponse, FollowUpResponse, ] # Holds the smallest valid payload for each schema model. _MINIMAL_PAYLOADS: dict[type[BaseModel], dict[str, Any]] = { QueryClassification: {"serviceable": True}, - ChartItem: {"title": "Population of California"}, - ChartConfigResponse: {"should_render": False}, - DataValidationResponse: {"data_found": True}, FollowUpResponse: {"questions": []}, } -def test_a_chart_config_response_parses_nested_charts() -> None: - # Test: Parsing a complete chart config response. - # Situation: The model returns JSON with one fully populated chart. - # Expectation: The JSON parses into a ChartConfigResponse with all - # provided chart fields populated and omitted optional fields set to - # None. - response = ChartConfigResponse.model_validate_json( - """{ - "should_render": true, - "charts": [{ - "viz_type": "line", - "title": "Population over time", - "variable_dcids": ["Count_Person"], - "place_dcids": ["geoId/06"], - "date": "2020" - }] - }""" - ) - - assert response.should_render is True - assert response.charts is not None - chart = response.charts[0] - assert chart.viz_type == "line" - assert chart.title == "Population over time" - assert chart.variable_dcids == ["Count_Person"] - assert chart.place_dcids == ["geoId/06"] - assert chart.date == "2020" - assert chart.parent_place is None - assert chart.child_place_type is None - - @pytest.mark.parametrize( ("model", "required_field"), [ (QueryClassification, "serviceable"), - (ChartItem, "title"), - (ChartConfigResponse, "should_render"), - (DataValidationResponse, "data_found"), (FollowUpResponse, "questions"), ], ) @@ -107,33 +63,6 @@ def test_a_missing_required_field_is_rejected( assert model.model_json_schema()["required"] == [required_field] -@pytest.mark.parametrize("viz_type", get_args(ChartVizType)) -def test_every_supported_viz_type_is_accepted(viz_type: str) -> None: - # Test: Accepted chart types. - # Situation: ChartItem is validated with each visualization type in - # ChartVizType. - # Expectation: Validation succeeds and preserves viz_type for every - # supported value. - chart = ChartItem.model_validate({"title": "t", "viz_type": viz_type}) - - assert chart.viz_type == viz_type - - -def test_an_unsupported_viz_type_is_rejected() -> None: - # Test: Chart type enforcement. - # Situation: ChartItem is given a visualization type ("histogram") that - # the Data Commons web components do not support. - # Expectation: Validation fails on viz_type, and the JSON schema - # enumerates only the supported types. - with pytest.raises(ValidationError) as raised: - ChartItem.model_validate({"title": "t", "viz_type": "histogram"}) - - assert raised.value.errors()[0]["loc"] == ("viz_type",) - viz_schema = ChartItem.model_json_schema()["properties"]["viz_type"] - enums = [option.get("enum") for option in viz_schema["anyOf"]] - assert list(get_args(ChartVizType)) in enums - - def test_query_classification_populates_defaults() -> None: # Test: Default values on `QueryClassification` when only `serviceable` is # supplied. diff --git a/narratives/agent/src/narratives_agent/server/routes/chat_test.py b/narratives/agent/src/narratives_agent/server/routes/chat_test.py index 174cf1da..645850be 100644 --- a/narratives/agent/src/narratives_agent/server/routes/chat_test.py +++ b/narratives/agent/src/narratives_agent/server/routes/chat_test.py @@ -42,13 +42,14 @@ from narratives_agent.config_model import AgentConfig from narratives_agent.server.routes import chat from narratives_agent.workflows import transcript as transcript_module +from narratives_agent.workflows.blocks import TextBlock from narratives_agent.workflows.chat_pipeline import Emit from narratives_agent.workflows.messages import DISCOVERY_STATUS from narratives_agent.workflows.sse_frames import ( + BlockContentFrame, SseFrame, StatusFrame, TerminalCompleteFrame, - TextFrame, ThoughtFrame, ) from narratives_agent.workflows.transcript import ( @@ -72,7 +73,7 @@ ) _FRAMES: list[SseFrame] = [ _STATUS, - TextFrame(text="About 68 million."), + BlockContentFrame(block=TextBlock(index=0, markdown="About 68 million.")), _UNSIGNED_COMPLETE, ] @@ -157,7 +158,8 @@ def test_stream_sends_each_event_as_a_typed_frame_byte_for_byte( b'data: {"phase": "discovery", "message": "Querying data tools..."}' b"\r\n\r\n" b"id: 2\r\nevent: content\r\n" - b'data: {"text": "About 68 million."}\r\n\r\n' + b'data: {"block": {"index": 0, "type": "text", ' + b'"markdown": "About 68 million."}}\r\n\r\n' b"id: 3\r\nevent: terminal\r\n" b'data: {"state": "complete", ' b'"idempotency_key": "test-idempotency-key"}\r\n\r\n' diff --git a/narratives/agent/src/narratives_agent/telemetry.py b/narratives/agent/src/narratives_agent/telemetry.py index 98ba8ef1..745f25bc 100644 --- a/narratives/agent/src/narratives_agent/telemetry.py +++ b/narratives/agent/src/narratives_agent/telemetry.py @@ -410,6 +410,9 @@ class TurnTelemetry: total_rows_fetched: Total observation rows across all fetched tables. processing_retries: Number of rejected `submit_specification` attempts that returned to discovery for a retry. + charts_emitted: Number of resolved chart blocks emitted during the turn. + charts_rejected: Number of chart blocks dropped by the block splitter or + chart resolver during the turn. """ turn_id: str = field(default_factory=lambda: uuid.uuid4().hex) @@ -421,11 +424,13 @@ class TurnTelemetry: # the turn was refused or ended before classification. pattern: QueryPattern | None = None # The scripts the turn's specification ran, in step order, and the counts - # of its fetches and processing retries. + # of its fetches, processing retries, and charts. scripts: list[str] = field(default_factory=list) tables_fetched: int = 0 total_rows_fetched: int = 0 processing_retries: int = 0 + charts_emitted: int = 0 + charts_rejected: int = 0 _start: float = field(default_factory=time.monotonic) _span: trace.Span = field( default_factory=lambda: _tracer.start_span("narratives.turn") @@ -511,6 +516,8 @@ def finish( "tables_fetched": self.tables_fetched, "total_rows_fetched": self.total_rows_fetched, "processing_retries": self.processing_retries, + "charts_emitted": self.charts_emitted, + "charts_rejected": self.charts_rejected, } if self.pattern is not None and state != TerminalState.REFUSED: record["pattern"] = self.pattern @@ -524,6 +531,8 @@ def finish( "narratives.tables_fetched": self.tables_fetched, "narratives.total_rows_fetched": self.total_rows_fetched, "narratives.processing_retries": self.processing_retries, + "narratives.charts_emitted": self.charts_emitted, + "narratives.charts_rejected": self.charts_rejected, **{ f"narratives.tokens.{key}": value for key, value in tokens.items() diff --git a/narratives/agent/src/narratives_agent/telemetry_test.py b/narratives/agent/src/narratives_agent/telemetry_test.py index 373bc9fd..117c5a31 100644 --- a/narratives/agent/src/narratives_agent/telemetry_test.py +++ b/narratives/agent/src/narratives_agent/telemetry_test.py @@ -304,6 +304,8 @@ def test_turn_summary_records_operations_without_content( turn.tables_fetched = 2 turn.total_rows_fetched = 1_200 turn.processing_retries = 1 + turn.charts_emitted = 2 + turn.charts_rejected = 1 turn.tokens.add( telemetry.TokenCounts(prompt=10, candidates=0, thoughts=0, total=15) ) @@ -328,6 +330,8 @@ def test_turn_summary_records_operations_without_content( assert record["tables_fetched"] == 2 assert record["total_rows_fetched"] == 1_200 assert record["processing_retries"] == 1 + assert record["charts_emitted"] == 2 + assert record["charts_rejected"] == 1 assert record["tokens"] == { "prompt_token_count": 10, "candidates_token_count": 4, @@ -351,6 +355,8 @@ def test_turn_summary_records_operations_without_content( "tables_fetched", "total_rows_fetched", "processing_retries", + "charts_emitted", + "charts_rejected", } @@ -361,9 +367,9 @@ def test_repeated_phases_accumulate_and_analysis_fields_reach_the_span( # Situation: The module tracer records spans in memory; a turn times # the processing phase twice, as a retried specification does, with # the clock advanced by 100 ms and then 50 ms, and records the - # analysis counts before finishing. + # analysis and chart counts before finishing. # Expectation: The processing duration is the sum of both runs, and - # the turn span carries the scripts and the three counts. + # the turn span carries the scripts and the five counts. exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -384,6 +390,8 @@ def read_clock() -> float: turn.tables_fetched = 1 turn.total_rows_fetched = 30 turn.processing_retries = 1 + turn.charts_emitted = 1 + turn.charts_rejected = 2 turn.finish(telemetry.TerminalState.COMPLETE) @@ -398,6 +406,8 @@ def read_clock() -> float: assert attributes["narratives.tables_fetched"] == 1 assert attributes["narratives.total_rows_fetched"] == 30 assert attributes["narratives.processing_retries"] == 1 + assert attributes["narratives.charts_emitted"] == 1 + assert attributes["narratives.charts_rejected"] == 2 @pytest.mark.parametrize( diff --git a/narratives/agent/src/narratives_agent/workflows/chart_config.py b/narratives/agent/src/narratives_agent/workflows/chart_config.py deleted file mode 100644 index 38cb3144..00000000 --- a/narratives/agent/src/narratives_agent/workflows/chart_config.py +++ /dev/null @@ -1,235 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Generates chart configurations and validates whether to display charts.""" - -import json -import logging -from typing import Any - -from narratives_agent.config_model import AgentConfig, GeminiRole -from narratives_agent.gemini.client import ( - GeminiBlockedError, - GeminiError, - async_gemini_request, -) -from narratives_agent.gemini.schemas import ( - CHART_CONFIG_SCHEMA, - DATA_VALIDATION_SCHEMA, -) -from narratives_agent.telemetry import TokenUsage - -logger = logging.getLogger(__name__) - -# Max characters of the synthesis text sent to the data-validation check. -SYNTHESIS_PREVIEW_LENGTH = 2000 - -# Max characters of a parsed payload quoted in error logs. -_LOG_EXCERPT_LENGTH = 300 - - -async def get_chart_config( - mcp_results: str, - user_message: str, - config: AgentConfig, - token_usage: TokenUsage | None = None, -) -> dict[str, Any]: - """Get chart configuration using structured output. - - Supports multiple charts for variables with different units/scales. - - Args: - mcp_results: The gathered tool results to chart. - user_message: The user's query. - config: The agent configuration; the request runs on - `gemini.chart_config_model`. - token_usage: Optional turn totals that receive the call's token - counts. - - Returns: - The parsed chart configuration, or `{"should_render": False}` if the - request fails or its response cannot be parsed. - """ - role = config.resolve_role_config(GeminiRole.CHART_CONFIG) - - prompt = f"""Based on the data query and results, determine chart \ -configurations. - -User Query: {user_message} - -Data Results: -{mcp_results if mcp_results else "No data results"} - -Instructions: -1. Extract variable DCIDs and place DCIDs from the results -2. Analyze variable units from source_metadata and data scales from the values -3. Group variables that can be meaningfully compared on the same Y-axis: - - Same unit type (e.g., both INR, both counts, both percentages) - - Similar magnitude (within ~100x of each other) -4. Create SEPARATE charts for incompatible variable groups: - - Different unit types should be separate (e.g., "Count" vs "INR" vs \ -"Percentage") - - Vastly different scales should be separate (e.g., millions vs trillions) -5. MAXIMUM 3 charts - if more groups exist, prioritize most relevant to the \ -query -6. Choose appropriate viz_type for each chart (line for time series, bar for \ -comparison) -7. Give each chart a descriptive title related to data it is showing but do \ -NOT include year/date in the title. -8. For ALL bar charts, ALWAYS include a date field: - - date can be in formats YYYY, YYYY-MM, or YYYY-MM-DD (e.g., "2021", \ -"2022", "2021-01", "2022-01-01") - - prefer explicit date from user query or tool result context - - if no reliable explicit date is available, omit date - -Set should_render to false if no meaningful data for visualization.""" - - try: - response = await async_gemini_request( - messages=[{"role": "user", "parts": [{"text": prompt}]}], - system_instruction=( - "You are a data visualization expert. Extract chart " - "configurations from data results, grouping compatible " - "variables together and separating incompatible ones into " - "multiple charts." - ), - model=role.model, - temperature=role.temperature, - thinking_level=role.thinking_level, - response_schema=CHART_CONFIG_SCHEMA, - token_usage=token_usage, - ) - except GeminiBlockedError as error: - logger.warning( - "Chart config blocked (%s); rendering without charts", - error.block_reason, - ) - return {"should_render": False} - except GeminiError: - # The client has already logged the redacted upstream error. - logger.warning("Chart config request failed; rendering without charts") - return {"should_render": False} - - try: - if response.text.strip(): - chart_config = json.loads(response.text) - # The model can return valid JSON that is not an object, such as - # `null` or `[]`. The caller expects a dict, so anything else is - # treated as "no charts". - if not isinstance(chart_config, dict): - logger.error( - "Chart config returned %s instead of an object", - type(chart_config).__name__, - ) - return {"should_render": False} - # The configuration is derived from the user's query, so it is - # logged only at DEBUG, below the agent's INFO log level. - logger.debug( - "Chart config result: %s", json.dumps(chart_config, indent=2) - ) - return chart_config - logger.warning( - "Chart config returned no text (finish reason: %s); rendering " - "without charts", - response.finish_reason, - ) - except Exception as error: - logger.error("Chart config parse error: %s", error) - - return {"should_render": False} - - -async def validate_data_response( - synthesis_text: str, - user_message: str, - config: AgentConfig, - token_usage: TokenUsage | None = None, -) -> bool: - """Validate via Gemini whether the synthesis text contains chartable data. - - Returns False so that charts are suppressed unless explicitly confirmed: - with a warning if the validation request fails or is blocked (the client - has already logged the upstream failure), and with an error if the - response is malformed JSON or omits the boolean `data_found` field. - - Args: - synthesis_text: The streamed answer to judge. - user_message: The user's query. - config: The agent configuration; the request runs on - `gemini.chart_validation_model`. - token_usage: Optional turn totals that receive the call's token - counts. - """ - role = config.resolve_role_config(GeminiRole.CHART_VALIDATION) - - prompt = f"""User asked: {user_message} - -Response given: -{synthesis_text[:SYNTHESIS_PREVIEW_LENGTH]} - -Did this response contain actual data/statistics that answer the user's \ -question? -Return false if the response says data is "not available", "not found", \ -"doesn't exist", or similar.""" - - try: - response = await async_gemini_request( - messages=[{"role": "user", "parts": [{"text": prompt}]}], - system_instruction=( - "You validate if a response contains actual data." - ), - model=role.model, - temperature=role.temperature, - thinking_level=role.thinking_level, - response_schema=DATA_VALIDATION_SCHEMA, - token_usage=token_usage, - ) - except GeminiBlockedError as error: - logger.warning( - "Data validation blocked (%s) on %s, hiding charts.", - error.block_reason, - role.model, - ) - return False - except GeminiError: - # The client has already logged the redacted upstream error. - logger.warning( - "Data validation request to %s failed, hiding charts.", role.model - ) - return False - if not response.text.strip(): - logger.error( - "Data validation got no text from %s (finish reason: %s), hiding " - "charts.", - role.model, - response.finish_reason, - ) - return False - - try: - result = json.loads(response.text) - except Exception as error: - logger.error("Data validation parse error, hiding charts: %s", error) - return False - - verdict = result.get("data_found") if isinstance(result, dict) else None - if not isinstance(verdict, bool): - logger.error( - "Data validation returned no data_found boolean, hiding charts. " - "Parsed: %s", - str(result)[:_LOG_EXCERPT_LENGTH], - ) - return False - return verdict diff --git a/narratives/agent/src/narratives_agent/workflows/chart_config_test.py b/narratives/agent/src/narratives_agent/workflows/chart_config_test.py deleted file mode 100644 index 9da677c9..00000000 --- a/narratives/agent/src/narratives_agent/workflows/chart_config_test.py +++ /dev/null @@ -1,296 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tests for the chart configuration workflow. - -Verifies that `validate_data_response` fails closed (returns `False`) when the -Gemini validation call raises `GeminiError` or `GeminiBlockedError`, or -returns invalid JSON or a JSON object without a boolean `data_found` field, -and returns the model's boolean verdict when `data_found` is present. Verifies -that `get_chart_config` returns a dict whatever JSON value the model produces -or error the client raises, and that both calls request a thinking level that -`gemini-3.8-flash` accepts. -""" - -import json -import logging -from collections.abc import Callable, Coroutine -from typing import Any - -import pytest -from google.genai import types - -from narratives_agent.config_model import AgentConfig -from narratives_agent.gemini.client import GeminiBlockedError, GeminiError -from narratives_agent.gemini.response import GeminiResponse -from narratives_agent.workflows import chart_config - -_SYNTHESIS = "Population reached 39,538,223 in 2020." -_QUESTION = "What is the population of California?" -_RESULTS = "Tool: get_observations\nResult: 39538223" -_CONFIG = AgentConfig() - -type GeminiAnswers = Callable[[GeminiResponse | Exception], None] - - -def _reply(text: str) -> GeminiResponse: - """Returns a model response whose text is `text`.""" - return GeminiResponse(text=text) - - -@pytest.fixture -def gemini_answers(monkeypatch: pytest.MonkeyPatch) -> GeminiAnswers: - """Stub `chart_config.async_gemini_request` with a response or an error. - - An exception is raised by the stub, as the Gemini client raises on a - failed or blocked request. - """ - - def install(response: GeminiResponse | Exception) -> None: - async def request(**kwargs: Any) -> GeminiResponse: - if isinstance(response, Exception): - raise response - return response - - monkeypatch.setattr(chart_config, "async_gemini_request", request) - - return install - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "error", - [ - pytest.param(GeminiError("HTTP 404: model not found"), id="error"), - pytest.param(GeminiBlockedError("SAFETY"), id="blocked"), - ], -) -async def test_a_failed_or_blocked_call_hides_the_charts( - gemini_answers: GeminiAnswers, - caplog: pytest.LogCaptureFixture, - error: Exception, -) -> None: - # Test: Chart suppression when a Gemini call fails or is blocked. - # Situation: async_gemini_request raises `GeminiError` or - # `GeminiBlockedError`. - # Expectation: validate_data_response returns False and get_chart_config - # returns `{"should_render": False}`, so no charts are shown. Each call - # logs one WARNING, not an ERROR, because the client already logged the - # redacted upstream error; a block names its reason. - gemini_answers(error) - - with caplog.at_level(logging.WARNING, logger=chart_config.logger.name): - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is False - ) - assert await chart_config.get_chart_config( - _RESULTS, _QUESTION, _CONFIG - ) == {"should_render": False} - - assert [record.levelno for record in caplog.records] == [ - logging.WARNING, - logging.WARNING, - ] - assert ("SAFETY" in caplog.text) is isinstance(error, GeminiBlockedError) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "response", - [ - pytest.param(GeminiResponse(finish_reason="MAX_TOKENS"), id="empty"), - pytest.param(GeminiResponse(text=" \n"), id="blank-text"), - ], -) -async def test_an_empty_response_hides_the_charts( - gemini_answers: GeminiAnswers, - caplog: pytest.LogCaptureFixture, - response: GeminiResponse, -) -> None: - # Test: Chart suppression when Gemini returns nothing to parse. - # Situation: async_gemini_request returns a reply with no text, or with - # only whitespace. - # Expectation: validate_data_response returns False and logs one ERROR; - # get_chart_config returns `{"should_render": False}` and logs one - # WARNING rather than a parse error, because there is nothing to parse. - # Both log lines name the finish reason. - gemini_answers(response) - - with caplog.at_level(logging.WARNING, logger=chart_config.logger.name): - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is False - ) - assert await chart_config.get_chart_config( - _RESULTS, _QUESTION, _CONFIG - ) == {"should_render": False} - - assert [record.levelno for record in caplog.records] == [ - logging.ERROR, - logging.WARNING, - ] - assert all( - f"finish reason: {response.finish_reason}" in record.getMessage() - for record in caplog.records - ) - - -@pytest.mark.asyncio -async def test_a_response_that_is_not_json_hides_the_charts( - gemini_answers: GeminiAnswers, -) -> None: - # Test: Chart suppression when the validation response text is not valid - # JSON. - # Situation: The response text is unparseable prose instead of a JSON - # object. - # Expectation: validate_data_response catches the parse error and returns - # False. - gemini_answers(_reply("I think so, yes.")) - - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is False - ) - - -@pytest.mark.asyncio -async def test_a_verdict_missing_the_field_hides_the_charts( - gemini_answers: GeminiAnswers, -) -> None: - # Test: Chart suppression when the parsed JSON omits `data_found`. - # Situation: The model returns a valid JSON object without the required - # boolean `data_found` property. - # Expectation: validate_data_response returns False instead of defaulting - # a missing field to True. - gemini_answers(_reply(json.dumps({"reason": "unsure"}))) - - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is False - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("payload", ["null", "[]", "false", "123"]) -async def test_a_non_object_json_payload_hides_the_charts( - gemini_answers: GeminiAnswers, payload: str -) -> None: - # Test: Chart suppression when the validation response is valid non-object - # JSON. - # Situation: json.loads succeeds and returns None, a list, a bool, or an - # int rather than a dict. - # Expectation: validate_data_response returns False without raising - # AttributeError. - gemini_answers(_reply(payload)) - - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is False - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("data_found", [True, False]) -async def test_a_clear_verdict_is_passed_through( - gemini_answers: GeminiAnswers, data_found: bool -) -> None: - # Test: Pass-through of a valid boolean `data_found` verdict. - # Situation: The model returns a JSON object with `data_found` set to True - # or False. - # Expectation: validate_data_response returns the exact boolean value from - # the response. - gemini_answers(_reply(json.dumps({"data_found": data_found}))) - - assert ( - await chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ) - is data_found - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("payload", ["null", "[]", "false", "123"]) -async def test_a_non_object_chart_config_is_replaced_by_no_charts( - gemini_answers: GeminiAnswers, payload: str -) -> None: - # Test: `get_chart_config` never returns a value that is not a dict. - # Situation: The chart-config response is valid JSON that is None, a - # list, a bool, or an int rather than an object. - # Expectation: `get_chart_config` returns `{"should_render": False}` - # rather than the parsed value, so the caller can index the result. - gemini_answers(_reply(payload)) - - assert await chart_config.get_chart_config( - _RESULTS, _QUESTION, _CONFIG - ) == {"should_render": False} - - -@pytest.mark.asyncio -async def test_an_object_chart_config_is_passed_through( - gemini_answers: GeminiAnswers, -) -> None: - # Test: Pass-through of a well-formed chart configuration. - # Situation: The chart-config response is a JSON object. - # Expectation: `get_chart_config` returns that object unchanged. - config = {"should_render": True, "charts": [{"title": "Population"}]} - gemini_answers(_reply(json.dumps(config))) - - assert ( - await chart_config.get_chart_config(_RESULTS, _QUESTION, _CONFIG) - == config - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "call", - [ - lambda: chart_config.get_chart_config(_RESULTS, _QUESTION, _CONFIG), - lambda: chart_config.validate_data_response( - _SYNTHESIS, _QUESTION, _CONFIG - ), - ], - ids=["get_chart_config", "validate_data_response"], -) -async def test_chart_calls_request_the_low_thinking_level( - monkeypatch: pytest.MonkeyPatch, - call: Callable[[], Coroutine[Any, Any, object]], -) -> None: - # Test: The thinking level that the chart calls send to Gemini. - # Situation: The chart configuration and chart validation calls run. - # Expectation: Each requests `ThinkingLevel.LOW`. `gemini-3.8-flash` - # rejects `MINIMAL` with a 400 error. These calls catch that error and - # return no charts, so the failure would be silent. - levels: list[types.ThinkingLevel | None] = [] - - async def request(**kwargs: Any) -> GeminiResponse: - levels.append(kwargs.get("thinking_level")) - raise GeminiError("stubbed") - - monkeypatch.setattr(chart_config, "async_gemini_request", request) - - await call() - - assert levels == [types.ThinkingLevel.LOW] diff --git a/narratives/agent/src/narratives_agent/workflows/chat_pipeline.py b/narratives/agent/src/narratives_agent/workflows/chat_pipeline.py index a5753b69..3fa96a4d 100644 --- a/narratives/agent/src/narratives_agent/workflows/chat_pipeline.py +++ b/narratives/agent/src/narratives_agent/workflows/chat_pipeline.py @@ -14,29 +14,28 @@ """Runs the agent loop, the phases of one chat turn, as native asyncio code. `run_turn` drives a turn through its phases in order: classification, the -discovery loop and specification processing, streamed synthesis, chart -configuration, turn signing, and follow-up questions. Each phase reports -progress by calling the turn's `emit` callback with one typed SSE frame at a -time (an `SseFrame` from `sse_frames`, which declares each frame's event name -and payload) and hands its results to the next phase as typed values rather -than shared mutable state. +discovery loop and specification processing, streamed synthesis with +interleaved text and chart blocks, turn signing, and follow-up questions. Each +phase reports progress by calling the turn's `emit` callback with one typed +SSE frame at a time (an `SseFrame` from `sse_frames`, which declares each +frame's event name and payload) and hands its results to the next phase as +typed values rather than shared mutable state. Every turn that is not canceled ends in one `terminal` frame (`TerminalCompleteFrame`, `TerminalRefusedFrame`, or `TerminalErrorFrame`). Every phase reads the conversation through the verified `Transcript`, and completed turns are appended and signed by `finalize_turn` before the -`complete` frame is sent. Chart configuration runs as an `asyncio.Task` -concurrently with synthesis and is awaited, with a timeout, once synthesis -has finished streaming. +`complete` frame is sent. """ import asyncio import logging -from collections.abc import Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Sequence from dataclasses import dataclass, field, replace from narratives_agent.analytics.digest import AnalysisDigest from narratives_agent.analytics.provenance import Source +from narratives_agent.analytics.tables import ObservationTable from narratives_agent.config import load_config from narratives_agent.config_model import AgentConfig, GeminiRole from narratives_agent.gemini.client import ( @@ -45,15 +44,19 @@ GeminiStreamError, async_gemini_stream, ) +from narratives_agent.gemini.response import StreamChunk from narratives_agent.gemini.schemas import QueryClassification, QueryPattern -from narratives_agent.json_value import JsonObject, JsonValue +from narratives_agent.json_value import JsonObject from narratives_agent.mcp import client as mcp_client from narratives_agent.mcp.capabilities import warn_on_missing_tools from narratives_agent.settings import get_settings from narratives_agent.telemetry import ErrorType, TerminalState, TurnTelemetry -from narratives_agent.workflows.chart_config import ( - get_chart_config, - validate_data_response, +from narratives_agent.workflows.blocks import BlockSplitter, TextBlock +from narratives_agent.workflows.chart_defaults import build_default_charts +from narratives_agent.workflows.charts import ( + ChartSpecification, + ResolvedChart, + resolve_chart_specification, ) from narratives_agent.workflows.classify import ( Clarification, @@ -68,10 +71,8 @@ DiscoveryResult, McpLoopError, execute_discovery_loop, - format_count, ) from narratives_agent.workflows.messages import ( - CHART_CONFIG_STATUS, CHAT_DISABLED, CLASSIFY_STATUS, INTERNAL_ERROR, @@ -83,19 +84,20 @@ SYNTHESIS_STATUS, TRUNCATED_MESSAGE, ) -from narratives_agent.workflows.pattern_guidance import append_pattern_guidance from narratives_agent.workflows.processing import ( StepRefusalDiagnostic, execute_specification, ) +from narratives_agent.workflows.prompt_contracts import ( + compose_system_instruction, +) from narratives_agent.workflows.specification import ( AnalysisSpecification, - ComputeDerivedMetricStep, build_default_specification, get_step_output_kind, ) from narratives_agent.workflows.sse_frames import ( - ChartConfigFrame, + BlockContentFrame, DataStatusFrame, FollowUpsFrame, SourcesFrame, @@ -105,11 +107,15 @@ TerminalCompleteFrame, TerminalErrorFrame, TerminalRefusedFrame, - TextFrame, ThoughtFrame, ToolCallFrame, ) from narratives_agent.workflows.state_slots import build_state_slots +from narratives_agent.workflows.synthesis import ( + SYNTHESIS_CHUNK_TIMEOUT_SECONDS, + build_synthesis_inputs, + render_synthesis_message, +) from narratives_agent.workflows.transcript import ( ConversationStateSlots, SignedWindow, @@ -121,22 +127,28 @@ logger = logging.getLogger(__name__) -# Maximum seconds to wait for the background chart task after synthesis ends. -_CHART_CONFIG_JOIN_TIMEOUT_SECONDS = 5 - # Maximum seconds to wait for the MCP tool list before aborting the turn. _MCP_TOOLS_TIMEOUT_SECONDS = 20 -# Maximum places with data named per table in the synthesis table inventory. -_INVENTORY_PLACE_CAP = 25 +# Maximum places named in a follow-up topic before keeping the variable alone. +_FOLLOW_UP_PLACE_CAP = 3 # Fallback source cited when the digest has data but no table reported a URL. _FALLBACK_SOURCE = Source(name="Data Commons", url="https://datacommons.org/") -_NO_CHARTS: JsonObject = {"should_render": False} - +# Generic fallback note when the discovery loop ends without tables or an +# unavailable report; passed to synthesis and excluded from follow-up topics. _NO_OBSERVATIONS_NOTE = "No observations were retrieved for the question." +# Prefixes of missing-data notes that report a processing failure rather than a +# data gap, excluded from follow-up topics. +_STEP_FAILED_NOTE_PREFIX = "An analysis step failed: " +_DEFAULT_FAILED_NOTE_PREFIX = "The default analysis also failed: " +_DIAGNOSTIC_NOTE_PREFIXES = ( + _STEP_FAILED_NOTE_PREFIX, + _DEFAULT_FAILED_NOTE_PREFIX, +) + # Default pattern recorded in telemetry when the classifier call failed. _UNCLASSIFIED_PATTERN: QueryPattern = "exploratory" @@ -186,31 +198,55 @@ class TurnContext: cache: TurnCache = field(default_factory=TurnCache) +def _build_empty_digest() -> AnalysisDigest: + """Returns the digest of a turn that computed nothing.""" + return AnalysisDigest(results=(), sources=()) + + @dataclass class DiscoveryPhaseResult: """Holds the discovery phase outputs passed to later phases. - `context` is the rendering of the digest and the tables it rests on - that synthesis and the chart planner read. `sources` is the single - authority on citation numbering; see `run_discovery_phase`. `tables` - are the cached tables the specification used, in order of first use, - from which the state slots are built. + `digest` is the analysis the turn is answered from and `spec` is the + specification that produced it, or `None` when the turn has none. + `sources` is the single authority on citation numbering; see + `run_discovery_phase`. `tables` are the cached tables the specification + used, in order of first use, from which the table inventory and the + state slots are built. Attributes: - context: Rendered digest and table inventory text for synthesis. + digest: The resolved `AnalysisDigest` for the turn. + spec: The `AnalysisSpecification` that produced `digest`, or `None` + when no specification ran. has_data: Whether the resolved digest holds computed data. tool_calls: Serialized tool-call records streamed to the client. sources: Ordered provenance sources for citation numbering. tables: Cached tables used by the specification, in first-use order. - chart_task: Background task computing chart configuration, if started. """ - context: str = "" + digest: AnalysisDigest = field(default_factory=_build_empty_digest) + spec: AnalysisSpecification | None = None has_data: bool = False tool_calls: list[JsonObject] = field(default_factory=list) sources: list[Source] = field(default_factory=list) tables: tuple[CachedTable, ...] = () - chart_task: asyncio.Task[JsonObject] | None = None + + +@dataclass(frozen=True) +class SynthesisPhaseResult: + """Holds what the synthesis phase streamed. + + `text` is the answer that is signed: the text blocks joined by blank + lines, with every chart block left out. `charts` are the resolved + charts in the order their blocks were emitted. + + Attributes: + text: The signed prose text with chart blocks excluded. + charts: Resolved charts in emission order. + """ + + text: str + charts: tuple[ResolvedChart, ...] = () def _emit_status(ctx: TurnContext, status: tuple[StatusPhase, str]) -> None: @@ -277,7 +313,7 @@ async def _run_default_specification( ) if retry is None: break - notes.append(f"An analysis step failed: {outcome.message}") + notes.append(f"{_STEP_FAILED_NOTE_PREFIX}{outcome.message}") specification = retry logger.info( "Default specification step %s (%s) refused; running %s", @@ -291,7 +327,7 @@ async def _run_default_specification( replace(outcome, missing_notes=(*outcome.missing_notes, *notes)), specification, ) - notes.append(f"The default analysis also failed: {outcome.message}") + notes.append(f"{_DEFAULT_FAILED_NOTE_PREFIX}{outcome.message}") return ( AnalysisDigest.from_results( ( @@ -338,7 +374,9 @@ async def _resolve_digest( ) notes: list[str] = [] if loop.last_diagnostic is not None: - notes.append(f"An analysis step failed: {loop.last_diagnostic.message}") + notes.append( + f"{_STEP_FAILED_NOTE_PREFIX}{loop.last_diagnostic.message}" + ) specification = build_default_specification(ctx.classification, ctx.cache) if specification is None: notes.append(_NO_OBSERVATIONS_NOTE) @@ -385,114 +423,20 @@ def _collect_tables_used( return tuple(table for table in tables if table is not None) -def _describe_derived_inputs( - specification: AnalysisSpecification | None, -) -> dict[str, str]: - """Maps each derived table's step id to a description of its inputs. - - Args: - specification: The specification that answered the turn, or `None`. - - Returns: - A dictionary mapping derived step IDs to `"derived from and - "`. - """ - if specification is None: - return {} - return { - step.step_id: f"derived from {step.left} and {step.right}" - for step in specification.steps - if isinstance(step, ComputeDerivedMetricStep) - } - - -def render_analysis_context( - digest: AnalysisDigest, - specification: AnalysisSpecification | None, - tables: Sequence[CachedTable], -) -> str: - """Renders the digest and its tables as the text synthesis reads. - - This adapter feeds the existing synthesis prompt and chart planner, - which read one block of text, a rendering of the digest in place of raw - tool results. Everything in the text is computed or fetched by the - server: the digest's results and missing-data notes, then one line per - table. The specification's `findings` are not included, because they - are the model's reading of five sample rows and the digest holds the - computed figures. The chart planner extracts variable DCIDs, place - DCIDs, parent places, and child place types from the text, so the - inventory lists each fetched table's identity and places alongside the - computed figures. A derived table is listed by name only, marked as not - a Data Commons variable, because its DCID names no series the chart web - components could fetch. - - Args: - digest: The computed `AnalysisDigest` for the turn. - specification: The `AnalysisSpecification` that produced `digest`, or - `None` when no specification ran. - tables: The cached tables used by `specification`. - - Returns: - The formatted digest and table inventory text. - """ - lines = [digest.format_for_prompt()] - if tables: - lines.append("\nTables used:") - derived_inputs = _describe_derived_inputs(specification) - for cached in tables: - table_id = cached.table_id - table = cached.table - dates = sorted({row.date for row in table.rows}) - # table.places_with_data is already sorted and deduplicated by - # ObservationTable. - places = table.places_with_data - named = ", ".join( - f"{table.place_names.get(place, place)} ({place})" - for place in places[:_INVENTORY_PLACE_CAP] - ) - if len(places) > _INVENTORY_PLACE_CAP: - extra = len(places) - _INVENTORY_PLACE_CAP - named += f", and {format_count(extra, 'other place')}" - scope = ( - f"every {table.child_place_type} in {table.parent_place_dcid}" - if table.child_place_type and table.parent_place_dcid - else format_count(len(places), "place") - ) - span = f"{dates[0]} to {dates[-1]}" if dates else "no dates" - if table_id in derived_inputs: - identity = ( - f"{table_id} ({derived_inputs[table_id]}; not a " - f"Data Commons variable, do not chart it): " - f"{table.variable_name}" - ) - else: - identity = ( - f"{table_id}: {table.variable_name} ({table.variable_dcid})" - ) - unit = table.unit or "unknown" - if table.scaling_factor is not None: - unit += f" (scaling factor {table.scaling_factor:g})" - lines.append( - f"- {identity}; unit: {unit}; {scope}; dates {span}; places with " - f"data: {named or 'none'}" - ) - return "\n".join(lines) - - async def run_discovery_phase(ctx: TurnContext) -> DiscoveryPhaseResult: """Finds and analyzes the data, the discovery and processing phases. Runs the discovery loop, resolves the digest the turn is answered from, streams the tool calls, the data status, and the provenance sources, - records the phase's telemetry counts, and starts chart configuration in - the background when the digest has data. + and records the phase's telemetry counts. Args: ctx: The turn context holding the transcript, config, telemetry, and cache. Returns: - The populated `DiscoveryPhaseResult` for synthesis and chart planning. + The populated `DiscoveryPhaseResult` for synthesis, state slots, and + follow-up generation. Raises: TurnAbortedError: If no MCP tools are available or the loop fails. @@ -546,6 +490,8 @@ async def run_discovery_phase(ctx: TurnContext) -> DiscoveryPhaseResult: ctx.emit(ToolCallFrame(tool_call=call)) digest, specification = await _resolve_digest(ctx, loop) + result.digest = digest + result.spec = specification result.has_data = digest.has_data result.tables = _collect_tables_used(specification, ctx.cache) ctx.telemetry.scripts = ( @@ -577,106 +523,71 @@ async def run_discovery_phase(ctx: TurnContext) -> DiscoveryPhaseResult: result.sources = [_FALLBACK_SOURCE] if result.sources: ctx.emit(SourcesFrame(sources=result.sources)) - - result.context = render_analysis_context( - digest, specification, result.tables - ) - # Only generate charts when the digest holds computed figures. - if result.has_data: - result.chart_task = asyncio.create_task( - get_chart_config( - result.context, - ctx.transcript.current_query, - ctx.config, - ctx.telemetry.tokens, - ), - name="chart-config", - ) - else: - logger.info( - "Chart config skipped: the digest has no data (truncated=%s, " - "unavailable=%s)", - loop.truncated, - loop.unavailable is not None, - ) return result -def _build_synthesis_message( - user_message: str, discovery: DiscoveryPhaseResult -) -> str: - """Builds the final user message that carries the analysis. - - Args: - user_message: The user's current query text. - discovery: The completed discovery phase result. - - Returns: - The formatted user prompt combining numbered sources and data results. - """ - context_parts = [] - if discovery.context: - # Number sources 1..N in the order streamed to the UI as `sources`. - if discovery.sources: - numbered_sources = "\n".join( - f"[{number}] {source.name or source.url} - {source.url}" - for number, source in enumerate(discovery.sources, start=1) - ) - context_parts.append( - "**NUMBERED SOURCES - cite these numbers, exactly as given:**\n" - f"{numbered_sources}" - ) - context_parts.append(f"**DATA RESULTS:**\n{discovery.context}") - context = ( - "\n".join(context_parts) - if context_parts - else "No additional context available." - ) - return ( - f"User Query: {user_message}\n\n{context}\n\n" - "Please provide a comprehensive response combining all available " - "information." - ) - - async def run_synthesis_phase( ctx: TurnContext, discovery: DiscoveryPhaseResult -) -> str: - """Streams the narrative answer, the synthesis phase. +) -> SynthesisPhaseResult: + """Streams the answer as text and chart blocks, the synthesis phase. + + The user message is rendered from the typed inputs at the last moment, + the stream is split into blocks as it arrives, and every chart block is + resolved against the turn's table inventory and cache before it is + emitted. When the digest has data and no chart block was emitted, at + most one default chart from the primary result is appended. Opening the + stream and reading each chunk are bounded by + `SYNTHESIS_CHUNK_TIMEOUT_SECONDS`, so a stalled stream ends the turn + rather than holding it open for the whole HTTP timeout. Args: ctx: The turn context holding the transcript, config, and telemetry. discovery: The discovery phase result to synthesize into prose. Returns: - The complete answer text. + The signed text and the resolved charts. Raises: TurnAbortedError: If the synthesis request fails, its stream breaks - part-way, or it returns no answer text. + or stalls part-way, or it returns no prose. GeminiBlockedError: If Gemini blocks the synthesis prompt or the - candidate before any answer text has streamed. + candidate before any answer block has been emitted. """ ctx.emit(StatusFrame(phase="synthesis", message=SYNTHESIS_STATUS)) config = ctx.config role = config.resolve_role_config(GeminiRole.SYNTHESIS) + inputs = build_synthesis_inputs( + ctx.transcript.current_query, + discovery.digest, + discovery.spec, + discovery.tables, + discovery.sources, + ) # Put the verified window first, then the query with the analysis context. messages = transcript_contents( - ctx.transcript, - _build_synthesis_message(ctx.transcript.current_query, discovery), + ctx.transcript, render_synthesis_message(inputs) ) try: - stream = await async_gemini_stream( - messages=messages, - system_instruction=append_pattern_guidance( - config.prompts.synthesis, ctx.classification, "synthesis" - ), - model=role.model, - temperature=role.temperature, - thinking_level=role.thinking_level, - token_usage=ctx.telemetry.tokens, - include_thoughts=True, + async with asyncio.timeout(SYNTHESIS_CHUNK_TIMEOUT_SECONDS): + stream = await async_gemini_stream( + messages=messages, + system_instruction=compose_system_instruction( + config.prompts.synthesis, "synthesis", ctx.classification + ), + model=role.model, + temperature=role.temperature, + thinking_level=role.thinking_level, + token_usage=ctx.telemetry.tokens, + include_thoughts=True, + ) + except TimeoutError as error: + logger.warning( + "Synthesis stream stalled for %.0f seconds", + SYNTHESIS_CHUNK_TIMEOUT_SECONDS, ) + raise TurnAbortedError( + ErrorType.SYNTHESIS_STREAM_ERROR, SYNTHESIS_FAILED + ) from error except GeminiBlockedError: # Propagates to `run_turn`, which ends the turn as `refused`. raise @@ -687,18 +598,88 @@ async def run_synthesis_phase( ErrorType.SYNTHESIS_ERROR, SYNTHESIS_FAILED ) from error - full_text = "" + allowed_table_ids = {entry.table_id for entry in discovery.tables} + + def resolve(spec: ChartSpecification) -> ResolvedChart | None: + if not discovery.has_data: + logger.info("Chart block rejected: turn has no data") + return None + for table_id in spec.tables: + if table_id not in allowed_table_ids: + logger.info( + "Chart block rejected: %s is not in the table inventory", + table_id, + ) + return None + return resolve_chart_specification( + spec, ctx.cache, discovery.digest, discovery.spec + ) + + splitter = BlockSplitter( + on_block=lambda block: ctx.emit(BlockContentFrame(block=block)), + resolve_chart=resolve, + ) + try: + await _consume_synthesis_stream(ctx, stream, splitter) + splitter.finish() + if not splitter.text: + logger.warning("Synthesis returned no answer text") + raise TurnAbortedError(ErrorType.SYNTHESIS_EMPTY, SYNTHESIS_FAILED) + if discovery.has_data and not splitter.charts: + for default in build_default_charts( + discovery.digest, discovery.spec, ctx.cache + ): + splitter.append_chart(default) + finally: + ctx.telemetry.charts_emitted = splitter.charts_emitted + ctx.telemetry.charts_rejected = splitter.charts_rejected + return SynthesisPhaseResult(splitter.text, tuple(splitter.charts)) + + +async def _consume_synthesis_stream( + ctx: TurnContext, + stream: AsyncGenerator[StreamChunk], + splitter: BlockSplitter, +) -> None: + """Feeds the synthesis stream to the splitter and relays its thoughts. + + Args: + ctx: The turn context holding the `emit` callback. + stream: The open Gemini synthesis stream yielding thought and text + chunks. + splitter: The block splitter that buffers paragraphs and chart blocks + and emits each completed block. + + Raises: + TurnAbortedError: If the stream breaks part-way, a chunk does not + arrive within `SYNTHESIS_CHUNK_TIMEOUT_SECONDS`, or Gemini blocks + the candidate after at least one answer block has been emitted. + GeminiBlockedError: If Gemini blocks the candidate before any answer + block has been emitted. + """ try: - async for chunk in stream: + while True: + try: + async with asyncio.timeout(SYNTHESIS_CHUNK_TIMEOUT_SECONDS): + chunk = await anext(stream) + except StopAsyncIteration: + return + except TimeoutError as error: + logger.warning( + "Synthesis stream stalled for %.0f seconds", + SYNTHESIS_CHUNK_TIMEOUT_SECONDS, + ) + raise TurnAbortedError( + ErrorType.SYNTHESIS_STREAM_ERROR, SYNTHESIS_FAILED + ) from error if chunk.kind == "thought": ctx.emit(ThoughtFrame(thought=chunk.text, phase="synthesis")) else: - full_text += chunk.text - ctx.emit(TextFrame(text=chunk.text)) + splitter.feed(chunk.text) except GeminiBlockedError as error: - if not full_text.strip(): + if not splitter.text_blocks and not splitter.charts: raise - # Once answer text has streamed to the client, a mid-stream block + # Once an answer block has streamed to the client, a mid-stream block # leaves the response truncated and is reported as a stream error. raise TurnAbortedError( ErrorType.SYNTHESIS_STREAM_ERROR, SYNTHESIS_FAILED @@ -709,114 +690,90 @@ async def run_synthesis_phase( ) from error finally: await stream.aclose() - if not full_text.strip(): - logger.warning("Synthesis returned no answer text") - raise TurnAbortedError(ErrorType.SYNTHESIS_EMPTY, SYNTHESIS_FAILED) - return full_text -async def resolve_chart_config( - ctx: TurnContext, discovery: DiscoveryPhaseResult, full_text: str -) -> JsonObject: - """Determines whether to display charts and returns their configuration. - - Charts are best-effort because the answer has already streamed when this - phase runs. They are hidden when a second model call finds that the - answer does not contain data or when that call fails, and they are - omitted when chart configuration fails or does not finish within - `_CHART_CONFIG_JOIN_TIMEOUT_SECONDS`. +def _format_topic_scope(table: ObservationTable | None) -> str: + """Formats where a variable was measured as a suffix for a follow-up topic. Args: - ctx: The turn context holding the transcript, config, and telemetry. - discovery: The discovery phase result holding the background chart task. - full_text: The complete synthesized answer text. + table: The observation table for the variable, or `None` when none is + cached. Returns: - The resolved chart configuration JSON object. + A suffix such as `" by County"` or `" in California, Texas"`, or `""` + when no place scope applies. """ - if discovery.chart_task is None: - return dict(_NO_CHARTS) - ctx.emit(StatusFrame(phase="chart_config", message=CHART_CONFIG_STATUS)) - show_charts = True - if full_text: - try: - show_charts = await validate_data_response( - full_text, - ctx.transcript.current_query, - ctx.config, - ctx.telemetry.tokens, - ) - except Exception as error: - logger.warning( - "Chart validation failed with %s; hiding charts", - type(error).__name__, - ) - show_charts = False - else: - if not show_charts: - logger.info("Chart validation found no data; hiding charts") - try: - # `wait_for` cancels the task when the timeout expires. - chart_config = await asyncio.wait_for( - discovery.chart_task, _CHART_CONFIG_JOIN_TIMEOUT_SECONDS - ) - except TimeoutError: - logger.warning( - "Chart config did not finish within %d seconds; rendering " - "without charts", - _CHART_CONFIG_JOIN_TIMEOUT_SECONDS, - ) - chart_config = dict(_NO_CHARTS) - except Exception: - # `_settle_chart_task` logs the task exception when the turn ends. - chart_config = dict(_NO_CHARTS) - if not show_charts: - chart_config["hide_charts"] = True - return chart_config + if table is None: + return "" + if table.child_place_type and table.parent_place_dcid: + return f" by {table.child_place_type}" + names = [ + table.place_names.get(place, place) + for place in table.places_with_data[:_FOLLOW_UP_PLACE_CAP] + ] + return f" in {', '.join(names)}" if names else "" + +def collect_follow_up_topics( + discovery: DiscoveryPhaseResult, charts: Sequence[ResolvedChart] +) -> list[str]: + """Returns the topics that ground the follow-up questions. -def _collect_chart_topics(chart_config: Mapping[str, JsonValue]) -> list[str]: - """Returns the chart titles that ground the follow-up questions. + The topics are the variables the digest's results computed figures + for, each with the places or the child place type it was measured + over, then the titles of the charts shown, then the missing-data notes + that name a data gap (excluding processing diagnostics and the generic + no-observations fallback note). Duplicates are dropped in order of + first appearance. Args: - chart_config: The resolved chart configuration dictionary. + discovery: The completed discovery phase result. + charts: The resolved charts emitted during synthesis. Returns: - A list of non-empty chart title strings. + Deduplicated follow-up topic strings in first-appearance order. """ - charts = chart_config.get("charts") + tables_by_variable: dict[str, ObservationTable] = {} + for cached in discovery.tables: + tables_by_variable.setdefault(cached.table.variable_dcid, cached.table) topics: list[str] = [] - if isinstance(charts, list): - for chart in charts: - if isinstance(chart, dict): - chart_title = chart.get("title") - if isinstance(chart_title, str) and chart_title: - topics.append(chart_title) - title = chart_config.get("title") - if not topics and isinstance(title, str) and title: - # Legacy single-chart shape. - topics.append(title) - return topics + for result in discovery.digest.results: + if result.status != "ok": + continue + for dcid in result.variables: + name = result.variable_names.get(dcid, dcid) + scope = _format_topic_scope(tables_by_variable.get(dcid)) + topics.append(f"{name}{scope}") + topics.extend(chart.title for chart in charts if chart.title) + topics.extend( + note + for note in discovery.digest.missing_notes + if note != _NO_OBSERVATIONS_NOTE + and not note.startswith(_DIAGNOSTIC_NOTE_PREFIXES) + ) + return list(dict.fromkeys(topics)) async def run_followups( - ctx: TurnContext, chart_config: Mapping[str, JsonValue] + ctx: TurnContext, + discovery: DiscoveryPhaseResult, + charts: Sequence[ResolvedChart], ) -> None: - """Emits follow-up questions grounded in the chart topics. + """Emits follow-up questions grounded in the digest and the charts. This phase runs after the terminal event so the UI can display the answer and charts immediately and show the follow-up questions when they - arrive. If there are no chart topics or generation fails, nothing is - emitted. + arrive. If there are no topics or generation fails, nothing is emitted. Args: ctx: The turn context holding the transcript, config, and telemetry. - chart_config: The resolved chart configuration dictionary. + discovery: The completed discovery phase result. + charts: The resolved charts emitted during synthesis. """ try: follow_ups = await generate_follow_up_questions( ctx.transcript.current_query, - _collect_chart_topics(chart_config), + collect_follow_up_topics(discovery, charts), ctx.config, ctx.telemetry.tokens, ) @@ -849,18 +806,17 @@ async def run_turn(transcript: Transcript, emit: Emit) -> None: Runs classification first: refuses adversarial or off-domain queries, streams and signs a clarification question when disambiguation is - required, or admits the query into the discovery loop, synthesis, chart - configuration, and turn signing. A model call that Gemini blocks in any - of those phases ends the turn as `refused` with reason `safety_blocked`; - the phases let `GeminiBlockedError` propagate here rather than each - translating it. Every non-canceled turn ends in one `terminal` frame - (`complete`, `refused`, or `error`) and records its outcome and - classified pattern in telemetry. Only `asyncio.CancelledError` is - re-raised to the caller; cancellation emits no terminal frame because the - client has already disconnected, and it is recorded as `canceled` unless - the turn has already completed. Follow-up questions are emitted after the - terminal frame, and the background chart task is settled whenever the - turn ends so no orphaned task outlives the request. + required, or admits the query into the discovery loop, synthesis, and + turn signing. A model call that Gemini blocks in any of those phases + ends the turn as `refused` with reason `safety_blocked`; the phases let + `GeminiBlockedError` propagate here rather than each translating it. + Every non-canceled turn ends in one `terminal` frame (`complete`, + `refused`, or `error`) and records its outcome and classified pattern + in telemetry. Only `asyncio.CancelledError` is re-raised to the caller; + cancellation emits no terminal frame because the client has already + disconnected, and it is recorded as `canceled` unless the turn has + already completed. Follow-up questions are emitted after the terminal + frame. Args: transcript: The verified transcript holding the current query, @@ -871,7 +827,6 @@ async def run_turn(transcript: Transcript, emit: Emit) -> None: telemetry = TurnTelemetry() state: TerminalState | None = None error_type: ErrorType | RefusedReason | None = None - chart_task: asyncio.Task[JsonObject] | None = None try: if not get_settings().enable_chat: state, error_type = TerminalState.REFUSED, ErrorType.CHAT_DISABLED @@ -907,17 +862,13 @@ async def run_turn(transcript: Transcript, emit: Emit) -> None: # Processing runs inside discovery are also timed under their own phase. with telemetry.phase("discovery"): discovery = await run_discovery_phase(ctx) - chart_task = discovery.chart_task with telemetry.phase("synthesis"): - full_text = await run_synthesis_phase(ctx, discovery) - with telemetry.phase("chart_config"): - chart_config = await resolve_chart_config(ctx, discovery, full_text) - emit(ChartConfigFrame(chart_config=chart_config)) + synthesis = await run_synthesis_phase(ctx, discovery) with telemetry.phase("finalize"): window = await _sign_turn( ctx, - full_text, - lambda: build_state_slots(discovery.tables, chart_config), + synthesis.text, + lambda: build_state_slots(discovery.tables, synthesis.charts), ) state = TerminalState.COMPLETE emit( @@ -926,7 +877,7 @@ async def run_turn(transcript: Transcript, emit: Emit) -> None: ) ) with telemetry.phase("follow_ups"): - await run_followups(ctx, chart_config) + await run_followups(ctx, discovery, synthesis.charts) except TurnRefusedError as refusal: state, error_type = TerminalState.REFUSED, refusal.reason emit( @@ -968,26 +919,27 @@ async def run_turn(transcript: Transcript, emit: Emit) -> None: ) ) finally: - _settle_chart_task(chart_task) telemetry.finish(state or TerminalState.ERROR, error_type) async def run_clarification( ctx: TurnContext, question: str ) -> SignedWindow | None: - """Streams `question` as the turn response and signs it with empty slots. + """Answers the turn with the classifier's clarification question. - No `data_status`, `sources`, `chart_config`, or `follow_ups` events are - emitted because no observations were retrieved. + The question is emitted as the one text block of the answer and the + turn is signed with empty state slots, so that the user's reply arrives + with this exchange as context. No data, sources, charts, or follow-ups + are emitted, because nothing was retrieved. Args: ctx: The turn context holding the transcript, config, and telemetry. question: The clarification question to stream to the user. Returns: - The signed conversation window, or `None` if signing failed. + The signed window, or `None` when the turn could not be signed. """ - ctx.emit(TextFrame(text=question)) + ctx.emit(BlockContentFrame(block=TextBlock(index=0, markdown=question))) with ctx.telemetry.phase("finalize"): return await _sign_turn(ctx, question, ConversationStateSlots) @@ -1028,33 +980,13 @@ async def _sign_turn( return None -def _settle_chart_task(task: asyncio.Task[JsonObject] | None) -> None: - """Cancels the chart task if it is still running, or logs its error. - - This is the one place a chart task failure is logged, whether - `resolve_chart_config` contained it or synthesis failed before the task - was awaited. Retrieving the exception also prevents the "exception was - never retrieved" warning that asyncio logs when a failed task that no - one awaited is garbage collected. - - Args: - task: The background chart configuration task, or `None`. - """ - if task is None: - return - if not task.done(): - task.cancel() - elif not task.cancelled() and (error := task.exception()) is not None: - logger.warning("Chart config failed: %s", type(error).__name__) - - __all__ = [ "DiscoveryPhaseResult", "Emit", + "SynthesisPhaseResult", "TurnAbortedError", "TurnContext", - "render_analysis_context", - "resolve_chart_config", + "collect_follow_up_topics", "run_clarification", "run_discovery_phase", "run_followups", diff --git a/narratives/agent/src/narratives_agent/workflows/chat_pipeline_test.py b/narratives/agent/src/narratives_agent/workflows/chat_pipeline_test.py index d25e7086..20210b61 100644 --- a/narratives/agent/src/narratives_agent/workflows/chat_pipeline_test.py +++ b/narratives/agent/src/narratives_agent/workflows/chat_pipeline_test.py @@ -14,40 +14,42 @@ """Tests for `run_turn` in the chat pipeline. Verifies that: -1. A complete turn emits the answer, the chart config, a `complete` +1. A complete turn emits the answer as text and chart blocks, a `complete` terminal event carrying the idempotency key and the signed window, then the follow-ups, and is recorded as `complete`. The signed turn verifies on the next request, and an answer too long to sign is left unsigned. 2. The discovery phase runs the discovery loop, answers from its digest or from a default specification over the cache when the loop ends without one, streams the tool calls, the data status, and the digest's sources, - renders the digest and a table inventory for synthesis and the chart - planner, and records the scripts, fetch counts, and retries in - telemetry. + passes the digest and the table inventory to synthesis, and records the + scripts, fetch counts, and retries in telemetry. 3. A missing configuration, a failed discovery loop, an empty or slow tool - list, and a failed, broken, or empty synthesis end the turn with one - `error` terminal event carrying a fixed message and the machine-readable - `reason`, and are recorded as `error` with that reason. Configuration and - MCP failures end the turn before synthesis. When the MCP server is missing - required tools, the discovery phase logs a warning and continues the turn. + list, and a failed, broken, stalled, or empty synthesis end the turn with + one `error` terminal event carrying a fixed message and the + machine-readable `reason`, and are recorded as `error` with that reason. + Configuration and MCP failures end the turn before synthesis. When the MCP + server is missing required tools, the discovery phase logs a warning and + continues the turn. 4. An unexpected exception ends the turn as `internal_error`; a follow-up failure after the terminal event leaves the turn `complete`. -5. A chart config that misses its join timeout is canceled and the turn - completes without charts; a chart config that failed is retrieved. -6. Cancellation mid-turn cancels the background chart task and is recorded - as `canceled`; cancellation after the terminal event stays `complete`. +5. A chart block in the synthesis stream is resolved against the turn cache + and emitted in place while its markdown fence is stripped from the signed + text; a rejected chart block is dropped and, when no chart block is + emitted for a turn with data, a default chart is appended. +6. Cancellation mid-turn is recorded as `canceled`; cancellation after the + terminal event stays `complete`. 7. Synthesis receives the verified window before the current query, an `error` terminal event carries no signing fields, and a failure while signing leaves the delivered turn complete and unsigned. 8. The classify phase runs first behind a `classify` status frame: a refusal ends the turn in one `refused` terminal frame without an `error` - key and without any model call, a clarification is streamed as the - answer and signed with empty state slots, a classification carries its - pattern directive into the discovery loop and synthesis and its pattern - into telemetry, and a classifier failure proceeds as `exploratory` - without a directive. A model call that Gemini blocks in the classify - phase, the discovery loop, or synthesis ends the turn as `refused` with - reason `safety_blocked`. + key and without any model call, a clarification is emitted as the one + text block and signed with empty state slots, a classification carries + its pattern directive into the discovery loop and synthesis and its + pattern into telemetry, and a classifier failure proceeds as + `exploratory` without a directive. A model call that Gemini blocks in the + classify phase, the discovery loop, or synthesis ends the turn as + `refused` with reason `safety_blocked`. 9. The serialized frames of a complete turn match the recorded wire sequence, byte for byte. """ @@ -61,9 +63,14 @@ import pytest -from narratives_agent.analytics.digest import AnalysisDigest +from narratives_agent.analytics.digest import ( + AnalysisDigest, + RefusedScriptResult, +) from narratives_agent.analytics.provenance import Provenance, Source +from narratives_agent.analytics.ranking import rank_cohort from narratives_agent.analytics.tables import ObservationRow, ObservationTable +from narratives_agent.analytics.timeseries import analyze_timeseries from narratives_agent.config_model import AgentConfig, Prompts from narratives_agent.gemini.client import ( GeminiBlockedError, @@ -83,7 +90,7 @@ ) from narratives_agent.workflows import chat_pipeline from narratives_agent.workflows import transcript as transcript_module -from narratives_agent.workflows.chat_pipeline import render_analysis_context +from narratives_agent.workflows.charts import PlaceList, ResolvedLineChart from narratives_agent.workflows.classify import ( Admitted, Clarification, @@ -110,6 +117,7 @@ StepRefusalDiagnostic, execute_specification, ) +from narratives_agent.workflows.prompt_contracts import load_prompt_contract from narratives_agent.workflows.specification import ( AnalysisSpecification, UnavailableReport, @@ -131,9 +139,13 @@ _QUESTION = "What is the population of France?" _IDEMPOTENCY_KEY = "test-idempotency-key" _ANSWER = "About 68 million." -_CHARTS: JsonObject = { - "should_render": True, - "charts": [{"title": "Population"}], +# Default chart of the series summary over `t1`: a line chart over France. +_DEFAULT_CHART: JsonObject = { + "viz_type": "line", + "title": "Total Population", + "variable_dcids": ["Count_Person"], + "place_dcids": ["country/FRA"], + "facet_ids": {"Count_Person": "f1"}, } _FOLLOW_UPS = ["And Spain?"] _TRANSCRIPT = Transcript( @@ -153,6 +165,8 @@ "result": json.dumps({"table_id": "t1", "row_count": 2}), "status": "success", } +# Scopes of a complete turn: one for the fetched table, with the span of its +# dates, and one for the default chart, which carries no date. _SCOPES: JsonObject = { "scopes": [ { @@ -161,7 +175,14 @@ "child_place_type": None, "variables": {"Count_Person": "Total Population"}, "date_range": ["2022", "2023"], - } + }, + { + "places": {"country/FRA": "France"}, + "parent_place": None, + "child_place_type": None, + "variables": {"Count_Person": "Total Population"}, + "date_range": None, + }, ] } _CLASSIFICATION = QueryClassification.model_validate( @@ -262,12 +283,11 @@ class _Turn: synthesis_calls: int = 0 synthesis_messages: list[list[JsonObject]] = field(default_factory=list) synthesis_instructions: list[str] = field(default_factory=list) - chart_contexts: list[str] = field(default_factory=list) + answer: str = _ANSWER + follow_up_topics: list[list[str]] = field(default_factory=list) loop_classifications: list[QueryClassification | None] = field( default_factory=list ) - chart_canceled: asyncio.Event = field(default_factory=asyncio.Event) - chart_release: asyncio.Event = field(default_factory=asyncio.Event) synthesis_release: asyncio.Event = field(default_factory=asyncio.Event) follow_ups_release: asyncio.Event = field(default_factory=asyncio.Event) @@ -287,6 +307,31 @@ def list_content_payloads(self, key: str) -> list[JsonValue]: if event == "content" and key in data ] + def list_blocks(self) -> list[JsonObject]: + """Returns the `block` payloads streamed, in order.""" + return [ + dict(block) + for block in self.list_content_payloads("block") + if isinstance(block, dict) + ] + + def list_texts(self) -> list[str]: + """Returns the markdown of the text blocks streamed, in order.""" + return [ + str(block["markdown"]) + for block in self.list_blocks() + if block["type"] == "text" + ] + + def list_charts(self) -> list[JsonObject]: + """Returns the charts of the chart blocks streamed, in order.""" + return [ + dict(chart) + for block in self.list_blocks() + if block["type"] == "chart" + and isinstance(chart := block.get("chart"), dict) + ] + def get_synthesis_context(self) -> str: """Returns the text of the final user message synthesis received.""" (messages,) = self.synthesis_messages @@ -315,7 +360,6 @@ async def _run_accepted_loop(cache: TurnCache, **_: object) -> DiscoveryResult: def turn(monkeypatch: pytest.MonkeyPatch) -> _Turn: """Stubs every external call of a turn so each stage succeeds at once.""" stub = _Turn() - stub.chart_release.set() stub.synthesis_release.set() stub.follow_ups_release.set() @@ -353,7 +397,7 @@ async def run_discovery( async def stream_answer() -> AsyncGenerator[StreamChunk]: await stub.synthesis_release.wait() - yield StreamChunk("text", _ANSWER) + yield StreamChunk("text", stub.answer) async def start_synthesis( *, @@ -366,30 +410,15 @@ async def start_synthesis( stub.synthesis_instructions.append(system_instruction) return stream_answer() - async def build_chart_config( - results: str, question: str, config: AgentConfig, usage: TokenUsage - ) -> JsonObject: - stub.chart_contexts.append(results) - try: - await stub.chart_release.wait() - except asyncio.CancelledError: - stub.chart_canceled.set() - raise - return dict(_CHARTS) - - async def validate( - text: str, question: str, config: AgentConfig, usage: TokenUsage - ) -> bool: - return True - async def generate_follow_ups( question: str, topics: list[str], config: AgentConfig, usage: TokenUsage, ) -> list[str]: + stub.follow_up_topics.append(list(topics)) await stub.follow_ups_release.wait() - return list(_FOLLOW_UPS) + return list(_FOLLOW_UPS) if topics else [] monkeypatch.setattr(chat_pipeline, "TurnTelemetry", create_telemetry) monkeypatch.setattr( @@ -404,8 +433,6 @@ async def generate_follow_ups( monkeypatch.setattr(mcp_client, "async_get_tools", get_tools) monkeypatch.setattr(chat_pipeline, "execute_discovery_loop", run_discovery) monkeypatch.setattr(chat_pipeline, "async_gemini_stream", start_synthesis) - monkeypatch.setattr(chat_pipeline, "get_chart_config", build_chart_config) - monkeypatch.setattr(chat_pipeline, "validate_data_response", validate) monkeypatch.setattr( chat_pipeline, "generate_follow_up_questions", generate_follow_ups ) @@ -453,18 +480,22 @@ async def test_complete_turn_emits_the_answer_then_the_follow_ups( turn: _Turn, ) -> None: # Test: The event sequence of a turn in which every phase succeeds. - # Situation: Discovery, synthesis, chart config, and follow-ups all - # succeed. - # Expectation: The answer text and then the chart config are emitted as - # content, followed by one `complete` terminal event with the - # request's idempotency key, the signed turn, the state slots built - # from the fetched table, and a one-turn window, and last the - # follow-ups; the turn is recorded as `complete`. + # Situation: Discovery, synthesis, and follow-ups all succeed, and the + # answer carries no chart block of its own. + # Expectation: The answer is emitted as one text block and then the + # default chart as a chart block, followed by one `complete` terminal + # event with the request's idempotency key, the signed turn, the + # state slots built from the fetched table and the chart, and a + # one-turn window, and last the follow-ups; telemetry counts the one + # chart, and the turn is recorded as `complete`. await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) + assert turn.list_blocks() == [ + {"index": 0, "type": "text", "markdown": _ANSWER}, + {"index": 1, "type": "chart", "chart": _DEFAULT_CHART}, + ] content = [data for event, data in turn.events if event == "content"] - assert {"text": _ANSWER} in content - assert content[-1] == {"chart_config": _CHARTS} + assert content[-1] == {"block": turn.list_blocks()[1]} event, terminal = turn.events[-2] assert event == "terminal" signature = terminal["hmac"] @@ -490,6 +521,9 @@ async def test_complete_turn_emits_the_answer_then_the_follow_ups( {"follow_up_questions": _FOLLOW_UPS}, ) assert turn.list_event_names().count("terminal") == 1 + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == 1 + assert turn.telemetry.charts_rejected == 0 assert turn.finished == [(TerminalState.COMPLETE, None)] @@ -504,11 +538,10 @@ async def test_the_discovery_phase_streams_its_results( # Expectation: The phase opens with a `discovery` status, relays the # thought with phase `discovery` and the processing status, streams # the tool call, a data status with data and no truncation, and the - # table's source as the one numbered source; synthesis and the chart - # planner receive the digest and the table inventory, prefixed with - # the numbered source, and not the model's finding; and telemetry - # records the script, the fetch counts, no retries, and the discovery - # phase. + # table's source as the one numbered source; synthesis receives the + # digest and the table inventory, prefixed with the numbered source, + # and not the model's finding; and telemetry records the script, the + # fetch counts, no retries, and the discovery phase. async def run_loop( transcript: Transcript, config: AgentConfig, @@ -551,9 +584,6 @@ async def run_loop( assert "Step s1: analyze_timeseries on Total Population" in context assert "France publishes an annual census estimate." not in context assert "- t1: Total Population (Count_Person); unit: Person" in context - (chart_context,) = turn.chart_contexts - assert chart_context.startswith("Step s1: analyze_timeseries") - assert chart_context in context assert turn.telemetry is not None assert turn.telemetry.scripts == ["analyze_timeseries"] assert turn.telemetry.tables_fetched == 1 @@ -607,8 +637,9 @@ async def test_an_unavailable_report_answers_without_data( # table. # Expectation: The data status reports no data with the no-data # message, no sources are streamed, synthesis receives the note as - # missing data, no chart task runs, the turn completes with empty - # state slots, and telemetry records no scripts. + # missing data, no chart block is emitted and no default chart is + # appended, the follow-ups are grounded in the note alone, the turn + # completes with empty state slots, and telemetry records no scripts. report = UnavailableReport( searched_indicators=["Count_Person_Martian"], searched_places=["country/FRA"], @@ -627,14 +658,12 @@ async def test_an_unavailable_report_answers_without_data( } ] assert turn.list_content_payloads("sources") == [] - assert turn.list_content_payloads("chart_config") == [ - {"should_render": False} - ] + assert turn.list_charts() == [] context = turn.get_synthesis_context() - assert "Missing data:" in context + assert "MISSING DATA:" in context assert "- " + report.format_note() in context assert "NUMBERED SOURCES" not in context - assert turn.chart_contexts == [] + assert turn.follow_up_topics == [[report.format_note()]] _, terminal = turn.events[-2] assert terminal["state_slots"] == {"scopes": []} assert turn.telemetry is not None @@ -681,8 +710,9 @@ async def test_a_truncated_loop_without_tables_reports_no_data( # Test: The fallback when the loop ends at a bound with nothing cached. # Situation: The loop returns truncated with an empty cache. # Expectation: The data status reports no data with the truncation - # message, synthesis receives the no-observations note, and the turn - # completes without charts. + # message, synthesis receives the no-observations note, no follow-up + # topics or `follow_ups` frame are emitted, and the turn completes + # without a chart block. _install_loop(monkeypatch, DiscoveryResult(truncated=True)) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) @@ -695,9 +725,9 @@ async def test_a_truncated_loop_without_tables_reports_no_data( } ] assert chat_pipeline._NO_OBSERVATIONS_NOTE in turn.get_synthesis_context() - assert turn.list_content_payloads("chart_config") == [ - {"should_render": False} - ] + assert turn.list_charts() == [] + assert turn.follow_up_topics == [[]] + assert "follow_ups" not in turn.list_event_names() assert turn.finished == [(TerminalState.COMPLETE, None)] @@ -742,109 +772,6 @@ async def test_exhausted_retries_fall_back_and_note_the_failure( assert turn.finished == [(TerminalState.COMPLETE, None)] -def test_render_analysis_context_lists_the_tables() -> None: - # Test: The adapter that feeds synthesis and the chart planner. - # Situation: A digest with no results, a specification with one finding - # whose second step derives a table from the first two, and three - # tables: a cohort of counties with more places than the inventory - # cap and a scaling factor, a single-place table without a unit, and - # the derived table, whose variable DCID the specification chose. - # Expectation: The text holds no trace of the finding, which is the - # model's note rather than a computed figure, then one inventory line - # per table naming its id, variable, unit with any scaling factor, - # scope, date span, and places, with the cohort's places capped and - # counted; the derived table is listed by name only, marked as derived - # from its inputs and not a Data Commons variable, so the chart - # planner cannot pick its DCID up. - places = tuple(f"geoId/{n:05d}" for n in range(30)) - cohort = CachedTable( - table_id="t1", - table=ObservationTable( - variable_dcid="Count_Person", - variable_name="Total Population", - unit="Person", - scaling_factor=1000.0, - provenance=Provenance(), - rows=tuple( - ObservationRow(place, f"County {place[-5:]}", "2023", 1.0) - for place in places - ), - requested_places=places, - place_names={place: f"County {place[-5:]}" for place in places}, - parent_place_dcid="geoId/06", - child_place_type="County", - ), - ) - single_table = ObservationTable( - variable_dcid="Median_Income", - variable_name="Median Income", - unit=None, - scaling_factor=None, - provenance=Provenance(), - rows=( - ObservationRow("geoId/06", "California", "2020", 1.0), - ObservationRow("geoId/06", "California", "2023", 2.0), - ), - requested_places=("geoId/06",), - place_names={"geoId/06": "California"}, - ) - single = CachedTable(table_id="t2", table=single_table) - derived = CachedTable( - table_id="s2", - table=replace( - single_table, - variable_dcid="Income_Per_Capita_Invented", - variable_name="Income per capita", - unit="USD", - ), - ) - specification = AnalysisSpecification.model_validate( - { - "steps": [ - {"step_id": "s1", "script": "rank_cohort", "table": "t1"}, - { - "step_id": "s2", - "script": "compute_derived_metric", - "left": "t2", - "right": "t1", - "derived_variable_dcid": "Income_Per_Capita_Invented", - }, - ], - "findings": ["Counties vary widely."], - } - ) - - text = render_analysis_context( - AnalysisDigest(results=(), sources=()), - specification, - (cohort, single, derived), - ) - - assert "Counties vary widely" not in text - assert "findings" not in text - cohort_line, single_line, derived_line = text.split("Tables used:\n")[ - 1 - ].splitlines() - assert cohort_line.startswith( - "- t1: Total Population (Count_Person); unit: Person (scaling factor " - "1000); every County in geoId/06; dates 2023 to 2023; places with " - "data: County 00000 (geoId/00000)" - ) - assert cohort_line.endswith( - f", and {30 - chat_pipeline._INVENTORY_PLACE_CAP} other places" - ) - assert single_line == ( - "- t2: Median Income (Median_Income); unit: unknown; 1 place; " - "dates 2020 to 2023; places with data: California (geoId/06)" - ) - assert derived_line == ( - "- s2 (derived from t2 and t1; not a Data Commons variable, do not " - "chart it): Income per capita; unit: USD; 1 place; dates 2020 to " - "2023; places with data: California (geoId/06)" - ) - assert "Income_Per_Capita_Invented" not in text - - def _make_state_cohort( table_id: str, variable_dcid: str, places: str ) -> CachedTable: @@ -1096,8 +1023,7 @@ async def test_a_signing_failure_completes_the_turn_unsigned( ) -> None: # Test: A signing failure after the answer is delivered leaves the turn # complete and unsigned. - # Situation: Signing raises once the answer and chart config have been - # emitted. + # Situation: Signing raises once the answer's blocks have been emitted. # Expectation: The turn still ends `complete`, without signing fields, # and is recorded as `complete`. async def fail_finalize(*args: object, **_: object) -> None: @@ -1198,126 +1124,270 @@ async def fail_synthesis(**_: object) -> AsyncGenerator[StreamChunk]: @pytest.mark.asyncio -async def test_slow_chart_config_is_canceled_and_charts_are_dropped( - monkeypatch: pytest.MonkeyPatch, turn: _Turn +async def test_a_chart_block_is_resolved_and_streamed_in_place( + turn: _Turn, ) -> None: - # Test: The chart-config join timeout. - # Situation: The join timeout is shortened to 10 ms and chart - # configuration never finishes. - # Expectation: The chart task is canceled, and the turn completes with - # `should_render` set to `False`. - monkeypatch.setattr( - chat_pipeline, "_CHART_CONFIG_JOIN_TIMEOUT_SECONDS", 0.01 + # Test: A chart block the model writes. + # Situation: The answer holds a paragraph, a chart block over `t1` + # naming the one place the table has, and a closing paragraph. + # Expectation: The blocks are streamed in order with the chart resolved + # to the wire shape, the signed answer is the two paragraphs joined by + # a blank line with the block stripped, no default chart is added, + # the chart contributes a state slot, the follow-ups are grounded in + # the variable and the chart's title, and telemetry counts one chart. + turn.answer = ( + "France is growing.\n\n" + '```chart\n{"kind": "line", "tables": ["t1"], ' + '"places": ["country/FRA"], "title": "Population of France"}\n```\n\n' + "It reached 68 million in 2023." ) - turn.chart_release.clear() await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - content = [data for event, data in turn.events if event == "content"] - assert content[-1] == {"chart_config": {"should_render": False}} - assert turn.chart_canceled.is_set() + chart = {**_DEFAULT_CHART, "title": "Population of France"} + assert turn.list_blocks() == [ + {"index": 0, "type": "text", "markdown": "France is growing."}, + {"index": 1, "type": "chart", "chart": chart}, + { + "index": 2, + "type": "text", + "markdown": "It reached 68 million in 2023.", + }, + ] + _, terminal = turn.events[-2] + assert terminal["state_slots"] == _SCOPES + stored = Turn.model_validate( + { + "turn_index": 0, + "idempotency_key": _IDEMPOTENCY_KEY, + "user_query": _QUESTION, + "model_response": ( + "France is growing.\n\nIt reached 68 million in 2023." + ), + "state_slots": terminal["state_slots"], + "hmac": terminal["hmac"], + } + ) + assert load_transcript(_Request("And Spain?", "key-2", [stored])).turns == [ + stored + ] + assert turn.follow_up_topics == [ + ["Total Population in France", "Population of France"] + ] + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == 1 + assert turn.telemetry.charts_rejected == 0 + + +@pytest.mark.asyncio +async def test_a_rejected_chart_block_gives_way_to_the_default_chart( + turn: _Turn, +) -> None: + # Test: The default-chart floor after a rejected block. + # Situation: The answer's one chart block names a table the turn never + # fetched. + # Expectation: The block is dropped without a trace in the text, the + # default chart is appended after the prose, and telemetry counts one + # rejected and one emitted chart. + turn.answer = ( + "France is growing.\n\n" + '```chart\n{"kind": "line", "tables": ["t9"], "title": "Ghost"}\n```' + ) + + await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) + + assert turn.list_texts() == ["France is growing."] + assert turn.list_charts() == [_DEFAULT_CHART] + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == 1 + assert turn.telemetry.charts_rejected == 1 assert turn.finished == [(TerminalState.COMPLETE, None)] @pytest.mark.asyncio -async def test_a_failed_chart_task_leaves_the_turn_complete( +async def test_a_chart_block_naming_an_unused_cached_table_is_rejected( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: - # Test: Containment of a chart-configuration failure. - # Situation: Chart configuration raises after the answer has streamed. - # Expectation: The turn completes with `should_render` set to `False` - # rather than ending as an internal error. - async def fail_chart_config(*args: object) -> JsonObject: - raise RuntimeError("chart config failed") + # Test: Rejecting a chart block that names a cached table outside the + # table inventory. + # Situation: The loop caches an unused `t1` for Germany and `t2` for + # France, accepts a series specification over `t2` alone, and + # synthesis emits a chart block naming `t1`. + # Expectation: The `t1` chart block is dropped, the default chart for + # `t2` is appended in its place, and telemetry counts one rejected + # and one emitted chart. + unused = _make_table("t1", places=("country/DEU",)) + unused = replace(unused, table=replace(unused.table, facet_id="f_unused")) + accepted = _make_table("t2", places=("country/FRA",)) + spec = AnalysisSpecification.model_validate( + { + "steps": [ + {"step_id": "s1", "script": "analyze_timeseries", "table": "t2"} + ] + } + ) - monkeypatch.setattr(chat_pipeline, "get_chart_config", fail_chart_config) + async def run_loop( + transcript: Transcript, + config: AgentConfig, + tools: Sequence[mcp_client.McpTool], + telemetry: TurnTelemetry, + cache: TurnCache, + **_: object, + ) -> DiscoveryResult: + cache.store_fetched(unused) + cache.store_fetched(accepted) + digest = await execute_specification(spec, cache) + assert isinstance(digest, AnalysisDigest) + return DiscoveryResult(specification=spec, digest=digest) + + monkeypatch.setattr(chat_pipeline, "execute_discovery_loop", run_loop) + turn.answer = ( + "France is growing.\n\n" + '```chart\n{"kind": "line", "tables": ["t1"], ' + '"places": ["country/DEU"], "title": "Germany"}\n```' + ) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - content = [data for event, data in turn.events if event == "content"] - assert {"text": _ANSWER} in content - assert content[-1] == {"chart_config": {"should_render": False}} - terminals = [data for event, data in turn.events if event == "terminal"] - assert [terminal["state"] for terminal in terminals] == [ - TerminalState.COMPLETE - ] + assert turn.list_texts() == ["France is growing."] + assert turn.list_charts() == [_DEFAULT_CHART] + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == 1 + assert turn.telemetry.charts_rejected == 1 assert turn.finished == [(TerminalState.COMPLETE, None)] @pytest.mark.asyncio -async def test_a_failed_chart_validation_hides_the_charts( +async def test_a_chart_block_on_a_no_data_turn_is_rejected( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: - # Test: Containment of a chart-validation failure. - # Situation: The data-validation check raises after the answer has - # streamed, and chart configuration succeeds. - # Expectation: The turn completes with the charts hidden rather than - # ending as an internal error. - async def fail_validate(*args: object) -> bool: - raise RuntimeError("validation failed") + # Test: Rejecting a chart block when the turn has no computed data. + # Situation: A top-N turn ends at a bound with one state cohort table + # requesting four states with two dates for only one state, so its + # default `rank_cohort` step refuses (`has_data=False`) while `t1` + # stays cached, and synthesis emits a `line` chart block over `t1`. + # Expectation: `data_status` reports `has_data=False`, no chart block is + # emitted, and telemetry counts zero emitted and one rejected chart. + ranking = QueryClassification.model_validate( + { + "serviceable": True, + "pattern": "top_n_ranking", + "geographic_shape": "child_places", + "requires_computation": True, + } + ) - monkeypatch.setattr(chat_pipeline, "validate_data_response", fail_validate) + async def classify(*args: object, **_: object) -> Admitted: + return Admitted(ranking) + + sparse = _make_state_cohort("t1", "Median_Income", "01 02 03 04") + sparse = replace( + sparse, + table=replace( + sparse.table, + rows=( + ObservationRow("geoId/01", "geoId/01", "2022", 1.0), + ObservationRow("geoId/01", "geoId/01", "2023", 2.0), + ), + ), + ) + monkeypatch.setattr(chat_pipeline, "classify_query", classify) + _install_loop(monkeypatch, DiscoveryResult(truncated=True), sparse) + turn.answer = ( + "State income data is sparse.\n\n" + '```chart\n{"kind": "line", "tables": ["t1"], ' + '"places": ["geoId/01"], "title": "Sparse state"}\n```' + ) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - content = [data for event, data in turn.events if event == "content"] - assert content[-1] == {"chart_config": {**_CHARTS, "hide_charts": True}} - terminals = [data for event, data in turn.events if event == "terminal"] - assert [terminal["state"] for terminal in terminals] == [ - TerminalState.COMPLETE + assert turn.list_content_payloads("data_status") == [ + {"has_data": False, "truncated": True, "message": TRUNCATED_MESSAGE} ] + assert turn.list_texts() == ["State income data is sparse."] + assert turn.list_charts() == [] + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == 0 + assert turn.telemetry.charts_rejected == 1 assert turn.finished == [(TerminalState.COMPLETE, None)] @pytest.mark.asyncio -async def test_malformed_chart_entries_do_not_fail_the_turn( +async def test_a_stalled_synthesis_stream_fails_the_turn( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: - # Test: Tolerance of malformed entries in a chart configuration. - # Situation: Chart configuration returns an object whose `charts` list - # mixes a well-formed chart with a null entry, a string, and a chart - # whose title is not a string. - # Expectation: The turn completes, and the follow-ups are grounded in - # the one well-formed title. - topics: list[list[str]] = [] - malformed: JsonObject = { - "should_render": True, - "charts": [None, "Population", {"title": 7}, {"title": "GDP"}], - } + # Test: The per-chunk idle timeout of the synthesis stream. + # Situation: The timeout is shortened to 10 ms, and the stream yields + # one paragraph and then never yields again. + # Expectation: The first paragraph is emitted as a text block, and the + # turn ends in an `error` terminal event with reason + # `synthesis_stream_error` rather than waiting on the stream. + async def stream_stalled() -> AsyncGenerator[StreamChunk]: + yield StreamChunk("text", "First paragraph.\n\n") + await asyncio.Event().wait() - async def build_chart_config(*args: object) -> JsonObject: - return malformed + async def start_stalled_synthesis( + **_: object, + ) -> AsyncGenerator[StreamChunk]: + return stream_stalled() - async def generate_follow_ups( - question: str, - chart_topics: list[str], - config: AgentConfig, - usage: TokenUsage, - ) -> list[str]: - topics.append(chart_topics) - return list(_FOLLOW_UPS) + monkeypatch.setattr(chat_pipeline, "SYNTHESIS_CHUNK_TIMEOUT_SECONDS", 0.01) + monkeypatch.setattr( + chat_pipeline, "async_gemini_stream", start_stalled_synthesis + ) + + await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - monkeypatch.setattr(chat_pipeline, "get_chart_config", build_chart_config) + assert turn.list_texts() == ["First paragraph."] + terminal = _get_only_terminal(turn) + assert terminal["reason"] == ErrorType.SYNTHESIS_STREAM_ERROR + assert terminal["error"] == SYNTHESIS_FAILED + assert turn.finished == [ + (TerminalState.ERROR, ErrorType.SYNTHESIS_STREAM_ERROR) + ] + + +@pytest.mark.asyncio +async def test_a_synthesis_stream_stalled_on_open_fails_the_turn( + monkeypatch: pytest.MonkeyPatch, turn: _Turn +) -> None: + # Test: The first-chunk timeout when opening the synthesis stream. + # Situation: `SYNTHESIS_CHUNK_TIMEOUT_SECONDS` is shortened to 10 ms, + # and `async_gemini_stream` stalls awaiting the first chunk. + # Expectation: No block is emitted, and the turn ends in an `error` + # terminal event with reason `synthesis_stream_error` and + # `SYNTHESIS_FAILED` rather than waiting for the HTTP timeout. + async def hang_on_first_chunk( + **_: object, + ) -> AsyncGenerator[StreamChunk]: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + monkeypatch.setattr(chat_pipeline, "SYNTHESIS_CHUNK_TIMEOUT_SECONDS", 0.01) monkeypatch.setattr( - chat_pipeline, "generate_follow_up_questions", generate_follow_ups + chat_pipeline, "async_gemini_stream", hang_on_first_chunk ) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - assert topics == [["GDP"]] - assert turn.finished == [(TerminalState.COMPLETE, None)] + assert turn.list_blocks() == [] + terminal = _get_only_terminal(turn) + assert terminal["reason"] == ErrorType.SYNTHESIS_STREAM_ERROR + assert terminal["error"] == SYNTHESIS_FAILED + assert turn.finished == [ + (TerminalState.ERROR, ErrorType.SYNTHESIS_STREAM_ERROR) + ] @pytest.mark.asyncio -async def test_cancel_mid_turn_cancels_the_chart_task(turn: _Turn) -> None: +async def test_cancel_mid_turn_is_recorded_as_canceled(turn: _Turn) -> None: # Test: Cancellation while synthesis is streaming. - # Situation: Synthesis and chart configuration are both held open, and - # the turn is canceled once synthesis has started. + # Situation: Synthesis is held open, and the turn is canceled once it + # has started. # Expectation: `CancelledError` is raised to the caller without a - # terminal event, the background chart task is canceled with the turn, - # and the turn is recorded as `canceled`. + # terminal event, and the turn is recorded as `canceled`. turn.synthesis_release.clear() - turn.chart_release.clear() task = await _start_and_wait_for(turn, "status") async with asyncio.timeout(1): while turn.synthesis_calls == 0: @@ -1327,8 +1397,6 @@ async def test_cancel_mid_turn_cancels_the_chart_task(turn: _Turn) -> None: with pytest.raises(asyncio.CancelledError): await task - async with asyncio.timeout(1): - await turn.chart_canceled.wait() assert "terminal" not in turn.list_event_names() assert turn.finished == [(TerminalState.CANCELED, None)] @@ -1466,18 +1534,32 @@ async def return_partial_tools() -> list[mcp_client.McpTool]: assert turn.finished == [(TerminalState.COMPLETE, None)] +@pytest.mark.parametrize( + ("raw_stream", "expected_charts_emitted"), + [ + (" \n", 0), + ( + '```chart\n{"kind": "line", "tables": ["t1"], ' + '"title": "Total Population"}\n```\n', + 1, + ), + ], +) @pytest.mark.asyncio async def test_an_empty_synthesis_fails_the_turn( - monkeypatch: pytest.MonkeyPatch, turn: _Turn + monkeypatch: pytest.MonkeyPatch, + turn: _Turn, + raw_stream: str, + expected_charts_emitted: int, ) -> None: - # Test: Fail-closed handling of empty synthesis output. - # Situation: The synthesis stream finishes after yielding only - # whitespace. + # Test: Fail-closed handling of synthesis output without prose text. + # Situation: The synthesis stream finishes after yielding either only + # whitespace or only a chart block with no accompanying prose. # Expectation: The turn ends in an `error` terminal event with reason - # `synthesis_empty` and the fixed synthesis message, and no chart - # config is emitted. + # `synthesis_empty` and the fixed synthesis message, no text block or + # default chart is appended, and telemetry records `charts_emitted`. async def stream_blank() -> AsyncGenerator[StreamChunk]: - yield StreamChunk("text", " \n") + yield StreamChunk("text", raw_stream) async def start_blank_synthesis( **_: object, @@ -1493,7 +1575,10 @@ async def start_blank_synthesis( terminal = _get_only_terminal(turn) assert terminal["reason"] == ErrorType.SYNTHESIS_EMPTY assert terminal["error"] == SYNTHESIS_FAILED - assert not any("chart_config" in data for _, data in turn.events) + assert turn.list_texts() == [] + assert len(turn.list_charts()) == expected_charts_emitted + assert turn.telemetry is not None + assert turn.telemetry.charts_emitted == expected_charts_emitted assert turn.finished == [(TerminalState.ERROR, ErrorType.SYNTHESIS_EMPTY)] @@ -1502,15 +1587,16 @@ async def test_a_broken_synthesis_stream_fails_the_turn( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: # Test: Fail-closed handling of a synthesis stream that fails mid-stream. - # Situation: The synthesis stream yields part of the answer and then - # raises `GeminiStreamError` with upstream error text. - # Expectation: The partial text is emitted, and the turn ends in an - # `error` terminal event with reason `synthesis_stream_error` and the - # fixed message rather than the upstream text. + # Situation: The synthesis stream yields one paragraph and part of the + # next, then raises `GeminiStreamError` with upstream error text. + # Expectation: The completed paragraph is emitted as a text block, the + # partial paragraph is not, and the turn ends in an `error` terminal + # event with reason `synthesis_stream_error` and the fixed message + # rather than the upstream text. upstream = "HTTP 503: upstream detail" async def stream_broken() -> AsyncGenerator[StreamChunk]: - yield StreamChunk("text", "About ") + yield StreamChunk("text", "First paragraph.\n\nAbout ") raise GeminiStreamError(upstream) async def start_broken_synthesis( @@ -1524,7 +1610,7 @@ async def start_broken_synthesis( await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - assert ("content", {"text": "About "}) in turn.events + assert turn.list_texts() == ["First paragraph."] terminal = _get_only_terminal(turn) assert terminal["reason"] == ErrorType.SYNTHESIS_STREAM_ERROR assert terminal["error"] == SYNTHESIS_FAILED @@ -1538,7 +1624,7 @@ async def test_an_unexpected_exception_ends_the_turn_as_internal_error( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: # Test: Catch-all handling for unexpected exceptions inside a phase. - # Situation: `render_analysis_context` raises an unexpected + # Situation: Rendering the synthesis message raises an unexpected # `RuntimeError` containing internal details. # Expectation: The turn ends in an `error` terminal event with reason # `internal_error` and the fixed internal-error message. @@ -1547,7 +1633,7 @@ async def test_an_unexpected_exception_ends_the_turn_as_internal_error( def fail_render(*args: object) -> str: raise RuntimeError(detail) - monkeypatch.setattr(chat_pipeline, "render_analysis_context", fail_render) + monkeypatch.setattr(chat_pipeline, "render_synthesis_message", fail_render) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) @@ -1598,42 +1684,101 @@ async def fail_follow_ups( assert record.exc_info is None -@pytest.mark.asyncio -async def test_a_failed_chart_task_is_reaped_when_the_turn_fails( - monkeypatch: pytest.MonkeyPatch, - turn: _Turn, - caplog: pytest.LogCaptureFixture, -) -> None: - # Test: Cleanup of a failed chart task when a later phase fails. - # Situation: `get_chart_config` raises a `ValueError`, and synthesis - # fails after the chart task has finished. - # Expectation: The turn ends in an `error` terminal event, and the chart - # task's exception is retrieved and logged by type. - chart_failed = asyncio.Event() - - async def fail_chart_config( - results: str, question: str, config: AgentConfig, usage: TokenUsage - ) -> JsonObject: - chart_failed.set() - raise ValueError("chart detail") - - async def fail_late_synthesis( - **_: object, - ) -> AsyncGenerator[StreamChunk]: - await chart_failed.wait() - await asyncio.sleep(0) - raise GeminiError("upstream") +def test_collect_follow_up_topics_leaves_processing_diagnostics_out() -> None: + # Test: The missing-data notes that become follow-up topics. + # Situation: A digest with no results carries a data-gap note, a + # step-failure diagnostic, a default-analysis diagnostic, and the + # generic no-observations fallback note, and no chart was shown. + # Expectation: The data-gap note is the only topic. + digest = AnalysisDigest( + results=(), + sources=(), + missing_notes=( + "No county breakdown of median income exists.", + "An analysis step failed: the cohort has one place.", + "The default analysis also failed: the table has no dates.", + chat_pipeline._NO_OBSERVATIONS_NOTE, + ), + ) + discovery = chat_pipeline.DiscoveryPhaseResult( + digest=digest, + spec=None, + has_data=False, + tool_calls=[], + sources=[], + tables=(), + ) - monkeypatch.setattr(chat_pipeline, "get_chart_config", fail_chart_config) - monkeypatch.setattr( - chat_pipeline, "async_gemini_stream", fail_late_synthesis + topics = chat_pipeline.collect_follow_up_topics(discovery, []) + + assert topics == ["No county breakdown of median income exists."] + + +def test_collect_follow_up_topics_formats_scopes_and_deduplicates() -> None: + # Test: Scope formatting, refused-step filtering, and deduplication in + # `collect_follow_up_topics`. + # Situation: A digest holds an `OkScriptResult` over a four-place table, + # an `OkScriptResult` over a state cohort table, and a + # `RefusedScriptResult`, alongside two resolved charts and two data-gap + # notes that repeat earlier topics. + # Expectation: Named places are capped at three, cohort tables use + # `" by State"`, refused steps are skipped, and duplicate topics across + # results, chart titles, and missing notes are dropped in order. + multi_place = _make_table( + "t1", + places=( + "country/FRA", + "country/ESP", + "country/ITA", + "country/DEU", + ), ) + cohort = _make_state_cohort("t2", "Median_Income", "01 02 03 04") + refused = RefusedScriptResult( + step_id="s3", + script="rank_cohort", + variables=("Count_Unemployment",), + variable_names={"Count_Unemployment": "Unemployment"}, + provenance=Provenance(), + refusal_reason="no_observations", + ) + digest = AnalysisDigest.from_results( + ( + analyze_timeseries(multi_place.table, step_id="s1"), + rank_cohort(cohort.table, step_id="s2"), + refused, + ), + missing_notes=("Population Trend", "No county breakdown exists."), + ) + discovery = chat_pipeline.DiscoveryPhaseResult( + digest=digest, + spec=None, + has_data=True, + tool_calls=[], + sources=[], + tables=(multi_place, cohort), + ) + charts = [ + ResolvedLineChart( + title="Median_Income by State", + variable_dcids=("Median_Income",), + scope=PlaceList(("geoId/01",)), + ), + ResolvedLineChart( + title="Population Trend", + variable_dcids=("Count_Person",), + scope=PlaceList(("country/FRA",)), + ), + ] - with caplog.at_level(logging.WARNING, logger=chat_pipeline.__name__): - await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) + topics = chat_pipeline.collect_follow_up_topics(discovery, charts) - assert _get_only_terminal(turn)["reason"] == ErrorType.SYNTHESIS_ERROR - assert "Chart config failed: ValueError" in caplog.text + assert topics == [ + "Total Population in Deu, Esp, France", + "Median_Income by State", + "Population Trend", + "No county breakdown exists.", + ] @pytest.mark.asyncio @@ -1645,9 +1790,10 @@ async def test_the_classify_phase_opens_the_stream_and_guides_the_turn( # and every later phase succeeds. # Expectation: The stream opens with a `classify` status frame followed # by the `discovery` one, the classification reaches the discovery - # loop, the synthesis system instruction ends with the pattern's - # directive, and telemetry records the pattern and the classify phase - # duration. + # loop, the synthesis system instruction opens with the prompt, + # carries the synthesis contract, and ends with the synthesis half of + # the pattern's directive, and telemetry records the pattern and the + # classify phase duration. await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) assert turn.events[0] == ( @@ -1660,7 +1806,9 @@ async def test_the_classify_phase_opens_the_stream_and_guides_the_turn( assert statuses[:2] == ["classify", "discovery"] assert turn.loop_classifications == [_CLASSIFICATION] (instruction,) = turn.synthesis_instructions - assert instruction.startswith("Write prose.") + assert instruction.startswith( + f"Write prose.\n\n{load_prompt_contract('synthesis')}" + ) assert instruction.endswith( build_pattern_directive(_CLASSIFICATION, "synthesis") ) @@ -1725,7 +1873,7 @@ async def test_a_clarification_is_streamed_and_signed_with_empty_slots( ) -> None: # Test: The clarification outcome of the classify phase. # Situation: The classifier asks which Georgia the user means. - # Expectation: The question is streamed as the answer text, the turn + # Expectation: The question is emitted as the one text block, the turn # ends in a signed `complete` frame with empty state slots and a # one-turn window, no data, sources, charts, or follow-ups are # emitted, no discovery loop or synthesis runs, the pattern is @@ -1745,7 +1893,13 @@ async def clarify_classify(*args: object, **_: object) -> Clarification: await chat_pipeline.run_turn(transcript, turn.emit) assert turn.list_event_names() == ["status", "content", "terminal"] - assert turn.events[1][1] == {"text": _CLARIFICATION_QUESTION} + assert turn.events[1][1] == { + "block": { + "index": 0, + "type": "text", + "markdown": _CLARIFICATION_QUESTION, + } + } terminal = turn.events[2][1] assert terminal["state"] == TerminalState.COMPLETE assert terminal["turn_index"] == 0 @@ -1783,8 +1937,9 @@ async def test_a_classifier_failure_proceeds_as_exploratory( # Test: Fall-through when the classifier failed. # Situation: The classify phase admits the query without a verdict. # Expectation: The turn completes as usual, the discovery loop receives - # no classification, the synthesis system instruction carries no - # directive, and telemetry records the pattern as `exploratory`. + # no classification, the synthesis system instruction is the prompt + # and the synthesis contract with no directive, and telemetry records + # the pattern as `exploratory`. async def return_unclassified(*args: object, **_: object) -> Admitted: return Admitted(None) @@ -1793,7 +1948,9 @@ async def return_unclassified(*args: object, **_: object) -> Admitted: await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) assert turn.loop_classifications == [None] - assert turn.synthesis_instructions == ["Write prose."] + assert turn.synthesis_instructions == [ + f"Write prose.\n\n{load_prompt_contract('synthesis')}" + ] assert turn.telemetry is not None assert turn.telemetry.pattern == "exploratory" assert turn.finished == [(TerminalState.COMPLETE, None)] @@ -1836,7 +1993,7 @@ async def test_a_blocked_synthesis_prompt_refuses_the_turn( ) -> None: # Test: Mapping of a blocked prompt in synthesis. # Situation: The synthesis request is blocked before any text streams. - # Expectation: No answer text is emitted, the turn ends in a `refused` + # Expectation: No block is emitted, the turn ends in a `refused` # terminal frame with reason `safety_blocked`, and it is recorded as # `refused` with `safety_blocked`. async def raise_blocked_synthesis( @@ -1850,7 +2007,7 @@ async def raise_blocked_synthesis( await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - assert not any("text" in data for event, data in turn.events) + assert turn.list_blocks() == [] terminal = _get_only_terminal(turn) assert terminal["state"] == TerminalState.REFUSED assert terminal["reason"] == RefusedReason.SAFETY_BLOCKED @@ -1860,33 +2017,47 @@ async def raise_blocked_synthesis( @pytest.mark.asyncio -async def test_a_synthesis_stream_blocked_after_thoughts_refuses_the_turn( - monkeypatch: pytest.MonkeyPatch, turn: _Turn +@pytest.mark.parametrize( + "chunk", + [ + pytest.param( + StreamChunk("thought", "Reviewing data."), + id="after_thought", + ), + pytest.param( + StreamChunk("text", "Unfinished first paragraph"), + id="after_buffered_text_before_block", + ), + ], +) +async def test_a_synthesis_stream_blocked_before_any_block_refuses_the_turn( + monkeypatch: pytest.MonkeyPatch, turn: _Turn, chunk: StreamChunk ) -> None: - # Test: Mid-stream safety block in `run_synthesis_phase` before answer text. - # Situation: The synthesis stream yields a thought chunk and then raises + # Test: Mid-stream safety block in `run_synthesis_phase` before any answer + # block has been emitted. + # Situation: The synthesis stream yields either a thought chunk or partial + # paragraph text that has not yet completed a block, and then raises # `GeminiBlockedError("SAFETY")`. - # Expectation: No answer text is emitted, the turn ends in a `refused` - # terminal frame with `reason="safety_blocked"`, and telemetry records - # `("refused", "safety_blocked")`. - async def stream_blocked_after_thought() -> AsyncGenerator[StreamChunk]: - yield StreamChunk("thought", "Reviewing data.") + # Expectation: No `content {block}` frame is emitted, the turn ends in a + # `refused` terminal frame with `reason="safety_blocked"`, and telemetry + # records `("refused", "safety_blocked")`. + async def stream_blocked() -> AsyncGenerator[StreamChunk]: + yield chunk raise GeminiBlockedError("SAFETY") async def start_synthesis(**_: object) -> AsyncGenerator[StreamChunk]: - return stream_blocked_after_thought() + return stream_blocked() monkeypatch.setattr(chat_pipeline, "async_gemini_stream", start_synthesis) await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - assert ( - "thought", - {"thought": "Reviewing data.", "phase": "synthesis"}, - ) in turn.events - assert not any( - event == "content" and "text" in data for event, data in turn.events - ) + if chunk.kind == "thought": + assert ( + "thought", + {"thought": "Reviewing data.", "phase": "synthesis"}, + ) in turn.events + assert turn.list_blocks() == [] terminal = _get_only_terminal(turn) assert terminal["state"] == TerminalState.REFUSED assert terminal["reason"] == RefusedReason.SAFETY_BLOCKED @@ -1901,13 +2072,14 @@ async def test_a_synthesis_stream_blocked_after_text_fails_as_stream_error( monkeypatch: pytest.MonkeyPatch, turn: _Turn ) -> None: # Test: Mid-stream safety block in `run_synthesis_phase` after answer text. - # Situation: The synthesis stream yields answer text and then raises - # `GeminiBlockedError("SAFETY")`. - # Expectation: The partial text is emitted, and the turn ends in an `error` - # terminal frame with `reason="synthesis_stream_error"` and - # `SYNTHESIS_FAILED` rather than `refused`. + # Situation: The synthesis stream yields one paragraph and part of the next, + # then raises `GeminiBlockedError("SAFETY")`. + # Expectation: The completed paragraph is emitted as a text block, and the + # turn ends in an `error` terminal frame with + # `reason="synthesis_stream_error"` and `SYNTHESIS_FAILED` rather than + # `refused`. async def stream_blocked_after_text() -> AsyncGenerator[StreamChunk]: - yield StreamChunk("text", "About ") + yield StreamChunk("text", "First paragraph.\n\nAbout ") raise GeminiBlockedError("SAFETY") async def start_synthesis(**_: object) -> AsyncGenerator[StreamChunk]: @@ -1917,7 +2089,7 @@ async def start_synthesis(**_: object) -> AsyncGenerator[StreamChunk]: await chat_pipeline.run_turn(_TRANSCRIPT, turn.emit) - assert ("content", {"text": "About "}) in turn.events + assert turn.list_texts() == ["First paragraph."] terminal = _get_only_terminal(turn) assert terminal["state"] == TerminalState.ERROR assert terminal["reason"] == ErrorType.SYNTHESIS_STREAM_ERROR @@ -1968,7 +2140,7 @@ async def fail_tools() -> list[mcp_client.McpTool]: # The serialized frames of a complete turn whose stubs also report a thought # from discovery and from synthesis. The HMAC is deterministic because the # `signing_secret` fixture fixes the signing key. -_WIRE_HMAC = "6632f0de072e5eb60fc92d150cd93b07bf6600cdd18512a42af57790e178d7b0" +_WIRE_HMAC = "bab7873e1b550debfd0193a70934a1781d7cae19c5be8eea5e6dfe069f2917dc" _WIRE_FRAMES: list[tuple[EventName, str]] = [ ( "status", @@ -1991,15 +2163,18 @@ async def fail_tools() -> list[mcp_client.McpTool]: ("content", '{"sources": [{"name": "INSEE", "url": "https://insee.fr/"}]}'), ("status", '{"phase": "synthesis", "message": "Generating response..."}'), ("thought", '{"thought": "Writing the answer.", "phase": "synthesis"}'), - ("content", '{"text": "About 68 million."}'), ( - "status", - '{"phase": "chart_config", "message": "Preparing charts..."}', + "content", + '{"block": {"index": 0, "type": "text", ' + '"markdown": "About 68 million."}}', ), ( "content", - '{"chart_config": {"should_render": true, "charts": ' - '[{"title": "Population"}]}}', + '{"block": {"index": 1, "type": "chart", "chart": ' + '{"viz_type": "line", "title": "Total Population", ' + '"variable_dcids": ["Count_Person"], ' + '"place_dcids": ["country/FRA"], ' + '"facet_ids": {"Count_Person": "f1"}}}}', ), ( "terminal", @@ -2008,7 +2183,10 @@ async def fail_tools() -> list[mcp_client.McpTool]: '{"scopes": [{"places": {"country/FRA": "France"}, ' '"parent_place": null, "child_place_type": null, "variables": ' '{"Count_Person": "Total Population"}, "date_range": ' - '["2022", "2023"]}]}, "compacted_summary": null, "window": ' + '["2022", "2023"]}, {"places": {"country/FRA": "France"}, ' + '"parent_place": null, "child_place_type": null, "variables": ' + '{"Count_Person": "Total Population"}, "date_range": null}]}, ' + '"compacted_summary": null, "window": ' '[{"turn_index": 0, "idempotency_key": "test-idempotency-key", ' f'"hmac": "{_WIRE_HMAC}"}}]}}', ), diff --git a/narratives/agent/src/narratives_agent/workflows/contracts/synthesis.md b/narratives/agent/src/narratives_agent/workflows/contracts/synthesis.md new file mode 100644 index 00000000..f479b4fc --- /dev/null +++ b/narratives/agent/src/narratives_agent/workflows/contracts/synthesis.md @@ -0,0 +1,143 @@ + +### Citations + +- Place the step's `cite:` marker, `[1]` or a composite such as `[1, 2]`, + immediately after every bolded figure: `**5.6 million** [1]`, + `**64.25%** [1, 2]`. One marker at the end of a sentence or clause does + not cover the figures before it; each figure carries its own. +- A figure drawn from a step whose `cite:` is composite carries the whole + marker, `[1, 2]`, so every source that contributed is cited. +- Use only the numbers in the NUMBERED SOURCES block, reuse the same number + for the same source, never renumber, and never cite a number the block + does not list. Each listed source supplied some of the figures, so each + should carry at least one marker. +- The reader's citation list is rendered from that same block, so do not + write a Sources or Citations section of your own; end on the last + sentence of the analysis. + +### Writing the answer + +- Lead with the figure the question asked for, then the comparison, the + trend, or the spread the digest supports, then the caveats. +- Write in paragraphs. Use a short bulleted list for a handful of places or + categories. Never reproduce a whole cohort as a table or a list: name the + leading and trailing places the digest ranked, the median and the spread, + and let the chart carry the rest. +- Bold every figure and percentage. +- Report a figure at the precision the digest gives it. The UI shortens + long decimals for reading and keeps the full value on hover, so rounding + in the prose only removes what the reader would have hovered for. +- Calendar years by default; label a fiscal year as such when the series is + fiscal. +- Do not write section headings such as "From Data Commons". + +### Charts + +The chart block grammar follows these instructions. Use it as follows. + +- Place a chart block immediately after the paragraph that discusses what it + shows, with a blank line before and after it. A chart belongs where the + reader has just read the figures it illustrates. +- Chart a time series as a `line`; a value across every child place as a + `map`; a single value as a `highlight`; a composition or breakdown across + component tables as a `bar`. +- Chart a ranked cohort as a `ranking` with the tail the question asked for + (`bottom_n` for the lowest, `top_n` for the highest, both for both). When + the digest relates a second variable to a ranked cohort, chart each + variable as its own `ranking` of that cohort, and optionally pair the two + in a `scatter` across it. Use a `bar` of the named places for a handful + of places that are not a cohort. +- Chart the fetched tables the inventory lists, and chart a per-capita + value by naming the per-capita table the inventory marks chartable. It + renders per person even when the digest reports the rate per 1,000 or per + 100,000 people, so title the chart per person and state the digest's + scaled figures and scale in prose. A table marked "do not chart it" is + never charted; describe its figures in prose instead. +- Keep to the places the inventory lists for the table, or ask for the + cohort or a top or bottom slice. A `line` chart takes no `date` field; on + any other chart, omit `date` to use the aligned reference date. +- One chart per distinct point, at most three per answer, and none when the + digest holds no data. A chart block that breaks the grammar is dropped + without notice, so every figure must also be stated in the prose. + +### When the digest holds no data + +Say what was searched and what related data exists, drawn from the MISSING +DATA notes, in one or two short paragraphs. Include this sentence only when +no relevant data was found at all: "If you'd like to see this data in Data +Commons, you can +[submit a data request](https://docs.datacommons.org/contributing)." +Do not include it when data was found for some places, periods, or related +indicators. + +### Example + +The shape to follow. The places, figures, and source below are illustrative; +never reuse them as data. + +Colorado's population reached **5.9 million** [1] in 2023, growing +**11.6%** [1] from **5.3 million** [1] in 2013, about **61,200** [1] people +a year. + +```chart +{"kind": "line", "tables": ["t1"], "title": "Population of Colorado"} +``` + +Growth was fastest between 2013 and 2019 at **1.4%** [1] a year before +slowing after 2020. + +Minnesota's population stood at **5.7 million** [1] in 2023, up +**0.4 million** [1] (**7.5%** [1]) from 2013. + +Note that every bolded figure carries its own marker, and where it ends: on +the last sentence, with no citation list. `[1]` is row 1 of the NUMBERED +SOURCES block, and the reader sees that same row rendered beneath the answer. + +### Chart blocks + +Where a chart belongs, write a fenced code block tagged `chart` between two +paragraphs, holding one JSON object. The opening fence (```` ```chart ````) +and the closing fence (```` ``` ````) each stand on a line of their own; a +chart written inside a line of prose is not a chart. + +```chart +{"kind": "line", "tables": ["t1"], "title": "GDP, France and Germany", + "places": ["country/FRA", "country/DEU"]} +``` + +Fields every chart takes: +- `kind` (required): `line`, `bar`, `map`, `ranking`, `scatter`, or + `highlight`. +- `tables` (required): one or more table ids from the TABLE INVENTORY. +- `title` (required): a short descriptive title without the year. +- `unit`: a display label for the values. + +Fields by kind; a field not listed for that kind causes the block to be +rejected: +- Every kind except `line`: `date`, an observation date listed for the + tables, or `LATEST`; omit it to use the aligned reference date. A `line` + chart draws the full series and takes no `date` field. +- `line` and `bar`: tables that share a unit, and at most one of `places` + (a subset of the place DCIDs listed for those tables), `cohort: true` + (every child place of the tables' parent place), `top_n`, or `bottom_n` + (how many of the highest or lowest places to show). A `bar` also takes + `sort` (`ascending` or `descending`). +- `ranking`: one variable of a cohort, with `top_n`, `bottom_n`, or both, + equal when both are given. +- `map`: one variable of a cohort. +- `scatter`: exactly two tables, of two variables of one cohort. +- `highlight`: one variable, with `places` as a one-element list naming + the place to show; omit `places` when the table holds a single place. + +Rules: reference only table ids, place DCIDs, and dates that the TABLE +INVENTORY lists; a derived table marked "do not chart it" cannot appear in +`tables`; at most three chart blocks per answer; never place a chart block +inside a list or a table; always leave a blank line before and after a +chart block. A block that breaks these rules is dropped without notice, so +keep every figure in the prose as well. diff --git a/narratives/agent/src/narratives_agent/workflows/follow_up.py b/narratives/agent/src/narratives_agent/workflows/follow_up.py index 13b07e87..6fa3bd82 100644 --- a/narratives/agent/src/narratives_agent/workflows/follow_up.py +++ b/narratives/agent/src/narratives_agent/workflows/follow_up.py @@ -52,7 +52,8 @@ async def generate_follow_up_questions( Args: user_message: The user's query. - topics: The chart titles that ground the questions. + topics: The digest scopes, chart titles, and missing-data notes that + ground the questions. config: The agent configuration; `prompts.follow_up` is the system instruction and `gemini.follow_up_model` the model. token_usage: Optional turn totals that receive the call's token diff --git a/narratives/agent/src/narratives_agent/workflows/mcp_loop.py b/narratives/agent/src/narratives_agent/workflows/mcp_loop.py index ae650bdc..34c9403d 100644 --- a/narratives/agent/src/narratives_agent/workflows/mcp_loop.py +++ b/narratives/agent/src/narratives_agent/workflows/mcp_loop.py @@ -323,7 +323,7 @@ def _select_unique(values: list[str], cap: int) -> list[str]: return list(dict.fromkeys(values))[:cap] -def format_count(count: int, singular: str, plural: str | None = None) -> str: +def _format_count(count: int, singular: str, plural: str | None = None) -> str: """Formats `count` with its grammatical singular or plural noun. Args: @@ -423,9 +423,9 @@ def _build_server_report( listed = ", ".join(fetched) if len(unique_empty) > len(fetched): extra = len(unique_empty) - len(fetched) - listed += f", and {format_count(extra, 'other variable')}" + listed += f", and {_format_count(extra, 'other variable')}" tested = ( - f"the {format_count(len(unique_empty), 'variable')} fetched had " + f"the {_format_count(len(unique_empty), 'variable')} fetched had " f"no rows for the requested entities ({listed})" ) else: @@ -1172,5 +1172,4 @@ async def execute_discovery_loop( "McpLoopError", "StatusCallback", "execute_discovery_loop", - "format_count", ] diff --git a/narratives/agent/src/narratives_agent/workflows/messages.py b/narratives/agent/src/narratives_agent/workflows/messages.py index 2a361d65..387c6de1 100644 --- a/narratives/agent/src/narratives_agent/workflows/messages.py +++ b/narratives/agent/src/narratives_agent/workflows/messages.py @@ -21,7 +21,6 @@ DISCOVERY_STATUS = "Querying data tools..." PROCESSING_STATUS = "Analyzing data..." SYNTHESIS_STATUS = "Generating response..." -CHART_CONFIG_STATUS = "Preparing charts..." # Refusal messages carried on `terminal` SSE frames with `state == "refused"`. # Adversarial and safety-blocked turns always return a fixed string rather @@ -57,7 +56,6 @@ PER_CAPITA_CHART_TITLE_TEMPLATE = "{variable_name} per person" __all__ = [ - "CHART_CONFIG_STATUS", "CHAT_DISABLED", "CLASSIFY_STATUS", "DISCOVERY_STATUS", diff --git a/narratives/agent/src/narratives_agent/workflows/prompt_contracts.py b/narratives/agent/src/narratives_agent/workflows/prompt_contracts.py index cb4c0af8..38150ea6 100644 --- a/narratives/agent/src/narratives_agent/workflows/prompt_contracts.py +++ b/narratives/agent/src/narratives_agent/workflows/prompt_contracts.py @@ -36,7 +36,7 @@ ] # The packaged contracts under `narratives_agent.workflows.contracts`. -type ContractName = Literal["discovery"] +type ContractName = Literal["discovery", "synthesis"] _CONTRACTS_PACKAGE = "narratives_agent.workflows.contracts" diff --git a/narratives/agent/src/narratives_agent/workflows/prompt_contracts_test.py b/narratives/agent/src/narratives_agent/workflows/prompt_contracts_test.py index 71c5b041..f476c1b2 100644 --- a/narratives/agent/src/narratives_agent/workflows/prompt_contracts_test.py +++ b/narratives/agent/src/narratives_agent/workflows/prompt_contracts_test.py @@ -18,10 +18,16 @@ from typing import get_args import pytest +from pydantic import TypeAdapter from narratives_agent.config import strip_prompt_comments from narratives_agent.gemini.schemas import QueryClassification from narratives_agent.workflows import prompt_contracts +from narratives_agent.workflows.charts import ( + BarSort, + ChartKind, + ChartSpecification, +) from narratives_agent.workflows.mcp_loop import ( DISCOVERY_TOOL_NAMES, SERVER_FUNCTION_NAMES, @@ -173,3 +179,85 @@ def test_discovery_contract_tables_match_declared_tools_and_scripts() -> None: assert tools == DISCOVERY_TOOL_NAMES | SERVER_FUNCTION_NAMES assert scripts == set(mapping) + + +def test_composed_default_synthesis_prompt_includes_guidance_and_contract() -> ( + None +): + # Test: Composition of the shipped default `synthesis.md` guidance with + # the packaged synthesis contract. + # Situation: `defaults/prompts/synthesis.md` is read and stripped of HTML + # comments, then composed with the `synthesis` contract. + # Expectation: The composed instruction begins with the role heading from + # `defaults/prompts/synthesis.md`, includes the citation and chart-block + # sections from `workflows/contracts/synthesis.md`, and contains no HTML + # comments. + repo_root = Path(__file__).resolve().parents[4] + guidance = strip_prompt_comments( + (repo_root / "defaults" / "prompts" / "synthesis.md").read_text( + encoding="utf-8" + ) + ) + + composed = compose_system_instruction(guidance, "synthesis") + + assert composed == f"{guidance}\n\n{load_prompt_contract('synthesis')}" + assert composed.startswith("## Data Agent: Narrative Synthesis") + assert "\n\n### Citations\n\n" in composed + assert "\n\n### Chart blocks\n\n" in composed + assert " -### Role: Statistical Intelligence Synthesizer - -You are Data Agent, turning statistical observations into one clear, sourced -answer. - -### Your Task - -Answer the question from the data provided. Write it as prose, not as a -report: do not label a section "From Data Commons". - -### CRITICAL GROUNDING RULES - ZERO TOLERANCE FOR HALLUCINATION - -**YOU MUST FOLLOW THESE RULES EXACTLY:** - -1. **ONLY use information from the DATA RESULTS provided below** -2. **DO NOT use your training knowledge or make assumptions** -3. **Every factual statement MUST include an inline citation** -4. **If information is not in the provided context, say: "I don't have this specific data in the current dataset."** -5. **Judge availability by whether there are observation VALUES, not by source metadata. An empty `alternative_sources` array, or `source_id:unknown`, is normal for single-provenance series and does NOT mean the data is missing — report the values and attribute them to the provenance that is present. Treat data as unavailable only when the observation payload itself is empty (`"data":{}`).** -6. **ONLY cite numbers that appear EXACTLY in the data results - do not calculate or derive new numbers** -7. **Do NOT create comparison tables unless the data explicitly contains all values for each row** -8. **If asked for a sub-national breakdown when only country-level data exists, say so plainly: "A sub-national breakdown is not available for this indicator." Use the term the place itself uses — states, provinces, regions, departments.** -9. **NEVER expose technical details like field names, JSON structures, or internal identifiers to the user** -10. **Always state the year of every figure, and name the place. A number without its year and place is not an answer.** - -### Citation Format (NotebookLM Style) - -The context below contains a **NUMBERED SOURCES** block. Those numbers are the -citation numbers. They are not a suggestion and not a starting point: the -reader's Citations list is that same numbered list, rendered. `[2]` in your -prose and `[2]` beside the source are the same row, so a number you invent -points at the wrong source, or at nothing. - -The reader may see a different numeral than you wrote: across a conversation the -UI gives each source one number for the whole thread, so your `[1]` can render -as `[4]` when the first three sources were cited in earlier answers. That -remapping is the UI's, and it only works while your numbers match the block you -were given. Keep numbering from the block and nothing else. - -**Inline Citation Rules:** -- Place the marker immediately after the fact or statistic: `**3.8%** [2]`. -- Use **only** numbers that appear in the NUMBERED SOURCES block. Never renumber - them, never start again from [1], and never cite a number the block does not - list. -- Cite the number of the source the figure actually came from. Where you passed - `source_override`, that is the source you name. -- Reuse the same number every time you cite the same source. -- Every source in the block supplied some of the numbers below: the list is - built from what the retrieval calls actually returned, not from what was - offered to you. So each entry should carry at least one marker somewhere in - your answer. If you find yourself citing only one of several, check you have - not attributed one source's figures to another. - ---- - -### Do not write a Sources or Citations section - -Do not end the response with a `Sources` or `Citations` list, a `**Sources:**` -heading, or any other roll-call of the sources you cited. The reader already -gets one: the UI renders it from the same numbered list you were given, at the -foot of your answer. - -Anything you write here is a *second* sources section, numbered by you rather -than by the list, sitting directly above the real one. Two lists that disagree -is worse than either alone. End on the last sentence of the analysis. - -### Example Response (Data Query) - -The shape to follow. The place, figures and source below are placeholders — -never reuse them as data. - -"PLACE_NAME's population was **5.6 million** in 2023 [1]. Across all places -reporting this indicator, the median was **10.2 million** in the same year [1]." - -Note where it ends: on the last sentence, with no Citations list. `[1]` is row -1 of the NUMBERED SOURCES block it was given, and the reader sees that same row -1 rendered beneath the answer. - -### What NOT to do - STRICTLY PROHIBITED - -- Do NOT make claims without citations -- Do NOT use phrases like "typically", "generally", "approximately" without source -- Do NOT fill gaps with general knowledge -- Do NOT hallucinate statistics, percentages, or rankings -- Do NOT create place or sector comparisons unless data exists for ALL items -- Do NOT invent trends, drivers, or key factors not in the data -- Do NOT name specific places or organizations unless they appear in the data -- Do NOT present a partial set of countries as a world total without saying so -- Do NOT assume a currency, or convert values into one the source did not use - -### When to Include Data Request Link - -**ONLY include the data request link when NO relevant data was found at all.** - -Include: "If you'd like to see this data in Data Commons, you can [submit a data request](https://docs.datacommons.org/contributing)." - -**ONLY when:** -- The search found ZERO matching variables for the user's query -- The query topic is completely outside our dataset coverage - -**DO NOT include the link when:** -- You found data for SOME places (partial coverage is still useful) -- You found related data that answers the core question -- Data exists but for different time periods than requested -- You successfully provided numerical answers with citations - -### Response Structure - -**Rich Formatting Rules**: -- **Bold** for every numerical value and percentage -- Bullet lists for place or category comparisons -- Trigger a chart when you have a time series or more than three places - -**Number Formatting**: -- Report values in the **units and currency the source uses**. Name the currency - explicitly whenever a figure is monetary. -- Use plain international scale words — thousand, million, billion. Do not use - lakh or crore unless the user's own question uses them. -- Report a figure at the precision the source gives it. The UI rounds long - decimals to two places for reading and keeps the full value on hover, so - rounding here as well only destroys what the reader would have hovered for. -- Calendar years by default. Use a fiscal year only when the series is itself - fiscal, and label it as such. - -### Follow-Up Suggestions - -Base suggestions on the data you actually returned — a comparison across -countries, a trend over the available years, a related indicator from the same -provenance, or a sub-national breakdown where one exists. - -### Important Guidelines - -- Lead with the number the user asked for, then add trend and comparison. -- Prefer the most recent year with broad coverage over a more recent year - reported by only a few places. -- Be explicit about coverage limits. "42 countries reported this in 2022" is a - better answer than an unqualified global figure. - -### CRITICAL: Handling RAW MCP TOOL RESULTS - -When you receive **RAW MCP TOOL RESULTS**, you MUST: - -1. **Look for actual numeric data** in the tool results: - - Search for "time_series", "observations", "values", "value" keys - - Look inside nested structures like "place_observations" - - Extract dates and their corresponding values - -2. **If numeric data exists, ALWAYS include it in your response**: - - Present the actual numbers, with their place and year - - Show the trend if multiple years are present - - Attribute them to the provenance in `sourceMetadata` - -3. **If NO data was found** (empty results, errors, or tool failures): - - Explain what was searched - - Name any related data that was found - - Do not present the absence of data as a value of zero +### Role: Statistical Narrative Writer + +You are Data Agent. You turn a verified analysis into one clear, sourced +answer, written as prose with charts placed where they help. + +### What you receive + +The user message carries, in this order: + +1. **User Query**: the question, exactly as asked. +2. **NUMBERED SOURCES**: the sources that supplied the figures, numbered. + These numbers are the citation numbers. +3. **ANALYSIS DIGEST**: one block per analysis step, naming the script, the + variables, the unit, the coverage (how many places reported, on which + date), the method, the computed metrics as JSON, and any caveats. Every + number you may report is in here. +4. **MISSING DATA**: what was searched for and not found, when anything is + missing. +5. **TABLE INVENTORY**: the cached tables a chart block may reference, each + with its id, variable, unit, places, and dates. + +### Grounding rules + +1. **Report only figures that appear in the ANALYSIS DIGEST, exactly as + given.** Do not add, subtract, divide, average, round further, or + extrapolate. The server computed every change, share, rank, and growth + rate the answer can state; if a figure is not in the digest, the answer + does not have it. +2. **Carry every caveat.** A refused step, an undefined growth rate, a + partial cohort, or an alignment across mixed vintages must be stated in + plain words. Never supply an estimate where the digest refused one. +3. **State the place, the date, and the unit of every figure.** A number + without those three is not an answer. Use the reference date the digest + aligned to, and say when places were aligned to a shared year. +4. **State coverage.** "42 of 54 countries reported this in 2022" is a better + answer than an unqualified total. When the digest says a place is + missing, say so. +5. **Do not fill gaps from general knowledge.** When the MISSING DATA notes + say an indicator is unavailable, say that it is unavailable and name what + related data exists. Do not present the absence of data as zero. +6. **Never expose internals.** No table ids, DCIDs, script names, step ids, + field names, or JSON in the prose, and no remarks about what the system + can or cannot compute or which methods it has. The only caveats are the + ones the ANALYSIS DIGEST and the MISSING DATA notes state and any scope + statement the analytical directive requires. Name variables and places + as a reader would. +7. **Answer in the data's own units and currency.** Name the currency of + every monetary figure. Use plain scale words: thousand, million, billion. +8. **Name the places a group covers.** When the question names a group of + places, name the places the figures cover once, where the group first + matters. diff --git a/narratives/deploy/terraform-custom-datacommons/modules/main.tf b/narratives/deploy/terraform-custom-datacommons/modules/main.tf index 8199c6b8..ec5e2c4f 100644 --- a/narratives/deploy/terraform-custom-datacommons/modules/main.tf +++ b/narratives/deploy/terraform-custom-datacommons/modules/main.tf @@ -88,8 +88,8 @@ resource "google_cloud_run_v2_service" "dc_app_service" { resources { # All CPU-hungry work happens inside a request: the SSE stream holds - # the request open for its duration, and the chart-config thread is - # joined before the generator returns. Startup is covered by the boost. + # the request open for its duration through follow-up generation. + # Startup is covered by the boost. cpu_idle = true limits = { cpu = var.agent_cpu diff --git a/narratives/docs/architecture.md b/narratives/docs/architecture.md index 5a1810dd..73efd13f 100644 --- a/narratives/docs/architecture.md +++ b/narratives/docs/architecture.md @@ -66,6 +66,10 @@ When an answer includes a chart, the UI renders Data Commons web components their own chart data and place names directly from Data Commons (`datacommons.org` or the same-origin Data Commons website host). +Prose figures come from the digest while charts are drawn by the Data Commons +web components from their own fetch, pinned by `facet_ids`, so the reference +date and the per-capita denominator can differ between prose and chart. + ## MCP tool discovery & session recovery ### Automatic tool discovery @@ -110,9 +114,10 @@ agent loop in five steps: not signed and is excluded from the multi-turn window, and the UI renders the message as a notice rather than an error card. - **Clarification:** When a place name is ambiguous or required and omitted, - streams the clarification question as answer `text`, signs the turn with - empty state slots, and emits a signed `complete` terminal frame without - `data_status`, `sources`, `chart_config`, or `follow_ups` frames. + streams the clarification question as the answer's one text block, signs + the turn with empty state slots, and emits a signed `complete` terminal + frame without `data_status`, `sources`, chart block, or `follow_ups` + frames. - **Classified:** Continues into the discovery loop and synthesis with the analytical `pattern` (`single_point`, `top_n_ranking`, `distribution`, `trend_delta`, `outliers`, `correlation`, `composition`, @@ -138,12 +143,44 @@ agent loop in five steps: default specification is built from the cached tables and the query's pattern. If the MCP server cannot be reached, the turn stops and returns an error so the model never guesses numbers without data. -3. **Answer synthesis & chart selection (`synthesis` / `chart_config`):** - Gemini receives the rendered `AnalysisDigest`, the inventory of cached - tables the specification used, and the numbered sources (`[1]`, `[2]`, ...), - and writes a markdown response with inline citations while checking in - parallel whether the data can be shown in a chart. A chart is included only - if validation confirms that matching data exists. +3. **Answer synthesis (`synthesis`):** The synthesis phase + (`workflows/synthesis.py`, `workflows/blocks.py`, `workflows/charts.py`, and + `workflows/chart_defaults.py`) writes the answer from the digest alone. The + prompt holds the digest's computed results (with floats rounded to display + precision), the numbered sources, the missing-data notes, the pattern's + directive, and a table inventory naming each cached table's variable, unit, + places, and dates; externally sourced text (variable names, source names, and + missing-data notes) stays in typed fields until the message is rendered. The + model writes markdown and places each chart where it belongs as a fenced + block tagged `chart` holding one specification: the chart kind (`line`, + `bar`, `map`, `ranking`, `scatter`, or `highlight`), the table ids it draws + from, a title, and the places (an explicit subset, the table's cohort, or a + `top_n` or `bottom_n` slice), with an optional unit and, on single-date + charts, an optional date. + The server splits the stream at paragraph boundaries and at the closing fence + of each chart block, and emits one `content {block}` frame per completed + paragraph and per chart that resolves. Resolution is set membership against + the turn cache: an unknown table, a place or date the table does not hold, a + unit mismatch across the tables on one axis, or a kind the table shape cannot + support (a map without a cohort, a scatter without two variables) rejects the + chart, which is dropped and counted while the prose stands. Each resolved + chart carries `facet_ids` mapping its variables to the facets that served + them, matching the Data Commons web components' facet attributes (`facetId` + on highlight, `facetIds` or `facetMapping` on the other components). A + derived per-capita table is scaled automatically in the digest (per 1, + 1,000, or 100,000 people so the median non-zero magnitude is at least 1) and + is charted when referenced by its `s` table id as its numerator with the + web component's per-person normalization; other derived tables are not + chartable. At most three charts are accepted per answer. When the digest has + data and no chart is emitted, the default chart floor appends at most one + chart projected directly from the primary result (`scatter` for a cohort + `relate_cohorts` result, which is considered first, or the direct chart of + `analyze_timeseries`, `rank_cohort`, or `compute_composition`; + distributions, standalone derivations, and non-cohort relations receive + none). The signed `model_response` is the text blocks joined by blank lines, + with the chart blocks stripped, matching the text blocks emitted in + `content` frames. A stream that stalls for 60 seconds waiting for a chunk + ends the turn with `synthesis_stream_error`. 4. **Turn signing (`finalize`):** The server builds the turn's `state_slots` from the fetched tables the specification used and the displayed charts (`workflows/state_slots.py`), and signs the completed turn into the window @@ -159,9 +196,9 @@ for `heartbeat`) a sequential `id` starting at `1`: | Event | What it contains | | :--- | :--- | -| `status` | Current step (`classify`, `discovery`, `processing`, `synthesis`, or `chart_config`) and status message | +| `status` | Current step (`classify`, `discovery`, `processing`, or `synthesis`) and status message | | `thought` | Model reasoning summary from the `discovery` or `synthesis` step | -| `content` | Answer text (`{text}`), tool call details (`{tool_call}`), citations (`{sources}`), data status (`{data_status}`), or chart config (`{chart_config}`) | +| `content` | Answer block (`{block}` of type `text` or `chart`), tool call details (`{tool_call}`), citations (`{sources}`), or data status (`{data_status}`) | | `terminal` | Final turn status (`state: "complete"`, `"refused"`, or `"error"`), plus the signed conversation state on completion, `{reason, message}` (with no `error` key) on refusal, or `{error, reason}` on failure | | `follow_ups` | `{follow_up_questions}`, sent after `terminal` when available | | `heartbeat` | `{}`, sent every 15 seconds so proxies do not close idle connections | diff --git a/narratives/docs/deployment.md b/narratives/docs/deployment.md index 3f16331b..f4e8b66a 100644 --- a/narratives/docs/deployment.md +++ b/narratives/docs/deployment.md @@ -294,15 +294,23 @@ code-owned contract (`workflows/contracts/discovery.md`, composed by in `config/prompts/discovery.md` should add instance-specific domain and regional guidance rather than restating the tool or specification rules. +`synthesis.md` holds the narrative role, input description, and grounding rules +for the synthesis phase. The `[n]` and `[n, m]` citation markers, answer-writing +rules, and fenced `chart` block grammar live in a code-owned contract +(`workflows/contracts/synthesis.md`, composed by `workflows/prompt_contracts.py` +together with the per-pattern synthesis directive from +`workflows/pattern_guidance.py`) that the server appends after the configured +`synthesis` prompt on every request and that cannot be overridden. An override +in `config/prompts/synthesis.md` should add instance-specific tone, domain, or +grounding guidance rather than restating the citation or chart-block grammar. + `classify.md` governs the admission phase that runs at the start of every turn to decide whether a question is served, refused, or answered with a clarification question, and which analytical pattern it follows. Each LLM stage reads its own model key under `gemini` in `agent-config.json`: `gemini.classify_model` (default `gemini-3.5-flash-lite`), and -`gemini.discovery_model`, `gemini.synthesis_model`, -`gemini.chart_config_model`, `gemini.chart_validation_model`, -`gemini.follow_up_model`, and `gemini.compaction_model` (each defaulting to -`gemini-3.8-flash`). +`gemini.discovery_model`, `gemini.synthesis_model`, `gemini.follow_up_model`, +and `gemini.compaction_model` (each defaulting to `gemini-3.8-flash`). Prompts support runtime template placeholders: @@ -509,11 +517,15 @@ instead. For further diagnostics, inspect the Cloud Run service logs. Each completed chat turn writes a structured `chat_turn` JSON log entry keyed by `turn_id`, containing the terminal state, error reason (if any), phase latencies -(`classify`, `discovery`, `processing`, `synthesis`, `chart_config`, `finalize`, -`follow_ups`), tool names and counts, the analytical `scripts` run, -`tables_fetched`, `total_rows_fetched`, `processing_retries`, token usage, and -the analytical `pattern` assigned by the `classify` phase (user query text is -never logged). On a `refused` turn, `error_type` holds the refusal reason -(`adversarial`, `off_domain`, `safety_blocked`, or `chat_disabled`) and -`pattern` is omitted; if the classifier call itself fails, the turn falls -through with `pattern` set to `exploratory`. +(`classify`, `discovery`, `processing`, `synthesis`, `finalize`, `follow_ups`), +tool names and counts, token usage, the analytical `pattern` assigned by the +`classify` phase, and the analysis counts (`scripts`, `tables_fetched`, +`total_rows_fetched`, `processing_retries`, `charts_emitted`, and +`charts_rejected`; user query text is never logged). The turn span carries the +same analysis counts as `narratives.scripts`, `narratives.tables_fetched`, +`narratives.total_rows_fetched`, `narratives.processing_retries`, +`narratives.charts_emitted`, and `narratives.charts_rejected`. On a `refused` +turn, `error_type` holds the refusal reason (`adversarial`, `off_domain`, +`safety_blocked`, or `chat_disabled`) and `pattern` is omitted; if the +classifier call itself fails, the turn falls through with `pattern` set to +`exploratory`. diff --git a/narratives/docs/development.md b/narratives/docs/development.md index 39f8bc90..18be1387 100644 --- a/narratives/docs/development.md +++ b/narratives/docs/development.md @@ -216,7 +216,7 @@ narratives/ │ ├── telemetry.py OpenTelemetry and structured logging │ ├── server/ FastAPI app and routes (spa, chat, system) │ ├── mcp/ MCP client, capabilities, schemas, and parsers -│ ├── workflows/ Agent loop, discovery, fetch, spec, processing +│ ├── workflows/ Agent loop, discovery, processing, synthesis │ ├── analytics/ Tables, alignment, summaries, digest, scripts │ └── gemini/ Gemini SDK client and Pydantic schemas ├── ui/ React 19 + Vite frontend (managed with pnpm) diff --git a/narratives/schemas/agent-config.example.json b/narratives/schemas/agent-config.example.json index 4cc03b1a..bea30bd3 100644 --- a/narratives/schemas/agent-config.example.json +++ b/narratives/schemas/agent-config.example.json @@ -8,8 +8,6 @@ "classify_model": "gemini-3.5-flash-lite", "discovery_model": "gemini-3.8-flash", "synthesis_model": "gemini-3.8-flash", - "chart_config_model": "gemini-3.8-flash", - "chart_validation_model": "gemini-3.8-flash", "follow_up_model": "gemini-3.8-flash", "compaction_model": "gemini-3.8-flash" }, diff --git a/narratives/schemas/agent-config.schema.json b/narratives/schemas/agent-config.schema.json index 578c1115..ce2389b3 100644 --- a/narratives/schemas/agent-config.schema.json +++ b/narratives/schemas/agent-config.schema.json @@ -36,18 +36,8 @@ "type": "string", "default": "gemini-3.8-flash" }, - "chart_config_model": { - "description": "Model used for chart configuration, where chart specifications are selected from the gathered observations.", - "type": "string", - "default": "gemini-3.8-flash" - }, - "chart_validation_model": { - "description": "Model used for chart validation, which checks whether the synthesized answer contains chartable data.", - "type": "string", - "default": "gemini-3.8-flash" - }, "follow_up_model": { - "description": "Model used for follow-up generation, which suggests follow-up questions grounded in the chart topics.", + "description": "Model used for follow-up generation, which suggests follow-up questions grounded in the turn's digest, chart titles, and missing-data notes.", "type": "string", "default": "gemini-3.8-flash" },