diff --git a/src/notebooklm/_artifacts.py b/src/notebooklm/_artifacts.py index e0ab8eb2..675eebb1 100644 --- a/src/notebooklm/_artifacts.py +++ b/src/notebooklm/_artifacts.py @@ -988,6 +988,8 @@ async def generate_mind_map( self, notebook_id: str, source_ids: builtins.list[str] | None = None, + language: str = "en", + instructions: str | None = None, ) -> dict[str, Any]: """Generate an interactive mind map. @@ -997,6 +999,8 @@ async def generate_mind_map( Args: notebook_id: The notebook ID. source_ids: Source IDs to include. If None, uses all sources. + language: Language code (default: "en"). + instructions: Custom instructions for the generated mind map. Returns: Dictionary with 'mind_map' (JSON data) and 'note_id'. @@ -1014,7 +1018,7 @@ async def generate_mind_map( None, None, None, - ["interactive_mindmap", [["[CONTEXT]", ""]], ""], + ["interactive_mindmap", [["[CONTEXT]", instructions or ""]], language], None, [2, None, [1]], ] diff --git a/src/notebooklm/cli/generate.py b/src/notebooklm/cli/generate.py index 210c5a26..f8f74dab 100644 --- a/src/notebooklm/cli/generate.py +++ b/src/notebooklm/cli/generate.py @@ -994,14 +994,21 @@ async def _generate(): default=None, help="Notebook ID (uses current if not set)", ) +@click.argument("description", default="", required=False) +@click.option("--language", default=None, help="Output language (default: from config or 'en')") @click.option("--source", "-s", "source_ids", multiple=True, help="Limit to specific source IDs") @json_option @with_client -def generate_mind_map(ctx, notebook_id, source_ids, json_output, client_auth): +def generate_mind_map(ctx, notebook_id, description, language, source_ids, json_output, client_auth): """Generate mind map. \b Use --json for machine-readable output. + + \b + Example: + notebooklm generate mind-map "focus on causal relationships" + notebooklm generate mind-map --language zh-CN """ nb_id = require_notebook(notebook_id) @@ -1013,12 +1020,18 @@ async def _run(): # Show status spinner only for console output if json_output: result = await client.artifacts.generate_mind_map( - nb_id_resolved, source_ids=sources + nb_id_resolved, + source_ids=sources, + language=resolve_language(language), + instructions=description or None, ) else: with console.status("Generating mind map..."): result = await client.artifacts.generate_mind_map( - nb_id_resolved, source_ids=sources + nb_id_resolved, + source_ids=sources, + language=resolve_language(language), + instructions=description or None, ) _output_mind_map_result(result, json_output) diff --git a/tests/unit/cli/test_generate.py b/tests/unit/cli/test_generate.py index bb64ec7a..884dab7c 100644 --- a/tests/unit/cli/test_generate.py +++ b/tests/unit/cli/test_generate.py @@ -432,6 +432,35 @@ def test_generate_mind_map(self, runner, mock_auth): assert result.exit_code == 0 + def test_generate_mind_map_with_language_and_instructions(self, runner, mock_auth): + with patch_client_for_module("generate") as mock_client_cls: + mock_client = create_mock_client() + mock_client.artifacts.generate_mind_map = AsyncMock( + return_value={"mind_map": {"name": "Root", "children": []}, "note_id": "n1"} + ) + mock_client_cls.return_value = mock_client + + with patch("notebooklm.cli.helpers.fetch_tokens", new_callable=AsyncMock) as mock_fetch: + mock_fetch.return_value = ("csrf", "session") + result = runner.invoke( + cli, + [ + "generate", + "mind-map", + "--language", + "ja", + "-n", + "nb_123", + "focus on causal chains", + ], + ) + + assert result.exit_code == 0 + mock_client.artifacts.generate_mind_map.assert_awaited_once() + kwargs = mock_client.artifacts.generate_mind_map.await_args.kwargs + assert kwargs["language"] == "ja" + assert kwargs["instructions"] == "focus on causal chains" + # ============================================================================= # GENERATE REPORT TESTS diff --git a/tests/unit/test_source_selection.py b/tests/unit/test_source_selection.py index f58a9759..e8470727 100644 --- a/tests/unit/test_source_selection.py +++ b/tests/unit/test_source_selection.py @@ -509,6 +509,29 @@ async def test_generate_mind_map_source_encoding(self, mock_core, mock_notes_api assert source_ids_nested == [[["src_mm_1"]], [["src_mm_2"]]] + @pytest.mark.asyncio + async def test_generate_mind_map_language_and_instructions(self, mock_core, mock_notes_api): + """Test generate_mind_map forwards language and instructions into the RPC payload.""" + api = ArtifactsAPI(mock_core, mock_notes_api) + + mock_core.rpc_call.return_value = [['{"name": "Mind Map", "children": []}']] + + await api.generate_mind_map( + notebook_id="nb_123", + source_ids=["src_mm_1"], + language="zh-CN", + instructions="Focus on decision tradeoffs", + ) + + params = mock_core.rpc_call.call_args.args[1] + config = params[5] + + assert config == [ + "interactive_mindmap", + [["[CONTEXT]", "Focus on decision tradeoffs"]], + "zh-CN", + ] + @pytest.mark.asyncio async def test_suggest_reports_uses_get_suggested_reports(self, mock_core, mock_notes_api): """Test suggest_reports uses GET_SUGGESTED_REPORTS RPC."""