omarsol's picture
Archive private KB bundles for fast cold starts
c8c398a
|
Raw History Blame
19.7 kB

Python scripts for adding new courses and updating documentation

First workflow: Adding a New Course

Make sure you have the required environment variables set:

  • GEMINI_API_KEY or GOOGLE_API_KEY for context generation with Gemini
  • COHERE_API_KEY for embeddings
  • HF_TOKEN for HuggingFace uploads and downloads - access to the private HuggingFace dataset repo
  • GITHUB_TOKEN for accessing files via the GitHub API (used by the docs workflow only; the course workflow never calls GitHub)

Optional Gemini context-generation tuning:

  • GEMINI_CONTEXT_TPM_LIMIT - input token-per-minute quota to throttle against (defaults to 30000000)
  • GEMINI_CONTEXT_TPM_SAFETY_MARGIN - fraction of that quota to use before pausing (defaults to 0.8)
  • GEMINI_CONTEXT_CONCURRENCY - max concurrent context requests (defaults to 50)
  • GEMINI_CONTEXT_RETRY_ATTEMPTS - max tenacity attempts for transient Gemini API errors (defaults to 8)
  • GEMINI_CONTEXT_MODEL - Gemini model used for context generation (defaults to gemini-3.1-flash-lite)
  • GEMINI_CONTEXT_TPM_WINDOW_SECONDS - sliding-window length in seconds for the TPM throttle (defaults to 60)

1. Prepare the course data

  1. Make sure you have access to:

  2. In Notion, navigate to the main course page that contains the live lessons:

  3. Click on the three dots in the top right corner and select "Export"

  4. Select these options:

    • Export format: "Markdown & CSV"
    • Include databases: current view
    • Include content: Everything
    • Include subpages: yes
    • Create folders for subpages: yes
  5. Click on "Export"

  6. Once the export is complete, unzip file.

  7. Move the unzipped folder into the ai-tutor-app/data directory

  8. Rename the folder to the course name.

    • e.g. master_ai_for_work
  9. Open data/scraping_scripts/source_registry.py.

  10. Add the new course to the SOURCE_CONFIGS dictionary. Sources listed in this registry are active in the knowledge base. Also add the key to COURSE_SOURCE_KEYS in the same file: a course missing from that tuple is classified as documentation (its pages land under raw/docs/ and wiki/frameworks/ with the wrong picker group) and test_registry_classifies_every_source fails CI.

    example:

       "master_ai_for_work": {
          "base_url": "",
          "input_directory": "data/master_ai_for_work",  # Relative path to the directory that contains the Markdown files
          "output_file": "data/master_ai_for_work_data.jsonl", # The output file that will be created by the script
          "source_name": "master_ai_for_work",
          "use_include_list": False,
          "included_dirs": [],
          "excluded_dirs": [],
          "excluded_root_files": [],
          "included_root_files": [],
          "url_extension": "",
       },
    
    • The most important fields are:
      • input_directory: the relative path to the directory that contains the Markdown files
      • output_file: the name of the output file that will be created by the script
      • source_name: the name of the course (keep underscores, no spaces)
    • The other fields can stay as empty lists or empty strings.

2. Run the add_course_workflow.py script

uv run -m data.scraping_scripts.add_course_workflow --courses [COURSE_NAME]

example:

uv run -m data.scraping_scripts.add_course_workflow --courses master_ai_for_work

This script will guide you through the complete process, it will:

  1. Download the JSONL files from the other courses
  2. Extract the markdown content from each of the lessons and create a new JSONL file for the course
  3. Prompt you to manually add URLs to the course content, inside the newly created JSONL file (more details in step 3 below)
  4. Merge the course data into the main dataset
  5. Rebuild the KB artifacts (raw/generated/wiki; skip with --skip-kb)
  6. Add contextual information to document chunks before embedding (then prune contextual nodes to active sources)
  7. Create vector stores
  8. Upload databases to HuggingFace
  9. Confirm the course is configured in the central source registry

3. Add URLs to the course content + Manual Dataset Cleaning (Most important step)

  • After the script has processed the markdown files, it will prompt you to manually add URLs to the course content.

  • Answer "no" to the question "Have you added all the URLs?"

  • Open the newly created data/master_ai_for_work_data.jsonl file in a text editor and open the live course page in the browser.

  • If the JSON looks split into multiple wrapped lines in VS Code / Cursor, you can toggle word wrap off.

    • macOS: press βŒ₯ Option + Z
    • Windows/Linux: press Alt + Z
  • For each lesson in the course, starting at the beginning academy.towardsai.net, copy the URL, and add it to the url field in the JSONL file, in the corresponding line/lesson. Read the note below to know what urls to add.

Note: While you do this, now is the time to clean up the .jsonl file, remove any lines/lessons that should not be added to the RAG chatbot. example: "Course Admin and Syllabus", "Course Structure/Overview", "Course Outline", "Quiz", "Assigments", "Introduction to Module X" etc. You can also remove lines that are videos. Only actual lessons should be kept in the JSONL file, with the url field filled in.

  • What you can do is add the URLs for all the lessons you want to add to the RAG chatbot, and when done, remove all the json lines that have an empty url field.

4. Once done, run the script again and answer "yes" to the question "Have you added all the URLs?"

uv run -m data.scraping_scripts.add_course_workflow --courses master_ai_for_work

5. Its done

Run the chatbot locally to test if the course has been added correctly.

uv run -m app.api

Second workflow: Updating Documentation (GitHub API / llms.txt)

To update library documentation from GitHub repositories or llms.txt indexes:

uv run -m data.scraping_scripts.update_docs_workflow

Rebuilding the KB from scratch

If you've blown away local KB state (e.g. rm -rf data/kb data/chroma-db-all_sources) and want everything back, just run the docs workflow β€” it now handles fresh-build cases automatically.

uv run -m data.scraping_scripts.update_docs_workflow

What happens under the hood when local state is missing:

  1. ensure_required_files_exist downloads all_sources_data.jsonl, all_sources_contextual_nodes.pkl, and every per-source JSONL (including courses) from the private HF dataset towardsai-tutors/ai-tutor-data.
  2. Fresh docs are re-fetched from GitHub / llms.txt indexes.
  3. build_kb_artifacts.py regenerates data/kb/raw/ and data/kb/generated/ from the JSONL. Course raw pages are reconstructed from the JSONL too β€” you do not need the original Notion exports in data/<course_name>/.
  4. update_kb_wiki.py rebuilds data/kb/wiki/. When it detects an empty wiki/ it auto-promotes to a full seed (topic pages, recipes/errors index pages), so a clean rebuild produces the full wiki. On incremental runs, maintainer-authored prose outside <!-- AUTO-GENERATED --> markers is preserved; pass --seed-defaults to wipe and reseed everything.
  5. data/kb/AGENTS.md is overwritten from data/scraping_scripts/kb_agents_template.md (the canonical template, tracked in git so it survives any rm -rf data/).
  6. Chroma vector store is rebuilt by create_vector_stores.py.

Prerequisites for a from-scratch rebuild:

  • HF_TOKEN with access to towardsai-tutors/ai-tutor-data (so course JSONLs and the contextual nodes PKL can be downloaded).
  • GITHUB_TOKEN (rate-limit headroom for docs downloads).
  • GEMINI_API_KEY or GOOGLE_API_KEY (used by add_context_to_nodes.py).
  • COHERE_API_KEY (used by create_vector_stores.py).

What this workflow cannot restore: the original Notion-exported course markdown in data/<course_name>/. You don't need it for the runtime β€” the chunked content lives in the JSONL on HF β€” but if you want to re-run add_course_workflow.py against fresh course content, you'd need a new Notion export.

Editing KB agent guidance permanently

Do not edit data/kb/AGENTS.md directly. The file is overwritten in two places, both reading from the canonical template at data/scraping_scripts/kb_agents_template.md:

  • data.scraping_scripts.update_kb_wiki.write_agents_md β€” rewrites it during every update_docs_workflow.py run, before uploading to HuggingFace.
  • app.config.ensure_kb_agents_md β€” rewrites it on every runtime startup, after ensure_local_vector_db downloads the snapshot. This catches the case where the HF snapshot has a stale AGENTS.md (e.g. uploaded before a template change landed in git) and ensures the live file always matches the version checked into the current branch.

Edit the template, commit it, and either run the docs workflow (to refresh the HF snapshot) or just restart the chatbot (the template gets copied into place on startup either way).

Where data/kb/ actually lives

data/kb/ is gitignored (it's a 100MB+ build artifact). It's:

  • Built by data.scraping_scripts.build_kb_artifacts (from data/all_sources_data.jsonl) and data.scraping_scripts.update_kb_wiki.
  • Archived and uploaded to towardsai-tutors/ai-tutor-vector-db by data.scraping_scripts.upload_dbs_to_hf as kb.tar.gz. The local unpacked tree is preserved; legacy remote kb/** files are pruned only after the archive upload is verified.
  • Downloaded by app.config.ensure_local_vector_db on the first chatbot start (or any start where the local KB is missing).

Treat it the same way as data/chroma-db-all_sources/: never commit it, never edit it by hand, regenerate or re-download when you need it.

This will update all supported documentation sources. You can also specify specific sources:

uv run -m data.scraping_scripts.update_docs_workflow --sources transformers peft langchain

The workflow includes:

  1. Downloading documentation from GitHub (API) or llms.txt indexes
  2. Capturing per-source library versions (capture_source_versions.py)
  3. Processing markdown files to create JSONL data
  4. Rebuilding the KB artifacts (raw/generated/wiki; skip with --skip-kb)
  5. Adding contextual information to document chunks (then pruning contextual nodes to active sources)
  6. Creating vector stores
  7. Uploading vector db and new JSONL files to HuggingFace

Both workflows validate Hugging Face access up front, and the docs workflow now stops immediately on GitHub auth/API failures so it does not overwrite source JSONL files with incomplete data.

Individual Components

If you need to run specific steps individually:

  • GitHub to Markdown: github_to_markdown_ai_docs.py
  • llms.txt to Markdown: llms_txt_to_markdown_docs.py
  • Capture library versions: capture_source_versions.py
  • Process Markdown: process_md_files.py
  • Build KB artifacts: build_kb_artifacts.py / Update KB wiki: update_kb_wiki.py / Lint KB wiki references: lint_kb_wiki.py
  • Build public docs-only bundle: build_public_docs_bundle.py
  • Add Context: add_context_to_nodes.py
  • Create Vector Stores: create_vector_stores.py
  • Upload to Chroma Vector Store to HuggingFace: upload_dbs_to_hf.py
  • Upload JSONL files to HuggingFace: upload_data_to_hf.py
  • Retire a source from data and Chroma: retire_source_workflow.py

Retiring a Course or Documentation Source

Use the retirement workflow when a source should be removed from retrieval and from the Hugging Face data repositories:

uv run -m data.scraping_scripts.retire_source_workflow --sources 8-hour_primer --yes

The workflow:

  1. Downloads all_sources_data.jsonl and all_sources_contextual_nodes.pkl if they are missing locally.
  2. Removes rows/chunks whose source matches the retired source key.
  3. Removes the retired source from source_registry.py, so it is no longer active in future workflow runs or the UI source picker.
  4. Rebuilds data/chroma-db-all_sources.
  5. Uploads the rebuilt vector DB and updated aggregate data files.
  6. Deletes the retired per-source JSONL from towardsai-tutors/ai-tutor-data.

Run with --dry-run first to preview counts without changing files:

uv run -m data.scraping_scripts.retire_source_workflow --sources 8-hour_primer --dry-run

Use --keep-source-registry only if you want to remove existing chunks while leaving the source configured as active for a future rebuild.

Publishing the public docs-only bundle

The production bundle in towardsai-tutors/ai-tutor-vector-db is private and holds every source, including gated course content, so cold-starting the app needs an HF_TOKEN with read access. To let anyone run the chatbot without that token, we publish a public, docs-only bundle to towardsai-tutors/ai-tutor-vector-db-public (the 9 documentation/reference sources, no courses). The runtime (app.config.ensure_local_vector_db) downloads the full private bundle when a usable HF_TOKEN is present and falls back to the public one otherwise; the picker hides course sources when running on it.

This is a standalone, manual step β€” it is not wired into the docs/course workflows. Refresh the public bundle whenever the docs sources change materially:

# Refresh the prod bundle first (so data/chroma-db-all_sources, data/kb,
# data/all_sources_data.jsonl are current), then:
uv run -m data.scraping_scripts.build_public_docs_bundle              # build + upload
uv run -m data.scraping_scripts.build_public_docs_bundle --skip-upload  # build only, inspect first

It derives the docs-only bundle from the already-built prod artifacts (no re-embedding, $0 Cohere). Filtering is an allowlist: only sources in DOC_SOURCE_KEYS survive; anything unclassified fails closed (dropped, never published), and a registry invariant test keeps the doc/course split total. The script copies the Chroma collection and keeps only allowlisted chunks (scanning metadata, so chunks with missing/unknown source are dropped too), rebuilds the document-dict pkl and the BM25 index (gzipped JSON) from a docs-only view of all_sources_data.jsonl, and copies data/kb while pruning raw/courses, wiki/courses, and the non-doc rows from the generated indexes. The staged wiki is then publicized: the scaffolder regenerates AUTO-GENERATED blocks from the filtered manifest, course references in maintainer prose are pruned, and a (maintainer) entry is appended to wiki/log.md (see data/kb/MAINTAINER.md). Every build ends with a leak audit that fails the build if any course path, course key, or non-allowlisted row survives.

After the audit, the KB tree is packed into a single kb.tar.gz and the unpacked tree is dropped from staging: the public repo ships 6 files instead of ~3,000, which keeps a token-free cold start clear of HF's anonymous per-request rate limits (15 min unpacked vs a few min archived). The runtime extracts the archive after download (app.config._extract_kb_archive).

The private production publisher now uses the same archive layout: it creates data/kb.tar.gz transiently, uploads and verifies it, then prunes the legacy remote kb/** tree. It never removes or rewrites the canonical local data/kb/ directory.

The staged tree lives in data/public_docs_bundle/ (gitignored). Uploading needs an HF_TOKEN with write access to the towardsai-tutors org; the script creates the public repo on first run. Pass --include-contextual to also ship a docs-only all_sources_contextual_nodes.pkl (off by default to keep the download small).

Tips for New Team Members

  1. To update the AI Tutor with new content:

    • For new courses, use add_course_workflow.py
    • For updated documentation, use update_docs_workflow.py
  2. When adding URLs to course content:

    • Get the URLs from the live course platform
    • Add them to the generated JSONL file in the url field
    • Example URL format: https://academy.towardsai.net/courses/take/ai-business-professionals/multimedia/64072092-introduction-to-chatgpt-claude-and-gemini
    • Make sure every document has a valid URL
  3. By default, only new content will have context added to save time and resources. Use --process-all-context only if you need to regenerate context for all documents. Use --skip-data-upload if you don't want to upload data files to the private HuggingFace repo (they're uploaded by default).

  4. When adding a new course, verify that it appears in the UI:

    • Add the source to the registry structures in source_registry.py: COURSE_SOURCE_KEYS (classification; a course absent here is treated as documentation and fails CI), SOURCE_KEY_TO_LABEL, SOURCE_DISPLAY_INFO (what /api/tools renders), UI_SOURCE_KEYS (a source absent here never appears in the picker), and optionally DEFAULT_SELECTED_SOURCE_KEYS
    • Check that the new source appears in the source picker in the UI
    • Make sure it's properly included in the default selected sources if desired
    • Restart the API server to see the changes
  5. First time setup or missing files:

    • Both workflows automatically check for and download required data files:
      • all_sources_data.jsonl - Contains the raw document data
      • all_sources_contextual_nodes.pkl - Contains the context-enriched chunk manifest used for embedding
    • If the PKL file exists, the default behavior only processes new content; use --process-all-context to regenerate context for every document
    • You must have proper HuggingFace credentials with access to the private repository
  6. Vector-store creation:

    • create_vector_stores.py now shows terminal progress for embedding generation and Chroma upserts
    • all_sources_contextual_nodes.pkl stores chunk records for embedding; legacy pickles containing LlamaIndex nodes are still accepted by the runtime compatibility helpers
  7. KB wiki artifact creation:

    • build_kb_artifacts.py converts all_sources_data.jsonl into generated markdown under data/kb/raw/
    • update_kb_wiki.py refreshes wiki navigation pages under data/kb/wiki/, preserving maintainer prose outside <!-- AUTO-GENERATED --> markers
    • lint_kb_wiki.py then verifies every KB path and relative link the wiki pages reference (existence on disk plus corpus_manifest.jsonl membership for raw/ files) and fails the workflow before upload on any broken reference; run it standalone with uv run -m data.scraping_scripts.lint_kb_wiki [--kb-dir data/kb]
    • The docs and course workflows run these steps automatically unless --skip-kb is passed
    • upload_dbs_to_hf.py packs data/kb/ into kb.tar.gz, uploads it with the vector database, verifies the remote archive, and only then prunes the old unpacked remote KB files