|
| 1 | +"""Generate llms.txt, llms-full.txt, and per-page markdown (https://llmstxt.org/). |
| 2 | +
|
| 3 | +The hook publishes three artifacts into the built site: |
| 4 | +
|
| 5 | +- `llms.txt`: a markdown index of the documentation, one link per page, |
| 6 | + grouped by nav section. |
| 7 | +- a `.md` rendition of every prose page next to its HTML (e.g. |
| 8 | + `tutorial/tools/index.md`), which is what the llms.txt links point at. |
| 9 | +- `llms-full.txt`: every prose page concatenated for single-fetch consumption. |
| 10 | +
|
| 11 | +Page markdown is the source markdown with `--8<--` snippet includes resolved |
| 12 | +(so the `docs_src/` code examples appear inline) and relative links rewritten |
| 13 | +to absolute URLs. The API reference pages under `api/` are mkdocstrings stubs |
| 14 | +with no markdown source, so they are linked as rendered HTML from an Optional |
| 15 | +section instead of being embedded. |
| 16 | +
|
| 17 | +Incremental builds (`mkdocs build --dirty`) are rejected: they skip unmodified |
| 18 | +pages, which would silently truncate the generated artifacts. |
| 19 | +""" |
| 20 | + |
| 21 | +from __future__ import annotations |
| 22 | + |
| 23 | +import posixpath |
| 24 | +import re |
| 25 | +from pathlib import Path |
| 26 | + |
| 27 | +from mkdocs.config.defaults import MkDocsConfig |
| 28 | +from mkdocs.exceptions import PluginError |
| 29 | +from mkdocs.structure.files import File, Files |
| 30 | +from mkdocs.structure.nav import Navigation, Section |
| 31 | +from mkdocs.structure.pages import Page |
| 32 | + |
| 33 | +# Pages with no markdown source, linked as HTML under "## Optional". |
| 34 | +_OPTIONAL_PAGES = [ |
| 35 | + ("api/mcp/index.md", "mcp API reference", "Auto-generated API reference for the mcp package (rendered HTML)"), |
| 36 | + ( |
| 37 | + "api/mcp_types/index.md", |
| 38 | + "mcp-types API reference", |
| 39 | + "Auto-generated API reference for the mcp-types package (rendered HTML)", |
| 40 | + ), |
| 41 | +] |
| 42 | + |
| 43 | +_SNIPPET_LINE = re.compile(r'^(?P<indent>[ \t]*)--8<-- "(?P<path>[^"\n]+)"$', flags=re.MULTILINE) |
| 44 | +_MD_LINK = re.compile(r'(\]\()([^)\s]+\.md)(#[^)\s]*)?( +"[^"]*")?(\))') |
| 45 | + |
| 46 | +_page_markdown: dict[str, str] = {} |
| 47 | +_rendition_uris: set[str] = set() |
| 48 | +_nav: Navigation | None = None |
| 49 | +_files: Files | None = None |
| 50 | + |
| 51 | + |
| 52 | +def _site_url(config: MkDocsConfig) -> str: |
| 53 | + assert config.site_url is not None |
| 54 | + return config.site_url.rstrip("/") + "/" |
| 55 | + |
| 56 | + |
| 57 | +def _md_uri(file: File) -> str: |
| 58 | + return re.sub(r"\.html$", ".md", file.dest_uri) |
| 59 | + |
| 60 | + |
| 61 | +def on_config(config: MkDocsConfig) -> None: |
| 62 | + # `mkdocs serve` rebuilds reuse the imported module; start each build clean. |
| 63 | + global _nav, _files |
| 64 | + _page_markdown.clear() |
| 65 | + _rendition_uris.clear() |
| 66 | + _nav = _files = None |
| 67 | + |
| 68 | + |
| 69 | +def on_nav(nav: Navigation, config: MkDocsConfig, files: Files) -> None: |
| 70 | + global _nav, _files |
| 71 | + _nav = nav |
| 72 | + _files = files |
| 73 | + _rendition_uris.update(page.file.src_uri for page in nav.pages if not page.file.src_uri.startswith("api/")) |
| 74 | + |
| 75 | + |
| 76 | +def on_page_markdown(markdown: str, page: Page, config: MkDocsConfig, files: Files) -> str | None: |
| 77 | + if page.file.src_uri not in _rendition_uris: |
| 78 | + return None |
| 79 | + |
| 80 | + # Same anchor as the pymdownx.snippets `base_path` in mkdocs.yml. |
| 81 | + repo_root = Path(config.config_file_path).parent |
| 82 | + |
| 83 | + def include(match: re.Match[str]) -> str: |
| 84 | + indent, path = match["indent"], match["path"] |
| 85 | + try: |
| 86 | + content = (repo_root / path).read_text(encoding="utf-8").rstrip("\n") |
| 87 | + except OSError as exc: |
| 88 | + raise PluginError(f"llms_txt: cannot read snippet {path!r} in {page.file.src_uri}") from exc |
| 89 | + if indent: |
| 90 | + content = "\n".join(indent + line if line else line for line in content.split("\n")) |
| 91 | + return content |
| 92 | + |
| 93 | + resolved, substitutions = _SNIPPET_LINE.subn(include, markdown) |
| 94 | + if substitutions != sum("--8<--" in line for line in markdown.splitlines()): |
| 95 | + raise PluginError(f"llms_txt: unresolved snippet include in {page.file.src_uri}") |
| 96 | + |
| 97 | + site_url = _site_url(config) |
| 98 | + src_dir = posixpath.dirname(page.file.src_uri) |
| 99 | + |
| 100 | + def rewrite(match: re.Match[str]) -> str: |
| 101 | + opening, target, anchor, title, closing = match.groups() |
| 102 | + if "://" in target: |
| 103 | + return match.group(0) |
| 104 | + linked = files.get_file_from_path(posixpath.normpath(posixpath.join(src_dir, target))) |
| 105 | + if linked is None: |
| 106 | + return match.group(0) |
| 107 | + # Pages without a markdown rendition (the api/ stubs) link to their HTML instead. |
| 108 | + url = _md_uri(linked) if linked.src_uri in _rendition_uris else linked.url |
| 109 | + return f"{opening}{site_url}{url}{anchor or ''}{title or ''}{closing}" |
| 110 | + |
| 111 | + _page_markdown[page.file.src_uri] = _MD_LINK.sub(rewrite, resolved) |
| 112 | + return None |
| 113 | + |
| 114 | + |
| 115 | +def _section_pages(section: Section) -> list[Page]: |
| 116 | + pages: list[Page] = [] |
| 117 | + for child in section.children: |
| 118 | + if isinstance(child, Page) and child.file.src_uri in _rendition_uris: |
| 119 | + pages.append(child) |
| 120 | + elif isinstance(child, Section): |
| 121 | + pages.extend(_section_pages(child)) |
| 122 | + return pages |
| 123 | + |
| 124 | + |
| 125 | +def on_post_build(config: MkDocsConfig) -> None: |
| 126 | + assert _nav is not None and _files is not None |
| 127 | + missing = _rendition_uris - _page_markdown.keys() |
| 128 | + if missing: |
| 129 | + raise PluginError(f"llms_txt: pages skipped this build (is this a --dirty build?): {sorted(missing)}") |
| 130 | + |
| 131 | + site_dir = Path(config.site_dir) |
| 132 | + site_url = _site_url(config) |
| 133 | + |
| 134 | + top_level = [item for item in _nav.items if isinstance(item, Page) and item.file.src_uri in _rendition_uris] |
| 135 | + sections: list[tuple[str, list[Page]]] = [("Docs", top_level)] if top_level else [] |
| 136 | + for item in _nav.items: |
| 137 | + if isinstance(item, Section): |
| 138 | + pages = _section_pages(item) |
| 139 | + if pages: |
| 140 | + sections.append((item.title, pages)) |
| 141 | + |
| 142 | + index = [f"# {config.site_name}", "", f"> {config.site_description}", ""] |
| 143 | + full: list[str] = [] |
| 144 | + for title, pages in sections: |
| 145 | + index += [f"## {title}", ""] |
| 146 | + for page in pages: |
| 147 | + markdown = _page_markdown[page.file.src_uri] |
| 148 | + (site_dir / _md_uri(page.file)).write_text(markdown, encoding="utf-8") |
| 149 | + |
| 150 | + description = page.meta.get("description") |
| 151 | + tail = f": {description}" if description else "" |
| 152 | + index.append(f"- [{page.title}]({site_url}{_md_uri(page.file)}){tail}") |
| 153 | + |
| 154 | + body = re.sub(r"\A\s*# .+\n", "", markdown) |
| 155 | + full += [f"# {page.title}", "", f"Source: {page.canonical_url}", "", body.strip(), ""] |
| 156 | + index.append("") |
| 157 | + |
| 158 | + index += ["## Optional", ""] |
| 159 | + for src_uri, title, description in _OPTIONAL_PAGES: |
| 160 | + linked = _files.get_file_from_path(src_uri) |
| 161 | + if linked is None: |
| 162 | + raise PluginError(f"llms_txt: optional page {src_uri} not found") |
| 163 | + index.append(f"- [{title}]({site_url}{linked.url}): {description}") |
| 164 | + index.append("") |
| 165 | + |
| 166 | + (site_dir / "llms.txt").write_text("\n".join(index), encoding="utf-8") |
| 167 | + (site_dir / "llms-full.txt").write_text("\n".join(full), encoding="utf-8") |
0 commit comments