From f6f0c80fb0d78c424e926198cfe2a24801f0d2ba Mon Sep 17 00:00:00 2001 From: Nick Blumberg Date: Tue, 20 Oct 2026 14:56:48 -0700 Subject: [PATCH] refactor: decouple Gemini model configuration per role --- .../src/narratives_agent/config_model.py | 74 +++++--- .../src/narratives_agent/config_model_test.py | 166 ++++++++++++------ .../workflows/chart_config.py | 4 +- .../narratives_agent/workflows/follow_up.py | 2 +- .../narratives_agent/workflows/transcript.py | 2 +- narratives/docs/deployment.md | 12 +- narratives/schemas/agent-config.example.json | 7 +- narratives/schemas/agent-config.schema.json | 31 +++- 8 files changed, 204 insertions(+), 94 deletions(-) diff --git a/narratives/agent/src/narratives_agent/config_model.py b/narratives/agent/src/narratives_agent/config_model.py index 3fc188e4..385bd4f9 100644 --- a/narratives/agent/src/narratives_agent/config_model.py +++ b/narratives/agent/src/narratives_agent/config_model.py @@ -75,6 +75,11 @@ class RoleConfig: thinking_level: types.ThinkingLevel +# `GeminiConfig` attribute that holds the model name for each role. +_ROLE_MODEL_ATTRS: Mapping[GeminiRole, str] = { + role: f"{role.value}_model" for role in GeminiRole +} + # Sampling temperature of each role. agent-config.json cannot change these. _ROLE_TEMPERATURES: Mapping[GeminiRole, float] = { GeminiRole.CLASSIFY: 0.0, @@ -104,8 +109,7 @@ class RoleConfig: # The thinking levels that `thinking.discovery_level` and # `thinking.synthesis_level` accept, keyed by their configuration string. They # match the `enum` that `schemas/agent-config.schema.json` declares for both -# keys; `MINIMAL` is excluded because the default `gemini.discovery_model` -# rejects it. +# keys; `MINIMAL` is excluded because `DEFAULT_GEMINI_MODEL` rejects it. _CONFIGURABLE_THINKING_LEVELS: Mapping[str, types.ThinkingLevel] = { "low": types.ThinkingLevel.LOW, "medium": types.ThinkingLevel.MEDIUM, @@ -163,18 +167,22 @@ def _clamp_thinking_level( @dataclass(frozen=True) class GeminiConfig: - """Holds the `gemini` section: the model each phase calls and the key. - - `discovery_model` runs the tool loop, synthesis, chart configuration, - follow-up generation, and transcript compaction; `classify_model` runs - the classify phase. `api_key` is the local-development key from - `config.json`, which `get_gemini_api_key` in `config.py` applies after - Secret Manager; it is excluded from `repr` so that a logged or printed - configuration never carries the key. + """Holds the `gemini` section: the model each role calls and the key. + + Each `_model` field names the model for the matching `GeminiRole` + in `AgentConfig.resolve_role_config`. `api_key` is the local-development + key from `config.json`, which `get_gemini_api_key` in `config.py` + applies after Secret Manager; it is excluded from `repr` so that a + logged or printed configuration never carries the key. """ - discovery_model: str = DEFAULT_GEMINI_MODEL 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) @@ -253,21 +261,16 @@ def __post_init__(self) -> None: def resolve_role_config(self, role: GeminiRole) -> RoleConfig: """Returns the model, temperature, and thinking level of `role`. - `CLASSIFY` runs on `gemini.classify_model` and every other role on - `gemini.discovery_model`. `DISCOVERY` and `SYNTHESIS` think at - `thinking.discovery_level` and `thinking.synthesis_level`; the other - roles think at their level in `_FIXED_THINKING_LEVELS`. A level - below the lowest the role's model accepts is raised to that floor, - so the classifier runs at `LOW` on a model that rejects `MINIMAL` - rather than failing its call on every turn and leaving the turn - unclassified. Every role samples at its temperature in - `_ROLE_TEMPERATURES`. + Each role runs on its `gemini._model` field (`_ROLE_MODEL_ATTRS`). + `DISCOVERY` and `SYNTHESIS` think at `thinking.discovery_level` and + `thinking.synthesis_level`; the other roles think at their level in + `_FIXED_THINKING_LEVELS`. A level below the lowest the role's model + accepts is raised to that floor, so the classifier runs at `LOW` on + a model that rejects `MINIMAL` rather than failing its call on every + turn and leaving the turn unclassified. Every role samples at its + temperature in `_ROLE_TEMPERATURES`. """ - model = ( - self.gemini.classify_model - if role is GeminiRole.CLASSIFY - else self.gemini.discovery_model - ) + model: str = getattr(self.gemini, _ROLE_MODEL_ATTRS[role]) if role is GeminiRole.DISCOVERY: thinking_level = self.thinking.discovery_level elif role is GeminiRole.SYNTHESIS: @@ -433,11 +436,28 @@ def parse_agent_config(document: Mapping[str, Any]) -> AgentConfig: prompts = _read_section(document, "prompts") return AgentConfig( gemini=GeminiConfig( + classify_model=_read_model_name( + gemini, "classify_model", GeminiConfig.classify_model + ), discovery_model=_read_model_name( gemini, "discovery_model", GeminiConfig.discovery_model ), - classify_model=_read_model_name( - gemini, "classify_model", GeminiConfig.classify_model + 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 + ), + compaction_model=_read_model_name( + gemini, "compaction_model", GeminiConfig.compaction_model ), api_key=_read_string( gemini, "gemini", "api_key", GeminiConfig.api_key diff --git a/narratives/agent/src/narratives_agent/config_model_test.py b/narratives/agent/src/narratives_agent/config_model_test.py index 993373ef..677e7e24 100644 --- a/narratives/agent/src/narratives_agent/config_model_test.py +++ b/narratives/agent/src/narratives_agent/config_model_test.py @@ -33,7 +33,6 @@ """ import logging -from dataclasses import replace from typing import Any import pytest @@ -53,6 +52,8 @@ parse_agent_config, ) +_ABSENT = object() + def _warnings(caplog: pytest.LogCaptureFixture) -> list[str]: """Returns the warning messages logged by the parser.""" @@ -72,8 +73,15 @@ def test_an_empty_document_gives_every_default() -> None: parsed = parse_agent_config({}) assert parsed == AgentConfig() - assert parsed.gemini.discovery_model == DEFAULT_GEMINI_MODEL - assert parsed.gemini.classify_model == DEFAULT_GEMINI_CLASSIFY_MODEL + assert parsed.gemini == GeminiConfig( + 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, + ) assert parsed.mcp == McpConfig(server_url="", enabled=True) assert parsed.thinking == ThinkingLevels( discovery_level=types.ThinkingLevel.MEDIUM, @@ -93,8 +101,13 @@ def test_configured_values_are_carried_through() -> None: parsed = parse_agent_config( { "gemini": { - "discovery_model": " gemini-3-pro ", "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 ", }, "mcp": {"server_url": " https://mcp.example ", "enabled": False}, @@ -114,8 +127,13 @@ def test_configured_values_are_carried_through() -> None: assert parsed == AgentConfig( gemini=GeminiConfig( - discovery_model="gemini-3-pro", 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", ), mcp=McpConfig(server_url="https://mcp.example", enabled=False), @@ -133,39 +151,42 @@ def test_configured_values_are_carried_through() -> None: ) +@pytest.mark.parametrize("role", list(GeminiRole), ids=lambda role: role.value) @pytest.mark.parametrize( - ("gemini", "warned"), + ("value", "warned"), [ - ({}, False), - ({"discovery_model": None, "classify_model": None}, True), - ({"discovery_model": "", "classify_model": ""}, True), - ({"discovery_model": " ", "classify_model": " "}, True), - ({"discovery_model": 7, "classify_model": ["gemini"]}, True), + (_ABSENT, False), + (None, True), + ("", True), + (" ", True), + (123, True), ], ids=["absent", "null", "empty", "whitespace", "wrong type"], ) def test_an_unusable_model_name_gives_the_default( caplog: pytest.LogCaptureFixture, - gemini: dict[str, Any], + role: GeminiRole, + value: object, warned: bool, ) -> None: # Test: Model-name fallback in `parse_agent_config`. - # Situation: `gemini.discovery_model` and `gemini.classify_model` are - # absent, null, empty, whitespace-only, or not strings. - # Expectation: Both models are the defaults. An absent name is not - # reported; every other case logs one warning per model naming its key. + # Situation: `gemini._model` is absent, null, empty, whitespace-only, + # or not a string. + # Expectation: The role's model falls back to its default on `GeminiConfig`. + # An absent key is not reported; every other case logs one warning naming + # `gemini._model`. + key = f"{role.value}_model" + gemini: dict[str, Any] = {} if value is _ABSENT else {key: value} with caplog.at_level(logging.WARNING): parsed = parse_agent_config({"gemini": gemini}) - assert parsed.gemini.discovery_model == DEFAULT_GEMINI_MODEL - assert parsed.gemini.classify_model == DEFAULT_GEMINI_CLASSIFY_MODEL + default_model = getattr(GeminiConfig, key) + assert getattr(parsed.gemini, key) == default_model + assert parsed.resolve_role_config(role).model == default_model warnings = _warnings(caplog) if warned: - assert len(warnings) == 2 - assert any( - "'gemini.discovery_model'" in message for message in warnings - ) - assert any("'gemini.classify_model'" in message for message in warnings) + assert len(warnings) == 1 + assert f"'gemini.{key}'" in warnings[0] else: assert warnings == [] @@ -301,8 +322,8 @@ def test_an_unrecognized_thinking_level_gives_the_default( # Test: Rejection of thinking levels outside the schema's enum in # `parse_agent_config`. # Situation: `thinking.discovery_level` and `thinking.synthesis_level` are - # empty, whitespace, `"minimal"` (which the default - # `gemini.discovery_model` rejects), or an unknown string. + # empty, whitespace, `"minimal"` (which `DEFAULT_GEMINI_MODEL` rejects), + # or an unknown string. # Expectation: Both fields fall back to `types.ThinkingLevel.MEDIUM`, and # one warning per field names its dotted key. with caplog.at_level(logging.WARNING): @@ -368,10 +389,12 @@ def test_a_default_configuration_resolves_each_role_to_its_defaults( # Test: `AgentConfig.resolve_role_config` without configured overrides. # Situation: The configuration is `AgentConfig()`, and every `GeminiRole` # is resolved. - # Expectation: Each role resolves to its default model, temperature, and - # thinking level; the classifier is the only role on the classify model - # and the only role at `MINIMAL`. - assert AgentConfig().resolve_role_config(role) == expected + # Expectation: Each role resolves to its default model on `GeminiConfig`, + # temperature, and thinking level; the classifier is the only role on + # the classify model and the only role at `MINIMAL`. + resolved = AgentConfig().resolve_role_config(role) + assert resolved == expected + assert resolved.model == getattr(GeminiConfig, f"{role.value}_model") @pytest.mark.parametrize("role", list(GeminiRole), ids=lambda role: role.value) @@ -387,45 +410,76 @@ def test_every_role_has_a_temperature_and_a_thinking_level( assert isinstance(AgentConfig().resolve_role_config(role), RoleConfig) +@pytest.mark.parametrize("role", list(GeminiRole), ids=lambda role: role.value) +def test_each_role_model_is_configured_independently( + role: GeminiRole, +) -> None: + # Test: Independent per-role model selection in `parse_agent_config` and + # `AgentConfig.resolve_role_config`. + # Situation: Only `gemini._model` is set in the configuration + # document. + # Expectation: `role` resolves to `"custom-model-v1"`, and every other + # `GeminiRole` resolves to its default model on `GeminiConfig`. + config = parse_agent_config( + {"gemini": {f"{role.value}_model": "custom-model-v1"}} + ) + + assert config.resolve_role_config(role).model == "custom-model-v1" + for other in GeminiRole: + if other is not role: + assert config.resolve_role_config(other).model == getattr( + GeminiConfig, f"{other.value}_model" + ) + + def test_configured_values_apply_to_the_roles_they_govern() -> None: # Test: `AgentConfig.resolve_role_config` with configured models and levels. - # Situation: `gemini.classify_model`, `gemini.discovery_model`, - # `thinking.discovery_level`, and `thinking.synthesis_level` are all set. - # Expectation: The classify model reaches `CLASSIFY` alone and the - # discovery model every other role; `discovery_level` reaches `DISCOVERY` - # and `synthesis_level` reaches `SYNTHESIS`, while the fixed thinking - # levels of the other roles and every temperature are unchanged. The - # classify model accepts `MINIMAL`, so the clamp leaves that level in - # place. + # Situation: Every `gemini._model` field, `thinking.discovery_level`, + # and `thinking.synthesis_level` are set to distinct values. + # Expectation: Each role receives its own configured model; + # `discovery_level` reaches `DISCOVERY` and `synthesis_level` reaches + # `SYNTHESIS`, while the fixed thinking levels of the other roles and + # every temperature are unchanged. The classify model accepts `MINIMAL`, + # so the clamp leaves that level in place. config = AgentConfig( gemini=GeminiConfig( - discovery_model="gemini-custom", classify_model="gemini-3.6-flash" + 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", ), thinking=ThinkingLevels( discovery_level=types.ThinkingLevel.HIGH, synthesis_level=types.ThinkingLevel.LOW, ), ) - defaults = AgentConfig() assert config.resolve_role_config(GeminiRole.CLASSIFY) == RoleConfig( "gemini-3.6-flash", 0.0, types.ThinkingLevel.MINIMAL ) assert config.resolve_role_config(GeminiRole.DISCOVERY) == RoleConfig( - "gemini-custom", 0.2, types.ThinkingLevel.HIGH + "gemini-discovery-custom", 0.2, types.ThinkingLevel.HIGH ) assert config.resolve_role_config(GeminiRole.SYNTHESIS) == RoleConfig( - "gemini-custom", 0.3, types.ThinkingLevel.LOW + "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 + ) + assert config.resolve_role_config(GeminiRole.COMPACTION) == RoleConfig( + "gemini-compaction-custom", 0.1, types.ThinkingLevel.LOW ) - for role in ( - GeminiRole.CHART_CONFIG, - GeminiRole.CHART_VALIDATION, - GeminiRole.FOLLOW_UP, - GeminiRole.COMPACTION, - ): - assert config.resolve_role_config(role) == replace( - defaults.resolve_role_config(role), model="gemini-custom" - ) @pytest.mark.parametrize( @@ -476,14 +530,18 @@ def test_the_classify_level_is_raised_to_what_the_model_accepts( def test_the_clamp_never_lowers_a_level() -> None: # Test: The thinking-level clamp on roles already at or above the floor. - # Situation: `gemini.discovery_model` accepts `MINIMAL`, - # `thinking.discovery_level` is `HIGH`, and `thinking.synthesis_level` is - # `LOW`. + # Situation: `gemini.discovery_model`, `gemini.synthesis_model`, and + # `gemini.compaction_model` accept `MINIMAL`, `thinking.discovery_level` + # is `HIGH`, and `thinking.synthesis_level` is `LOW`. # Expectation: The configured levels and the fixed `LOW` of the # lightweight roles are unchanged; the clamp only raises a level that # is below the model's floor. config = AgentConfig( - gemini=GeminiConfig(discovery_model="gemini-3.6-flash"), + gemini=GeminiConfig( + discovery_model="gemini-3.6-flash", + synthesis_model="gemini-3.6-flash", + compaction_model="gemini-3.6-flash", + ), thinking=ThinkingLevels( discovery_level=types.ThinkingLevel.HIGH, synthesis_level=types.ThinkingLevel.LOW, diff --git a/narratives/agent/src/narratives_agent/workflows/chart_config.py b/narratives/agent/src/narratives_agent/workflows/chart_config.py index d1282b89..38cb3144 100644 --- a/narratives/agent/src/narratives_agent/workflows/chart_config.py +++ b/narratives/agent/src/narratives_agent/workflows/chart_config.py @@ -54,7 +54,7 @@ async def get_chart_config( mcp_results: The gathered tool results to chart. user_message: The user's query. config: The agent configuration; the request runs on - `gemini.discovery_model`. + `gemini.chart_config_model`. token_usage: Optional turn totals that receive the call's token counts. @@ -168,7 +168,7 @@ async def validate_data_response( synthesis_text: The streamed answer to judge. user_message: The user's query. config: The agent configuration; the request runs on - `gemini.discovery_model`. + `gemini.chart_validation_model`. token_usage: Optional turn totals that receive the call's token counts. """ diff --git a/narratives/agent/src/narratives_agent/workflows/follow_up.py b/narratives/agent/src/narratives_agent/workflows/follow_up.py index e33f877c..13b07e87 100644 --- a/narratives/agent/src/narratives_agent/workflows/follow_up.py +++ b/narratives/agent/src/narratives_agent/workflows/follow_up.py @@ -54,7 +54,7 @@ async def generate_follow_up_questions( user_message: The user's query. topics: The chart titles that ground the questions. config: The agent configuration; `prompts.follow_up` is the system - instruction and `gemini.discovery_model` the model. + instruction and `gemini.follow_up_model` the model. token_usage: Optional turn totals that receive the call's token counts. """ diff --git a/narratives/agent/src/narratives_agent/workflows/transcript.py b/narratives/agent/src/narratives_agent/workflows/transcript.py index 85b64eba..f900d800 100644 --- a/narratives/agent/src/narratives_agent/workflows/transcript.py +++ b/narratives/agent/src/narratives_agent/workflows/transcript.py @@ -508,7 +508,7 @@ async def finalize_turn( model_response: The complete answer text streamed to the browser. state_slots: The scopes extracted from the turn's tool calls. config: The agent configuration; compaction runs on - `gemini.discovery_model`. + `gemini.compaction_model`. token_usage: Optional turn totals that receive the compaction call's token counts. diff --git a/narratives/docs/deployment.md b/narratives/docs/deployment.md index efcbca0b..795b71b0 100644 --- a/narratives/docs/deployment.md +++ b/narratives/docs/deployment.md @@ -269,11 +269,13 @@ bucket-side failure never leaves a phase without its system instruction. `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. The classify -phase runs on the model configured in `gemini.classify_model` in -`agent-config.json` (which defaults to `gemini-3.5-flash-lite`), while the tool -loop, synthesis, chart selection, and follow-up generation use -`gemini.discovery_model`. +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`). Prompts support runtime template placeholders: diff --git a/narratives/schemas/agent-config.example.json b/narratives/schemas/agent-config.example.json index a55bc55d..4cc03b1a 100644 --- a/narratives/schemas/agent-config.example.json +++ b/narratives/schemas/agent-config.example.json @@ -5,8 +5,13 @@ "gemini": { "api_key": "REPLACE_ME_OR_USE_SECRET_MANAGER", "api_base": "https://generativelanguage.googleapis.com/v1beta/models", + "classify_model": "gemini-3.5-flash-lite", "discovery_model": "gemini-3.8-flash", - "classify_model": "gemini-3.5-flash-lite" + "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" }, "mcp": { diff --git a/narratives/schemas/agent-config.schema.json b/narratives/schemas/agent-config.schema.json index b200f11b..578c1115 100644 --- a/narratives/schemas/agent-config.schema.json +++ b/narratives/schemas/agent-config.schema.json @@ -21,15 +21,40 @@ "format": "uri", "default": "https://generativelanguage.googleapis.com/v1beta/models" }, + "classify_model": { + "description": "Model used for the classify phase. The classifier runs at minimal thinking when the model name (its final path segment, for a resource path) starts with one of the prefixes in _MINIMAL_THINKING_MODEL_PREFIXES in config_model.py, and at low thinking otherwise.", + "type": "string", + "default": "gemini-3.5-flash-lite" + }, "discovery_model": { "description": "Model used for the discovery phase, where entities and statistical variables are resolved from the MCP server.", "type": "string", "default": "gemini-3.8-flash" }, - "classify_model": { - "description": "Model used for the classify phase. The classifier runs at minimal thinking when the model name (its final path segment, for a resource path) starts with one of the prefixes in _MINIMAL_THINKING_MODEL_PREFIXES in config_model.py, and at low thinking otherwise.", + "synthesis_model": { + "description": "Model used for the synthesis phase, where the narrative answer is streamed from the gathered data.", "type": "string", - "default": "gemini-3.5-flash-lite" + "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.", + "type": "string", + "default": "gemini-3.8-flash" + }, + "compaction_model": { + "description": "Model used for transcript compaction, which summarizes evicted turns when the conversation window exceeds its limit.", + "type": "string", + "default": "gemini-3.8-flash" } } },