Algolia search for docs (#775)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

Author: petyosi

.github/workflows/main.yml

@@ -52,6 +52,8 @@ jobs:
           else
             uv run --no-sync mkdocs build
           fi
+        env:
+          ALGOLIA_WRITE_API_KEY: ${{ secrets.ALGOLIA_WRITE_API_KEY }}
 
   test:
     name: test on Python ${{ matrix.python-version }} and pydantic ${{ matrix.pydantic-version }}

docs/extra/tweaks.css

@@ -75,3 +75,7 @@ li.md-nav__item>a[href^="#logfire.configure("] {
   text-transform: none;
   font-size: 1.1em;
 }
+
+.md-search__output em {
+  color: var(--md-primary-fg-color);
+}

docs/javascripts/search-worker.js

@@ -0,0 +1,69 @@
+importScripts('https://cdn.jsdelivr.net/npm/algoliasearch@5.18.0/dist/algoliasearch.umd.min.js')
+
+const SETUP = 0
+const READY = 1
+const QUERY = 2
+const RESULT = 3
+
+
+const appID = 'KPPUDTIAVX';
+const apiKey = '1fc841595212a2c3afe8c24dd4cb8790';
+const indexName = 'logfire-docs';
+
+const client = algoliasearch.algoliasearch(appID, apiKey);
+
+self.onmessage = async (event) => {
+  if (event.data.type === SETUP) {
+    self.postMessage({ type: READY });
+  } else if (event.data.type === QUERY) {
+
+    const query = event.data.data
+
+    if (query === '') {
+      self.postMessage({
+        type: RESULT, data: {
+          items: []
+        }
+      });
+      return
+    }
+
+    const { results } = await client.search({
+      requests: [
+        {
+          indexName,
+          query,
+        },
+      ],
+    });
+
+    const hits = results[0].hits
+
+    // make navigation work with preview deployments
+    const stripDocsPathName = !(new URL(self.location.href).pathname.startsWith('/docs'));
+
+    const mappedGroupedResults = hits.reduce((acc, hit) => {
+      if (!acc[hit.pageID]) {
+        acc[hit.pageID] = []
+      }
+      acc[hit.pageID].push({
+        score: 1,
+        terms: {},
+        location: stripDocsPathName ? hit.abs_url.replace('/docs', '') : hit.abs_url,
+        title: hit.title,
+        text: hit._highlightResult.content.value,
+
+      })
+      return acc
+    }, {})
+
+
+
+
+    self.postMessage({
+      type: RESULT, data: {
+        items: Object.values(mappedGroupedResults)
+      }
+    });
+  }
+};

docs/overrides/main.html

@@ -1,5 +1,50 @@
 {% extends "base.html" %}
 
+{% block config %}
+      {%- set app = {
+        "base": base_url,
+        "features": features,
+        "translations": {},
+        "search": base_url + "/javascripts/search-worker.js" | url
+      } -%}
+
+      <!-- Versioning -->
+      {%- if config.extra.version -%}
+        {%- set mike = config.plugins.get("mike") -%}
+        {%- if not mike or mike.config.version_selector -%}
+          {%- set _ = app.update({ "version": config.extra.version }) -%}
+        {%- endif -%}
+      {%- endif -%}
+
+      <!-- Tags -->
+      {%- if config.extra.tags -%}
+        {%- set _ = app.update({ "tags": config.extra.tags }) -%}
+      {%- endif -%}
+
+      <!-- Translations -->
+      {%- set translations = app.translations -%}
+      {%- for key in [
+        "clipboard.copy",
+        "clipboard.copied",
+        "search.result.placeholder",
+        "search.result.none",
+        "search.result.one",
+        "search.result.other",
+        "search.result.more.one",
+        "search.result.more.other",
+        "search.result.term.missing",
+        "select.version"
+      ] -%}
+        {%- set _ = translations.update({ key: lang.t(key) }) -%}
+      {%- endfor -%}
+
+      <!-- Configuration -->
+      <script id="__config" type="application/json">
+        {{- app | tojson -}}
+      </script>
+  <!-- Add scripts that need to run afterwards here -->
+{% endblock %}
+
 {% block content %}
 {{ super() }}
 <script src="/flarelytics/client.js"></script>

docs/plugins/build_index.py

@@ -0,0 +1,71 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import Any, cast
+
+from algoliasearch.search_client import SearchClient
+from bs4 import BeautifulSoup
+from mkdocs.config import Config
+from mkdocs.structure.files import Files
+from mkdocs.structure.pages import Page
+
+records: list[dict[str, Any]] = []
+ALGOLIA_INDEX_NAME = 'logfire-docs'
+ALGOLIA_APP_ID = 'KPPUDTIAVX'
+ALGOLIA_WRITE_API_KEY = os.environ.get('ALGOLIA_WRITE_API_KEY')
+
+
+def on_page_content(html: str, page: Page, config: Config, files: Files) -> str:
+    if not ALGOLIA_WRITE_API_KEY:
+        return html
+
+    assert page.title is not None, 'Page title must not be None'  # type: ignore[reportUnknownMemberType]
+    title = cast(str, page.title)  # type: ignore[reportUnknownMemberType]
+
+    soup = BeautifulSoup(html, 'html.parser')
+
+    # Find all h1 and h2 headings
+    headings = soup.find_all(['h1', 'h2'])
+
+    # Process each section
+    for i in range(len(headings)):
+        current_heading = headings[i]
+        heading_id = current_heading.get('id', '')
+        section_title = current_heading.get_text().replace('¶', '').replace('dataclass', '').strip()
+
+        # Get content until next heading
+        content: list[str] = []
+        sibling = current_heading.find_next_sibling()
+        while sibling and sibling.name not in ['h1', 'h2']:
+            content.append(str(sibling))
+            sibling = sibling.find_next_sibling()
+
+        section_html = ''.join(content)
+
+        # Create anchor URL
+        anchor_url = f'{page.abs_url}#{heading_id}' if heading_id else page.abs_url
+
+        # Create record for this section
+        records.append(
+            {
+                'content': section_html,
+                'pageID': title,
+                'abs_url': anchor_url,
+                'title': f'{title} - {section_title}',
+                'objectID': anchor_url,
+            }
+        )
+
+    return html
+
+
+def on_post_build(config: Config) -> None:
+    if not ALGOLIA_WRITE_API_KEY:
+        return
+
+    client = SearchClient.create(ALGOLIA_APP_ID, ALGOLIA_WRITE_API_KEY)
+    index = client.init_index(ALGOLIA_INDEX_NAME)
+    # temporary filter the records from the index if the content is bigger than 10k characters
+    filtered_records = list(filter(lambda record: len(record['content']) < 9000, records))
+    print(f'Uploading {len(filtered_records)} out of {len(records)} records to Algolia...')
+    index.replace_all_objects(filtered_records, {'createIfNotExists': True}).wait()  # type: ignore[reportUnknownMemberType]

mkdocs.yml

@@ -308,3 +308,4 @@ plugins:
         "get-started/traces.md": "concepts.md"
 hooks:
   - docs/plugins/main.py
+  - docs/plugins/build_index.py

pyproject.toml

@@ -172,6 +172,8 @@ docs = [
     "mkdocstrings-python>=1.8.0",
     "mkdocs-redirects>=1.2.1",
     "griffe",
+    "bs4>=0.0.2",
+    "algoliasearch>=3,<4",
 ]
 
 [tool.inline-snapshot]