Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/notebooklm/_artifacts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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: Output language code.
instructions: Optional generation instructions.

Returns:
Dictionary with 'mind_map' (JSON data) and 'note_id'.
Expand All @@ -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]],
]
Expand Down
22 changes: 19 additions & 3 deletions src/notebooklm/cli/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -987,38 +987,54 @@ async def _generate():


@generate.command("mind-map")
@click.argument("description", default="", required=False)
@click.option(
"-n",
"--notebook",
"notebook_id",
default=None,
help="Notebook ID (uses current if not set)",
)
@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, description, notebook_id, 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 chronology" --language zh_Hans
"""
nb_id = require_notebook(notebook_id)

async def _run():
async with NotebookLMClient(client_auth) as client:
nb_id_resolved = await resolve_notebook_id(client, nb_id)
sources = await resolve_source_ids(client, nb_id_resolved, source_ids)
resolved_language = resolve_language(language)
instructions = description or None

# 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=resolved_language,
instructions=instructions,
)
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=resolved_language,
instructions=instructions,
)

_output_mind_map_result(result, json_output)
Expand Down
30 changes: 30 additions & 0 deletions tests/unit/cli/test_generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -432,6 +432,36 @@ def test_generate_mind_map(self, runner, mock_auth):

assert result.exit_code == 0

def test_generate_mind_map_with_language(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",
"zh_Hans",
"-n",
"nb_123",
],
)

assert result.exit_code == 0
mock_client.artifacts.generate_mind_map.assert_awaited_once_with(
"nb_123",
source_ids=None,
language="zh_Hans",
instructions=None,
)


# =============================================================================
# GENERATE REPORT TESTS
Expand Down
24 changes: 24 additions & 0 deletions tests/unit/test_source_selection.py
Original file line number Diff line number Diff line change
Expand Up @@ -509,6 +509,30 @@ 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_includes_language_and_instructions(
self, mock_core, mock_notes_api
):
"""Test generate_mind_map encodes language and instructions in config."""
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_Hans",
instructions="Focus on chronology",
)

params = mock_core.rpc_call.call_args.args[1]

assert params[5] == [
"interactive_mindmap",
[["[CONTEXT]", "Focus on chronology"]],
"zh_Hans",
]

@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."""
Expand Down
Loading