Refactor scraping scripts to enhance source management and contextual node processing. Introduce new helper functions for pruning inactive contextual nodes and managing llms.txt sources. Update existing scripts to utilize the new source registry for improved source handling and streamline the process of adding and updating courses. Enhance error handling and logging for better traceability during data processing.

data/scraping_scripts/README.md

@@ -4,11 +4,18 @@
 
 Make sure you have the required environment variables set:
 
-- `OPENAI_API_KEY` for the LLM
+- `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](https://huggingface.co/datasets/towardsai-tutors/ai-tutor-data/tree/main)
 - `GITHUB_TOKEN` for accessing files via the GitHub API
 
+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`)
+
 ## 1. Prepare the course data
 
 0. Make sure you have access to:
@@ -40,9 +47,10 @@ Make sure you have the required environment variables set:
 7. Rename the folder to the course name.
    - e.g. `master_ai_for_work`
 
-8. Open the `data/scraping_scripts/process_md_files.py` python file and locate the `SOURCE_CONFIGS` dictionary.
+8. Open `data/scraping_scripts/source_registry.py`.
 
-9. Add the new course to the `SOURCE_CONFIGS` dictionary.
+9. Add the new course to the `SOURCE_CONFIGS` dictionary. Sources listed in
+   this registry are active in the knowledge base.
 
    example:
 
@@ -70,13 +78,13 @@ Make sure you have the required environment variables set:
 ## 2. Run the add_course_workflow.py script
 
 ```bash
-uv run -m data.scraping_scripts.add_course_workflow --course [COURSE_NAME]
+uv run -m data.scraping_scripts.add_course_workflow --courses [COURSE_NAME]
 ```
 
 example:
 
 ```bash
-uv run -m data.scraping_scripts.add_course_workflow --course master_ai_for_work
+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:
@@ -88,7 +96,7 @@ This script will guide you through the complete process, it will:
    5. Add contextual information to document chunks before embedding
    6. Create vector stores
    7. Upload databases to HuggingFace
-   8. Update UI configuration
+   8. Confirm the course is configured in the central source registry
 
 ## 3. Add URLs to the course content + Manual Dataset Cleaning (Most important step)
 
@@ -109,7 +117,7 @@ example: "Course Admin and Syllabus", "Course Structure/Overview", "Course Outli
 ## 4. Once done, run the script again and answer "yes" to the question "Have you added all the URLs?"
 
 ```bash
-uv run -m data.scraping_scripts.add_course_workflow --course master_ai_for_work
+uv run -m data.scraping_scripts.add_course_workflow --courses master_ai_for_work
 ```
 
 ## 5. Its done
@@ -172,9 +180,11 @@ 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. Rebuilds `data/chroma-db-all_sources`.
-4. Uploads the rebuilt vector DB and updated aggregate data files.
-5. Deletes the retired per-source JSONL from `towardsai-tutors/ai-tutor-data`.
+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:
 
@@ -182,8 +192,8 @@ 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
 ```
 
-After retiring a source, remove its entry from `SOURCE_CONFIGS` and from the
-workflow download/upload lists so future merges cannot reintroduce it.
+Use `--keep-source-registry` only if you want to remove existing chunks while
+leaving the source configured as active for a future rebuild.
 
 ## Tips for New Team Members
 
@@ -200,9 +210,9 @@ workflow download/upload lists so future merges cannot reintroduce it.
 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 Gradio UI:
-   - The workflow automatically updates `scripts/setup.py` to include the new source and `scripts/main.py` to preselect it by default
+   - Add the source label and default-selection metadata in `source_registry.py`
    - Check that the new source appears in the dropdown menu in the UI
-   - Make sure it's properly included in the default selected sources
+   - Make sure it's properly included in the default selected sources if desired
    - Restart the Gradio app to see the changes
 
 5. First time setup or missing files:

data/scraping_scripts/add_context_to_nodes.py

@@ -1,26 +1,83 @@
 import asyncio
 import json
+import os
 import pickle
+import random
+import re
+import time
+from collections import deque
 from typing import List
 
-import instructor
 import tiktoken
-# from anthropic import AsyncAnthropic
 from dotenv import load_dotenv
+from google import genai
+from google.genai import types
+from google.genai.errors import APIError
 from jinja2 import Template
 from llama_index.core import Document
-from llama_index.core.ingestion import IngestionPipeline
-from llama_index.core.node_parser import SentenceSplitter
-from llama_index.core.schema import TextNode
-from openai import AsyncOpenAI
 from pydantic import BaseModel, Field
-from tenacity import retry, stop_after_attempt, wait_exponential
+from tenacity import retry, retry_if_exception, stop_after_attempt
 from tqdm.asyncio import tqdm
 
-from scripts.chroma_rag import ChunkRecord
+from scripts.chroma_rag import ChunkRecord, build_chunk_records, format_chunk_for_retrieval
 
 load_dotenv(".env")
 
+CONTEXT_MODEL = os.getenv("GEMINI_CONTEXT_MODEL", "gemini-3.1-flash-lite")
+CONTEXT_MAX_OUTPUT_TOKENS = 1000
+CONTEXT_TPM_LIMIT = int(os.getenv("GEMINI_CONTEXT_TPM_LIMIT", "30000000"))
+CONTEXT_TPM_SAFETY_MARGIN = float(
+    os.getenv("GEMINI_CONTEXT_TPM_SAFETY_MARGIN", "0.8")
+)
+CONTEXT_TPM_WINDOW_SECONDS = float(
+    os.getenv("GEMINI_CONTEXT_TPM_WINDOW_SECONDS", "60")
+)
+CONTEXT_RETRY_ATTEMPTS = int(os.getenv("GEMINI_CONTEXT_RETRY_ATTEMPTS", "8"))
+DEFAULT_SEMAPHORE_LIMIT = int(os.getenv("GEMINI_CONTEXT_CONCURRENCY", "50"))
+MAX_DOCUMENT_TOKENS = 120_000
+RETRYABLE_GENAI_STATUS_CODES = {408, 429, 500, 502, 503, 504}
+_genai_client: genai.Client | None = None
+_token_encoding = tiktoken.get_encoding("cl100k_base")
+
+
+class AsyncTokenWindowLimiter:
+    def __init__(self, tokens_per_window: int, window_seconds: float) -> None:
+        self.tokens_per_window = max(1, tokens_per_window)
+        self.window_seconds = window_seconds
+        self._events: deque[tuple[float, int]] = deque()
+        self._used_tokens = 0
+        self._lock = asyncio.Lock()
+
+    async def acquire(self, tokens: int) -> None:
+        tokens = min(max(1, tokens), self.tokens_per_window)
+
+        while True:
+            async with self._lock:
+                now = time.monotonic()
+                self._prune(now)
+
+                if self._used_tokens + tokens <= self.tokens_per_window:
+                    self._events.append((now, tokens))
+                    self._used_tokens += tokens
+                    return
+
+                oldest_at, _ = self._events[0]
+                delay = max(0.1, self.window_seconds - (now - oldest_at))
+
+            await asyncio.sleep(delay + random.uniform(0.1, 0.75))
+
+    def _prune(self, now: float) -> None:
+        while self._events and now - self._events[0][0] >= self.window_seconds:
+            _, tokens = self._events.popleft()
+            self._used_tokens -= tokens
+
+
+_input_token_limiter = AsyncTokenWindowLimiter(
+    tokens_per_window=int(CONTEXT_TPM_LIMIT * CONTEXT_TPM_SAFETY_MARGIN),
+    window_seconds=CONTEXT_TPM_WINDOW_SECONDS,
+)
+
+
 def create_docs(input_file: str) -> List[Document]:
     with open(input_file, "r") as f:
         documents: list[Document] = []
@@ -61,19 +118,90 @@ class SituatedContext(BaseModel):
     )
 
 
-# client = AsyncInstructor(
-#     client=AsyncAnthropic(),
-#     create=patch(
-#         create=AsyncAnthropic().beta.prompt_caching.messages.create,
-#         mode=Mode.ANTHROPIC_TOOLS,
-#     ),
-#     mode=Mode.ANTHROPIC_TOOLS,
-# )
-aclient = AsyncOpenAI()
-client: instructor.AsyncInstructor = instructor.from_openai(aclient)
+def get_genai_client() -> genai.Client:
+    global _genai_client
+
+    if _genai_client is None:
+        api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
+        if not api_key:
+            raise RuntimeError(
+                "GEMINI_API_KEY or GOOGLE_API_KEY must be set to add chunk context."
+            )
+        _genai_client = genai.Client(api_key=api_key)
+
+    return _genai_client
+
+
+def is_retryable_genai_error(exc: BaseException) -> bool:
+    return (
+        isinstance(exc, APIError)
+        and getattr(exc, "code", None) in RETRYABLE_GENAI_STATUS_CODES
+    )
+
+
+def count_input_tokens(content: str) -> int:
+    return len(_token_encoding.encode(content, disallowed_special=()))
+
+
+def extract_retry_delay_seconds(exc: BaseException | None) -> float | None:
+    if not isinstance(exc, APIError):
+        return None
+
+    retry_delay = _find_retry_delay(getattr(exc, "details", None))
+    if retry_delay is not None:
+        return retry_delay
+
+    message = getattr(exc, "message", None) or str(exc)
+    match = re.search(r"retry in ([0-9.]+)s", message, flags=re.IGNORECASE)
+    if match:
+        return float(match.group(1))
+
+    return None
+
+
+def _find_retry_delay(value) -> float | None:
+    if isinstance(value, dict):
+        if value.get("@type") == "type.googleapis.com/google.rpc.RetryInfo":
+            return _parse_duration_seconds(value.get("retryDelay"))
+        for item in value.values():
+            retry_delay = _find_retry_delay(item)
+            if retry_delay is not None:
+                return retry_delay
+    elif isinstance(value, list):
+        for item in value:
+            retry_delay = _find_retry_delay(item)
+            if retry_delay is not None:
+                return retry_delay
+    return None
+
+
+def _parse_duration_seconds(value) -> float | None:
+    if not isinstance(value, str):
+        return None
+
+    match = re.fullmatch(r"([0-9.]+)s", value)
+    if match:
+        return float(match.group(1))
 
+    return None
 
-@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
+
+def wait_for_genai_retry(retry_state) -> float:
+    exc = retry_state.outcome.exception() if retry_state.outcome else None
+    retry_delay = extract_retry_delay_seconds(exc)
+    if retry_delay is not None:
+        return retry_delay + random.uniform(1.0, 4.0)
+
+    exponential_delay = min(60, max(4, 2 ** retry_state.attempt_number))
+    return exponential_delay + random.uniform(0.5, 2.0)
+
+
+@retry(
+    retry=retry_if_exception(is_retryable_genai_error),
+    stop=stop_after_attempt(CONTEXT_RETRY_ATTEMPTS),
+    wait=wait_for_genai_retry,
+    reraise=True,
+)
 async def situate_context(doc: str, chunk: str) -> str:
     template = Template(
         """
@@ -88,85 +216,104 @@ Here is the chunk we want to situate within the whole document above:
 </chunk>
 
 Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.
-Answer only with the succinct context and nothing else.
+Return a title for the document and the succinct context.
 """
     )
 
     content = template.render(doc=doc, chunk=chunk)
-
-    response = await client.chat.completions.create(
-        model="gpt-4o-mini",
-        max_tokens=1000,
-        temperature=0,
-        messages=[
-            {
-                "role": "user",
-                "content": content,
-            }
-        ],
-        response_model=SituatedContext,
-    )
-    return response.context
-
-
-def node_to_chunk_record(node: TextNode) -> ChunkRecord:
-    doc_id = str(node.source_node.node_id)  # type: ignore[union-attr]
-    metadata = dict(node.metadata)
-    metadata["doc_id"] = doc_id
-    return ChunkRecord(
-        chunk_id=str(node.node_id),
-        doc_id=doc_id,
-        text=str(node.text),
-        metadata=metadata,
+    await _input_token_limiter.acquire(count_input_tokens(content))
+
+    response = await get_genai_client().aio.models.generate_content(
+        model=CONTEXT_MODEL,
+        contents=content,
+        config=types.GenerateContentConfig(
+            max_output_tokens=CONTEXT_MAX_OUTPUT_TOKENS,
+            response_mime_type="application/json",
+            response_schema=SituatedContext,
+        ),
     )
-
-
-async def process_chunk(node: TextNode, document_dict: dict[str, Document]) -> ChunkRecord:
-    doc_id: str = node.source_node.node_id  # type: ignore
+    if isinstance(response.parsed, SituatedContext):
+        return response.parsed.context
+    if isinstance(response.parsed, dict):
+        return SituatedContext.model_validate(response.parsed).context
+    if response.text:
+        return SituatedContext.model_validate_json(response.text).context
+    raise ValueError("Gemini returned no context response text.")
+
+
+def document_to_row(document: Document) -> dict:
+    return {
+        "doc_id": document.doc_id,
+        "content": document.get_content(),
+        "name": document.metadata["title"],
+        "url": document.metadata["url"],
+        "source": document.metadata["source"],
+        "retrieve_doc": document.metadata["retrieve_doc"],
+        "tokens": document.metadata["tokens"],
+    }
+
+
+async def process_chunk(
+    chunk_record: ChunkRecord,
+    document_dict: dict[str, Document],
+) -> ChunkRecord:
+    doc_id = chunk_record.doc_id
     doc: Document = document_dict[doc_id]
 
-    if doc.metadata["tokens"] > 120_000:
+    if doc.metadata["tokens"] > MAX_DOCUMENT_TOKENS:
         # Tokenize the document text
-        encoding = tiktoken.encoding_for_model("gpt-4o-mini")
+        encoding = tiktoken.get_encoding("cl100k_base")
         tokens = encoding.encode(doc.get_content())
 
         # Trim to 120,000 tokens
-        trimmed_tokens = tokens[:120_000]
+        trimmed_tokens = tokens[:MAX_DOCUMENT_TOKENS]
 
         # Decode back to text
         trimmed_text = encoding.decode(trimmed_tokens)
 
         # Update the document with trimmed text
         doc = Document(text=trimmed_text, metadata=doc.metadata)
-        doc.metadata["tokens"] = 120_000
-
-    context: str = await situate_context(doc.get_content(), node.text)
-    node.text = f"{node.text}\n\n{context}"
-    return node_to_chunk_record(node)
+        doc.metadata["tokens"] = MAX_DOCUMENT_TOKENS
+
+    context = await situate_context(doc.get_content(), chunk_record.text)
+    metadata = dict(chunk_record.metadata)
+    metadata["raw_text"] = chunk_record.text
+    contextual_text = (
+        f"{format_chunk_for_retrieval(chunk_record.text, metadata)}"
+        f"\n\nContext: {context}"
+    )
+    return ChunkRecord(
+        chunk_id=chunk_record.chunk_id,
+        doc_id=chunk_record.doc_id,
+        text=contextual_text,
+        metadata=metadata,
+    )
 
 
 async def process(
-    documents: List[Document], semaphore_limit: int = 50
+    documents: List[Document], semaphore_limit: int = DEFAULT_SEMAPHORE_LIMIT
 ) -> List[ChunkRecord]:
 
-    # From the document, we create chunks
-    pipeline = IngestionPipeline(
-        transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
-    )
-    all_nodes: list[TextNode] = pipeline.run(documents=documents, show_progress=True)
-    print(f"Number of nodes: {len(all_nodes)}")
+    chunk_records = build_chunk_records([document_to_row(doc) for doc in documents])
+    print(f"Number of chunks: {len(chunk_records)}")
 
     document_dict: dict[str, Document] = {doc.doc_id: doc for doc in documents}
 
+    print(
+        "Gemini context rate limits: "
+        f"{_input_token_limiter.tokens_per_window:,} input tokens/"
+        f"{CONTEXT_TPM_WINDOW_SECONDS:g}s, concurrency={semaphore_limit}"
+    )
+
     semaphore = asyncio.Semaphore(semaphore_limit)
 
-    async def process_with_semaphore(node):
+    async def process_with_semaphore(chunk_record: ChunkRecord):
         async with semaphore:
-            result = await process_chunk(node, document_dict)
+            result = await process_chunk(chunk_record, document_dict)
             await asyncio.sleep(0.1)
             return result
 
-    tasks = [process_with_semaphore(node) for node in all_nodes]
+    tasks = [process_with_semaphore(chunk_record) for chunk_record in chunk_records]
 
     results: List[ChunkRecord] = await tqdm.gather(*tasks, desc="Processing chunks")
 

data/scraping_scripts/add_course_workflow.py

@@ -6,8 +6,8 @@ This script guides you through adding or updating one or more courses in the AI
 
 1. Process course markdown files to create per-course JSONL data
 2. MANDATORY MANUAL STEP: Add URLs to each course JSONL
-3. Rebuild all_sources_data.jsonl from the per-source JSONLs listed in SOURCE_CONFIGS
-   (this naturally drops any retired sources no longer in SOURCE_CONFIGS)
+3. Rebuild all_sources_data.jsonl from active sources in source_registry.py
+   (this naturally drops any retired sources no longer in the registry)
 4. Optionally purge retired sources from the contextual-nodes PKL
 5. Add contextual information to document nodes (only new docs by default)
 6. Create vector stores
@@ -37,13 +37,21 @@ import os
 import pickle
 import subprocess
 import sys
-from pathlib import Path
 from typing import Dict, List, Set
 
 from dotenv import load_dotenv
 from huggingface_hub import hf_hub_download
 
+from data.scraping_scripts.contextual_node_pruning import (
+    prune_contextual_nodes_to_active_sources,
+)
 from data.scraping_scripts.hf_auth import HuggingFaceAuthError, validate_hf_access
+from data.scraping_scripts.source_registry import (
+    SOURCE_CONFIGS,
+    SOURCE_KEY_TO_LABEL,
+    required_data_files,
+    source_output_files,
+)
 from scripts.chroma_rag import get_chunk_record_doc_id
 
 # Load environment variables from .env file
@@ -69,27 +77,10 @@ def ensure_hf_access() -> None:
         sys.exit(1)
 
 
-def ensure_required_files_exist():
+def ensure_required_files_exist(sources_to_regenerate: List[str] | None = None):
     """Download required data files from HuggingFace if they don't exist locally."""
-    # List of files to check and download
-    required_files = {
-        # Critical files
-        "data/all_sources_data.jsonl": "all_sources_data.jsonl",
-        "data/all_sources_contextual_nodes.pkl": "all_sources_contextual_nodes.pkl",
-        # Documentation source files
-        "data/transformers_data.jsonl": "transformers_data.jsonl",
-        "data/peft_data.jsonl": "peft_data.jsonl",
-        "data/trl_data.jsonl": "trl_data.jsonl",
-        "data/llama_index_data.jsonl": "llama_index_data.jsonl",
-        "data/langchain_data.jsonl": "langchain_data.jsonl",
-        "data/openai_cookbooks_data.jsonl": "openai_cookbooks_data.jsonl",
-        # Course files
-        "data/tai_blog_data.jsonl": "tai_blog_data.jsonl",
-        "data/master_ai_for_work_data.jsonl": "master_ai_for_work_data.jsonl",
-        "data/agentic_ai_engineering_data.jsonl": "agentic_ai_engineering_data.jsonl",
-        "data/full_stack_ai_engineering_data.jsonl": "full_stack_ai_engineering_data.jsonl",
-        "data/beginner_python_for_ai_engineering_data.jsonl": "beginner_python_for_ai_engineering_data.jsonl",
-    }
+    required_files = required_data_files()
+    regenerated_source_files = source_output_files(sources_to_regenerate or [])
 
     # Critical files that must be downloaded
     critical_files = [
@@ -99,6 +90,14 @@ def ensure_required_files_exist():
 
     # Check and download each file
     for local_path, remote_filename in required_files.items():
+        if local_path in regenerated_source_files:
+            if not os.path.exists(local_path):
+                logger.info(
+                    "%s will be regenerated for this run; skipping HuggingFace download",
+                    remote_filename,
+                )
+            continue
+
         if not os.path.exists(local_path):
             logger.info(
                 f"{remote_filename} not found. Attempting to download from HuggingFace..."
@@ -164,9 +163,6 @@ def process_markdown_files(course_name: str) -> str:
 
     logger.info(f"Successfully processed markdown files for {course_name}")
 
-    # Determine the output file path from process_md_files.py
-    from data.scraping_scripts.process_md_files import SOURCE_CONFIGS
-
     if course_name not in SOURCE_CONFIGS:
         logger.error(f"Course {course_name} not found in SOURCE_CONFIGS")
         sys.exit(1)
@@ -206,12 +202,11 @@ def manual_url_addition(jsonl_path: str) -> None:
 
 
 def rebuild_all_sources(courses: List[str]) -> None:
-    """Rebuild all_sources_data.jsonl from the per-source JSONLs in SOURCE_CONFIGS.
+    """Rebuild all_sources_data.jsonl from active source JSONLs.
 
     Unlike the previous append-style merge, this drops any source whose entry
-    has been removed from SOURCE_CONFIGS (e.g. retired/renamed courses) and
-    reloads every other source from its own JSONL. Safe against stale rows
-    for sources like llm_developper / python_primer that were renamed.
+    has been removed from source_registry.py and reloads every other active
+    source from its own JSONL.
     """
     from data.scraping_scripts.process_md_files import combine_all_sources
 
@@ -415,85 +410,19 @@ def upload_to_huggingface(upload_jsonl: bool = False) -> None:
 
 
 def update_ui_files(course_name: str) -> None:
-    """Update setup.py and the default selected UI sources for the new course."""
-    logger.info(f"Updating UI files with new course: {course_name}")
-
-    # Get the source configuration for display name
-    from data.scraping_scripts.process_md_files import SOURCE_CONFIGS
-
-    if course_name not in SOURCE_CONFIGS:
-        logger.error(f"Course {course_name} not found in SOURCE_CONFIGS")
+    """Confirm the course is represented in the central source registry."""
+    if course_name not in SOURCE_KEY_TO_LABEL:
+        logger.warning(
+            "%s is not in source_registry.py UI metadata. Add it there if it "
+            "should appear in the app source picker.",
+            course_name,
+        )
         return
 
-    # Get a readable display name for the UI
-    display_name = course_name.replace("_", " ").title()
-
-    # Update setup.py - add to AVAILABLE_SOURCES, AVAILABLE_SOURCES_UI, and SOURCE_UI_TO_KEY
-    setup_path = Path("scripts/setup.py")
-    if setup_path.exists():
-        setup_content = setup_path.read_text()
-
-        # Check if already added
-        if f'"{course_name}"' in setup_content:
-            logger.info(f"Course {course_name} already in setup.py")
-        else:
-            # Add to AVAILABLE_SOURCES_UI
-            ui_list_start = setup_content.find("AVAILABLE_SOURCES_UI = [")
-            ui_list_end = setup_content.find("]", ui_list_start)
-            new_ui_content = (
-                setup_content[:ui_list_end]
-                + f'    "{display_name}",\n'
-                + setup_content[ui_list_end:]
-            )
-
-            # Add to AVAILABLE_SOURCES
-            sources_list_start = new_ui_content.find("AVAILABLE_SOURCES = [")
-            sources_list_end = new_ui_content.find("]", sources_list_start)
-            new_content = (
-                new_ui_content[:sources_list_end]
-                + f'    "{course_name}",\n'
-                + new_ui_content[sources_list_end:]
-            )
-
-            mapping_start = new_content.find("SOURCE_UI_TO_KEY = {")
-            mapping_end = new_content.find("}", mapping_start)
-            if f'"{display_name}": "{course_name}"' not in new_content:
-                new_content = (
-                    new_content[:mapping_end]
-                    + f'    "{display_name}": "{course_name}",\n'
-                    + new_content[mapping_end:]
-                )
-
-            # Write updated content
-            setup_path.write_text(new_content)
-            logger.info(f"Updated setup.py with {course_name}")
-    else:
-        logger.warning(f"setup.py not found at {setup_path}")
-
-    # Update main.py - add to the default selected source list
-    main_path = Path("scripts/main.py")
-    if main_path.exists():
-        main_content = main_path.read_text()
-
-        # Check if already added
-        if f'"{display_name}"' in main_content:
-            logger.info(f"Course {course_name} already in main.py")
-        else:
-            value_start = main_content.find("value=[")
-            value_end = main_content.find("]", value_start)
-            new_main_content = main_content
-            if value_start != -1 and value_end != -1:
-                new_main_content = (
-                    main_content[: value_start + 7]
-                    + f'        "{display_name}",\n'
-                    + main_content[value_start + 7 :]
-                )
-
-            # Write updated content
-            main_path.write_text(new_main_content)
-            logger.info(f"Updated main.py with {course_name}")
-    else:
-        logger.warning(f"main.py not found at {main_path}")
+    logger.info(
+        "%s is configured in source_registry.py; no setup.py/main.py edits needed.",
+        course_name,
+    )
 
 
 def main():
@@ -504,7 +433,7 @@ def main():
         "--courses",
         nargs="+",
         required=True,
-        help="One or more course names to process (must match SOURCE_CONFIGS)",
+        help="One or more course names to process (must match source_registry.py)",
     )
     parser.add_argument(
         "--purge-sources",
@@ -553,19 +482,19 @@ def main():
     args = parser.parse_args()
     courses: List[str] = args.courses
 
-    ensure_hf_access()
-
-    # Ensure required data files exist before proceeding
-    ensure_required_files_exist()
-
-    from data.scraping_scripts.process_md_files import SOURCE_CONFIGS
-
     # Validate every course up front so we fail fast
     for course_name in courses:
         if course_name not in SOURCE_CONFIGS:
             logger.error(f"Course {course_name} not found in SOURCE_CONFIGS")
             sys.exit(1)
 
+    ensure_hf_access()
+
+    # Keep untouched source JSONLs by downloading them when needed, but don't
+    # require first-time courses that this run is about to regenerate.
+    sources_to_regenerate = [] if args.skip_process_md else courses
+    ensure_required_files_exist(sources_to_regenerate=sources_to_regenerate)
+
     # Per-course: process markdown + manual URL addition
     for course_name in courses:
         course_jsonl_path = SOURCE_CONFIGS[course_name]["output_file"]
@@ -591,6 +520,8 @@ def main():
     if not args.skip_context:
         add_context_to_nodes(not args.process_all_context)
 
+    prune_contextual_nodes_to_active_sources()
+
     if not args.skip_vectors:
         create_vector_stores()
 

data/scraping_scripts/capture_source_versions.py

@@ -1,5 +1,5 @@
 """
-Capture the latest release tag + commit SHA for each documentation source and
+Capture the latest release tag + commit SHA or docs timestamp for each source and
 write them to `data/source_versions.json`, so the frontend can surface which
 library version is represented in the knowledge base (and how fresh it is).
 
@@ -7,8 +7,8 @@ Usage:
     uv run -m data.scraping_scripts.capture_source_versions
     uv run -m data.scraping_scripts.capture_source_versions --sources langchain peft
 
-This is called automatically from `update_docs_workflow.py` after the GitHub
-download step, but can also be run standalone to refresh the JSON.
+This is called automatically from `update_docs_workflow.py` after the download
+step, but can also be run standalone to refresh the JSON.
 
 Output shape (per source):
     {
@@ -31,6 +31,11 @@ from typing import Optional
 import requests
 from dotenv import load_dotenv
 
+try:
+    from data.scraping_scripts.source_registry import DOC_SOURCE_KEYS, SOURCE_CONFIGS
+except ModuleNotFoundError:
+    from source_registry import DOC_SOURCE_KEYS, SOURCE_CONFIGS
+
 load_dotenv()
 
 OUTPUT_PATH = Path("data/source_versions.json")
@@ -43,6 +48,8 @@ VERSION_REPOS: dict[str, tuple[str, str]] = {
     "trl": ("huggingface", "trl"),
     "llama_index": ("run-llama", "llama_index"),
     "langchain": ("langchain-ai", "langchain"),
+    "langgraph": ("langchain-ai", "langgraph"),
+    "deep_agents": ("langchain-ai", "deepagents"),
     "openai_cookbooks": ("openai", "openai-cookbook"),
 }
 
@@ -91,7 +98,22 @@ def fetch_default_branch_sha(owner: str, repo: str) -> Optional[str]:
     return sha[:7] if isinstance(sha, str) else None
 
 
-def capture_for_source(source: str) -> dict:
+def fetch_last_modified(url: str) -> Optional[str]:
+    try:
+        response = requests.head(url, timeout=30, allow_redirects=True)
+    except requests.RequestException as exc:
+        print(f"  ! network error: {exc}", file=sys.stderr)
+        return None
+    if not response.ok:
+        print(
+            f"  ! HTTP {response.status_code} for {url}: {response.text[:160]}",
+            file=sys.stderr,
+        )
+        return None
+    return response.headers.get("Last-Modified")
+
+
+def capture_github_source(source: str) -> dict:
     owner, repo = VERSION_REPOS[source]
     print(f"- {source}: querying {owner}/{repo}")
     return {
@@ -101,6 +123,25 @@ def capture_for_source(source: str) -> dict:
     }
 
 
+def capture_llms_txt_source(source: str) -> dict:
+    config = SOURCE_CONFIGS[source]
+    urls = [str(url) for url in config.get("llms_txt_urls", [])]
+    print(f"- {source}: querying llms.txt metadata")
+    return {
+        "version": fetch_last_modified(urls[0]) if urls else None,
+        "sha": None,
+        "indexedAt": datetime.now(timezone.utc).date().isoformat(),
+    }
+
+
+def capture_for_source(source: str) -> dict | None:
+    if source in VERSION_REPOS:
+        return capture_github_source(source)
+    if SOURCE_CONFIGS.get(source, {}).get("llms_txt_urls"):
+        return capture_llms_txt_source(source)
+    return None
+
+
 def load_existing() -> dict:
     if not OUTPUT_PATH.exists():
         return {}
@@ -115,10 +156,11 @@ def load_existing() -> dict:
 def capture(sources: list[str]) -> dict:
     versions = load_existing()
     for source in sources:
-        if source not in VERSION_REPOS:
+        metadata = capture_for_source(source)
+        if metadata is None:
             print(f"  ! skipping unknown source: {source}", file=sys.stderr)
             continue
-        versions[source] = capture_for_source(source)
+        versions[source] = metadata
 
     OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
     with OUTPUT_PATH.open("w", encoding="utf-8") as f:
@@ -133,8 +175,8 @@ def main() -> None:
     parser.add_argument(
         "--sources",
         nargs="+",
-        choices=list(VERSION_REPOS.keys()),
-        default=list(VERSION_REPOS.keys()),
+        choices=list(DOC_SOURCE_KEYS),
+        default=list(DOC_SOURCE_KEYS),
         help="Subset of sources to refresh (default: all).",
     )
     args = parser.parse_args()

data/scraping_scripts/contextual_node_pruning.py

@@ -0,0 +1,67 @@
+"""Helpers for keeping contextual-node pickles aligned with active sources."""
+
+from __future__ import annotations
+
+import logging
+import os
+import pickle
+from collections import Counter
+from typing import Any
+
+from data.scraping_scripts.source_registry import (
+    ACTIVE_SOURCE_KEYS,
+    CONTEXTUAL_NODES_PKL,
+)
+from scripts.chroma_rag import get_chunk_record_source
+
+logger = logging.getLogger(__name__)
+
+
+def prune_contextual_nodes_to_active_sources(
+    pkl_path: str = CONTEXTUAL_NODES_PKL,
+) -> Counter[str]:
+    """Remove contextual nodes whose source is not active in source_registry.py."""
+    if not os.path.exists(pkl_path):
+        logger.info("%s does not exist; no contextual nodes to prune", pkl_path)
+        return Counter()
+
+    with open(pkl_path, "rb") as handle:
+        nodes = pickle.load(handle)
+
+    kept_nodes: list[Any] = []
+    removed_counts: Counter[str] = Counter()
+    unknown_source = 0
+
+    for node in nodes:
+        try:
+            source = get_chunk_record_source(node)
+        except Exception:
+            unknown_source += 1
+            kept_nodes.append(node)
+            continue
+
+        if source in ACTIVE_SOURCE_KEYS:
+            kept_nodes.append(node)
+        else:
+            removed_counts[str(source)] += 1
+
+    if not removed_counts:
+        if unknown_source:
+            logger.info(
+                "No inactive contextual nodes pruned; kept %s nodes with unknown source",
+                unknown_source,
+            )
+        return removed_counts
+
+    with open(pkl_path, "wb") as handle:
+        pickle.dump(kept_nodes, handle)
+
+    logger.info(
+        "Pruned %s inactive contextual nodes from %s: %s",
+        sum(removed_counts.values()),
+        pkl_path,
+        dict(sorted(removed_counts.items())),
+    )
+    if unknown_source:
+        logger.info("Kept %s contextual nodes with unknown source", unknown_source)
+    return removed_counts

data/scraping_scripts/create_vector_stores.py

@@ -22,13 +22,22 @@ from dotenv import load_dotenv
 from tqdm.auto import tqdm
 
 try:
+    from data.scraping_scripts.source_registry import (
+        ACTIVE_SOURCE_KEYS,
+        vector_store_source_configs,
+    )
     from data.scraping_scripts.add_context_to_nodes import create_docs, process
     from scripts.chroma_rag import (
+        DEFAULT_COHERE_EMBED_BATCH_SIZE,
         DEFAULT_EMBED_MODEL,
+        BM25Index,
+        build_chunk_records,
         build_document_dict,
         embed_texts,
+        get_chunk_record_source,
         load_jsonl_documents,
         normalize_chunk_record,
+        save_bm25_index,
         save_document_dict,
     )
 except ModuleNotFoundError:
@@ -38,13 +47,22 @@ except ModuleNotFoundError:
     project_root = Path(__file__).resolve().parents[2]
     if str(project_root) not in sys.path:
         sys.path.insert(0, str(project_root))
+    from data.scraping_scripts.source_registry import (
+        ACTIVE_SOURCE_KEYS,
+        vector_store_source_configs,
+    )
     from data.scraping_scripts.add_context_to_nodes import create_docs, process
     from scripts.chroma_rag import (
+        DEFAULT_COHERE_EMBED_BATCH_SIZE,
         DEFAULT_EMBED_MODEL,
+        BM25Index,
+        build_chunk_records,
         build_document_dict,
         embed_texts,
+        get_chunk_record_source,
         load_jsonl_documents,
         normalize_chunk_record,
+        save_bm25_index,
         save_document_dict,
     )
 
@@ -54,55 +72,32 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-SOURCE_CONFIGS = {
-    "transformers": {
-        "input_file": "data/transformers_data.jsonl",
-        "db_name": "chroma-db-transformers",
-        "document_dict_file": "document_dict_transformers.pkl",
-    },
-    "peft": {
-        "input_file": "data/peft_data.jsonl",
-        "db_name": "chroma-db-peft",
-        "document_dict_file": "document_dict_peft.pkl",
-    },
-    "trl": {
-        "input_file": "data/trl_data.jsonl",
-        "db_name": "chroma-db-trl",
-        "document_dict_file": "document_dict_trl.pkl",
-    },
-    "llama_index": {
-        "input_file": "data/llama_index_data.jsonl",
-        "db_name": "chroma-db-llama_index",
-        "document_dict_file": "document_dict_llama_index.pkl",
-    },
-    "openai_cookbooks": {
-        "input_file": "data/openai_cookbooks_data.jsonl",
-        "db_name": "chroma-db-openai_cookbooks",
-        "document_dict_file": "document_dict_openai_cookbooks.pkl",
-    },
-    "langchain": {
-        "input_file": "data/langchain_data.jsonl",
-        "db_name": "chroma-db-langchain",
-        "document_dict_file": "document_dict_langchain.pkl",
-    },
-    "tai_blog": {
-        "input_file": "data/tai_blog_data.jsonl",
-        "db_name": "chroma-db-tai_blog",
-        "document_dict_file": "document_dict_tai_blog.pkl",
-    },
-    "all_sources": {
-        "input_file": "data/all_sources_data.jsonl",
-        "db_name": "chroma-db-all_sources",
-        "document_dict_file": "document_dict_all_sources.pkl",
-    },
-}
+SOURCE_CONFIGS = vector_store_source_configs()
 
 
 def load_or_create_chunk_records(source: str) -> list[Any]:
     config = SOURCE_CONFIGS[source]
     if source == "all_sources" and os.path.exists("data/all_sources_contextual_nodes.pkl"):
         with open("data/all_sources_contextual_nodes.pkl", "rb") as handle:
-            return pickle.load(handle)
+            records = pickle.load(handle)
+        active_records = []
+        skipped_count = 0
+        for record in records:
+            try:
+                record_source = get_chunk_record_source(record)
+            except Exception:
+                active_records.append(record)
+                continue
+            if record_source in ACTIVE_SOURCE_KEYS:
+                active_records.append(record)
+            else:
+                skipped_count += 1
+        if skipped_count:
+            logger.info(
+                "Skipped %s inactive contextual chunks while building all_sources",
+                skipped_count,
+            )
+        return active_records
 
     documents = create_docs(config["input_file"])
     return asyncio.run(process(documents))
@@ -113,7 +108,51 @@ def iter_batches(items: Sequence[Any], batch_size: int):
         yield items[start : start + batch_size]
 
 
-def process_source(source: str) -> None:
+def get_collection_ids(collection, batch_size: int = 5000) -> set[str]:
+    ids: set[str] = set()
+    offset = 0
+
+    while True:
+        result = collection.get(limit=batch_size, offset=offset, include=[])
+        batch_ids = result.get("ids", [])
+        if not batch_ids:
+            break
+
+        ids.update(str(chunk_id) for chunk_id in batch_ids)
+        offset += len(batch_ids)
+
+    return ids
+
+
+def write_retrieval_artifacts(
+    *,
+    config: dict[str, str],
+    document_rows: list[dict[str, Any]],
+    db_path: str,
+) -> int:
+    document_dict = build_document_dict(document_rows)
+    save_document_dict(
+        document_dict,
+        f"{db_path}/{config['document_dict_file']}",
+    )
+    bm25_records = build_chunk_records(document_rows)
+    save_bm25_index(
+        BM25Index.build(bm25_records),
+        f"{db_path}/{config['bm25_index_file']}",
+    )
+    return len(bm25_records)
+
+
+def process_source(
+    source: str,
+    *,
+    force_rebuild: bool = False,
+    skip_dense_embeddings: bool = False,
+    embed_batch_size: int = DEFAULT_COHERE_EMBED_BATCH_SIZE,
+    cohere_embed_inputs_per_minute: int | None = None,
+    cohere_embed_tpm_limit: int | None = None,
+    cohere_embed_rpm_limit: int | None = None,
+) -> None:
     config = SOURCE_CONFIGS[source]
     document_rows = load_jsonl_documents(config["input_file"])
     if not document_rows:
@@ -122,11 +161,23 @@ def process_source(source: str) -> None:
 
     db_name = config["db_name"]
     db_path = f"data/{db_name}"
-    if os.path.exists(db_path):
+    if force_rebuild and not skip_dense_embeddings and os.path.exists(db_path):
         shutil.rmtree(db_path)
 
     os.makedirs(db_path, exist_ok=True)
 
+    if skip_dense_embeddings:
+        bm25_count = write_retrieval_artifacts(
+            config=config,
+            document_rows=document_rows,
+            db_path=db_path,
+        )
+        print(
+            f"Indexed {bm25_count} BM25 chunks from {len(document_rows)} documents "
+            f"into {db_path}; skipped dense embedding updates"
+        )
+        return
+
     chunk_records = [
         normalize_chunk_record(record)
         for record in load_or_create_chunk_records(source)
@@ -136,9 +187,6 @@ def process_source(source: str) -> None:
         return
 
     chunk_ids = [record.chunk_id for record in chunk_records]
-    chunk_texts = [record.text for record in chunk_records]
-    chunk_metadatas = [record.metadata for record in chunk_records]
-
     logger.info(
         "Preparing %s chunks from %s documents for %s",
         len(chunk_records),
@@ -146,58 +194,108 @@ def process_source(source: str) -> None:
         source,
     )
 
-    cohere_client = cohere.ClientV2(api_key=os.environ["COHERE_API_KEY"])
-    logger.info("Generating embeddings for %s", source)
-    embeddings = embed_texts(
-        cohere_client,
-        chunk_texts,
-        input_type="search_document",
-        model=DEFAULT_EMBED_MODEL,
-        show_progress=True,
-        progress_desc=f"Embedding {source}",
-    )
-
     chroma_client = chromadb.PersistentClient(path=db_path)
     collection = chroma_client.get_or_create_collection(name=db_name)
     max_batch_size = chroma_client.get_max_batch_size()
+    existing_ids = set() if force_rebuild else get_collection_ids(collection)
+    desired_ids = set(chunk_ids)
+
+    stale_ids = sorted(existing_ids - desired_ids)
+    if stale_ids:
+        logger.info("Deleting %s stale embeddings for %s", len(stale_ids), source)
+        for batch_ids in iter_batches(stale_ids, max_batch_size):
+            collection.delete(ids=batch_ids)
+
+    reusable_ids = existing_ids & desired_ids
+    records_to_embed = [
+        record for record in chunk_records if record.chunk_id not in reusable_ids
+    ]
+
     logger.info(
-        "Writing %s embeddings to Chroma for %s in batches of up to %s",
-        len(chunk_ids),
+        "Reusing %s existing embeddings and generating %s embeddings for %s",
+        len(reusable_ids),
+        len(records_to_embed),
         source,
-        max_batch_size,
-    )
-    with tqdm(total=len(chunk_ids), desc=f"Upserting {source}", unit="chunk") as progress:
-        for batch_ids, batch_embeddings, batch_texts, batch_metadatas in zip(
-            iter_batches(chunk_ids, max_batch_size),
-            iter_batches(embeddings, max_batch_size),
-            iter_batches(chunk_texts, max_batch_size),
-            iter_batches(chunk_metadatas, max_batch_size),
-        ):
-            collection.upsert(
-                ids=batch_ids,
-                embeddings=batch_embeddings,
-                documents=batch_texts,
-                metadatas=batch_metadatas,
-            )
-            progress.update(len(batch_ids))
+    )
 
-    document_dict = build_document_dict(document_rows)
-    save_document_dict(
-        document_dict,
-        f"{db_path}/{config['document_dict_file']}",
+    if records_to_embed:
+        cohere_client = cohere.ClientV2(api_key=os.environ["COHERE_API_KEY"])
+        texts_to_embed = [record.text for record in records_to_embed]
+        ids_to_embed = [record.chunk_id for record in records_to_embed]
+        metadatas_to_embed = [record.metadata for record in records_to_embed]
+
+        logger.info("Generating embeddings for %s", source)
+        embeddings = embed_texts(
+            cohere_client,
+            texts_to_embed,
+            input_type="search_document",
+            model=DEFAULT_EMBED_MODEL,
+            batch_size=embed_batch_size,
+            max_inputs_per_minute=cohere_embed_inputs_per_minute,
+            max_tokens_per_minute=cohere_embed_tpm_limit,
+            max_requests_per_minute=cohere_embed_rpm_limit,
+            show_progress=True,
+            progress_desc=f"Embedding {source}",
+        )
+
+        logger.info(
+            "Writing %s new embeddings to Chroma for %s in batches of up to %s",
+            len(ids_to_embed),
+            source,
+            max_batch_size,
+        )
+        with tqdm(
+            total=len(ids_to_embed), desc=f"Upserting {source}", unit="chunk"
+        ) as progress:
+            for batch_ids, batch_embeddings, batch_texts, batch_metadatas in zip(
+                iter_batches(ids_to_embed, max_batch_size),
+                iter_batches(embeddings, max_batch_size),
+                iter_batches(texts_to_embed, max_batch_size),
+                iter_batches(metadatas_to_embed, max_batch_size),
+            ):
+                collection.upsert(
+                    ids=batch_ids,
+                    embeddings=batch_embeddings,
+                    documents=batch_texts,
+                    metadatas=batch_metadatas,
+                )
+                progress.update(len(batch_ids))
+
+    bm25_count = write_retrieval_artifacts(
+        config=config,
+        document_rows=document_rows,
+        db_path=db_path,
     )
 
     print(
-        f"Indexed {len(chunk_records)} chunks from {len(document_rows)} documents into {db_path}"
+        f"Indexed {len(chunk_records)} dense chunks and {bm25_count} BM25 chunks "
+        f"from {len(document_rows)} documents into {db_path}"
     )
 
 
-def main(sources: list[str]) -> None:
+def main(
+    sources: list[str],
+    *,
+    force_rebuild: bool = False,
+    skip_dense_embeddings: bool = False,
+    embed_batch_size: int = DEFAULT_COHERE_EMBED_BATCH_SIZE,
+    cohere_embed_inputs_per_minute: int | None = None,
+    cohere_embed_tpm_limit: int | None = None,
+    cohere_embed_rpm_limit: int | None = None,
+) -> None:
     for source in sources:
         if source not in SOURCE_CONFIGS:
             print(f"Unknown source: {source}")
             continue
-        process_source(source)
+        process_source(
+            source,
+            force_rebuild=force_rebuild,
+            skip_dense_embeddings=skip_dense_embeddings,
+            embed_batch_size=embed_batch_size,
+            cohere_embed_inputs_per_minute=cohere_embed_inputs_per_minute,
+            cohere_embed_tpm_limit=cohere_embed_tpm_limit,
+            cohere_embed_rpm_limit=cohere_embed_rpm_limit,
+        )
 
 
 if __name__ == "__main__":
@@ -210,5 +308,59 @@ if __name__ == "__main__":
         choices=SOURCE_CONFIGS.keys(),
         help="Specify one or more sources to process",
     )
+    parser.add_argument(
+        "--force-rebuild",
+        action="store_true",
+        help="Delete the existing Chroma directory and regenerate all embeddings.",
+    )
+    parser.add_argument(
+        "--skip-dense-embeddings",
+        action="store_true",
+        help=(
+            "Only write retrieval artifacts that do not require Cohere embeddings "
+            "(document dictionary and BM25 index). Existing Chroma data is left untouched."
+        ),
+    )
+    parser.add_argument(
+        "--embed-batch-size",
+        type=int,
+        default=DEFAULT_COHERE_EMBED_BATCH_SIZE,
+        help="Maximum number of chunks to include in one Cohere embed request.",
+    )
+    parser.add_argument(
+        "--cohere-embed-inputs-per-minute",
+        type=int,
+        default=None,
+        help=(
+            "Cohere embed input-per-minute limit before the safety margin. "
+            "Defaults to COHERE_EMBED_INPUTS_PER_MINUTE or 2000; use 0 to disable."
+        ),
+    )
+    parser.add_argument(
+        "--cohere-embed-tpm-limit",
+        type=int,
+        default=None,
+        help=(
+            "Cohere embed token-per-minute limit before the safety margin. "
+            "Defaults to COHERE_EMBED_TPM_LIMIT or 0; use 0 to disable."
+        ),
+    )
+    parser.add_argument(
+        "--cohere-embed-rpm-limit",
+        type=int,
+        default=None,
+        help=(
+            "Cohere embed request-per-minute limit before the safety margin. "
+            "Defaults to COHERE_EMBED_RPM_LIMIT or 0; use 0 to disable."
+        ),
+    )
     args = parser.parse_args()
-    main(args.sources)
+    main(
+        args.sources,
+        force_rebuild=args.force_rebuild,
+        skip_dense_embeddings=args.skip_dense_embeddings,
+        embed_batch_size=args.embed_batch_size,
+        cohere_embed_inputs_per_minute=args.cohere_embed_inputs_per_minute,
+        cohere_embed_tpm_limit=args.cohere_embed_tpm_limit,
+        cohere_embed_rpm_limit=args.cohere_embed_rpm_limit,
+    )

data/scraping_scripts/github_to_markdown_ai_docs.py

@@ -42,6 +42,11 @@ import requests
 from dotenv import load_dotenv
 from nbconvert import MarkdownExporter
 
+try:
+    from data.scraping_scripts.source_registry import GITHUB_SOURCE_KEYS
+except ModuleNotFoundError:
+    from source_registry import GITHUB_SOURCE_KEYS
+
 load_dotenv()
 
 # Configuration for different sources
@@ -113,6 +118,23 @@ SOURCE_CONFIGS = {
         ],
         "local_dir": "data/langchain_md_files",
     },
+    "langgraph": {
+        "owner": "langchain-ai",
+        "repo": "docs",
+        "path": "src/oss/langgraph",
+        "local_dir": "data/langgraph_md_files",
+    },
+    "deep_agents": {
+        "owner": "langchain-ai",
+        "repo": "docs",
+        "path": "src/oss/deepagents",
+        "local_dir": "data/deep_agents_md_files",
+    },
+}
+SOURCE_CONFIGS = {
+    source: config
+    for source, config in SOURCE_CONFIGS.items()
+    if source in GITHUB_SOURCE_KEYS
 }
 
 # GitHub Personal Access Token (replace with your own token)
@@ -125,6 +147,7 @@ if GITHUB_TOKEN:
 
 # Maximum number of retries
 MAX_RETRIES = 5
+SOURCE_EXTENSIONS_FILENAME = "_source_extensions.json"
 
 
 class GitHubAPIError(RuntimeError):
@@ -242,11 +265,22 @@ def download_file(file_url: str, file_path: str, retries: int = 0):
     #         f.write(markdown)
 
 
+def strip_notebook_outputs(notebook: nbformat.NotebookNode) -> nbformat.NotebookNode:
+    for cell in notebook.cells:
+        if cell.get("cell_type") == "code":
+            cell["outputs"] = []
+            cell["execution_count"] = None
+        if "attachments" in cell:
+            cell["attachments"] = {}
+    return notebook
+
+
 def convert_ipynb_to_md(ipynb_path: str, md_path: str):
     try:
         with open(ipynb_path, "r", encoding="utf-8") as f:
             notebook = nbformat.read(f, as_version=4)
 
+        notebook = strip_notebook_outputs(notebook)
         exporter = MarkdownExporter()
         markdown, _ = exporter.from_notebook_node(notebook)
 
@@ -260,7 +294,23 @@ def convert_ipynb_to_md(ipynb_path: str, md_path: str):
         print("Skipping this file and continuing with others...")
 
 
-def fetch_files(api_url: str, local_dir: str):
+def manifest_path(file_path: str, root_dir: str) -> str:
+    return os.path.relpath(file_path, root_dir).replace(os.sep, "/")
+
+
+def save_source_extension_manifest(local_dir: str, manifest: dict[str, str]) -> None:
+    path = os.path.join(local_dir, SOURCE_EXTENSIONS_FILENAME)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(dict(sorted(manifest.items())), f, indent=2)
+        f.write("\n")
+
+
+def fetch_files(
+    api_url: str,
+    local_dir: str,
+    source_extensions: dict[str, str],
+    root_dir: str,
+):
     files = get_files_in_directory(api_url)
     if isinstance(files, dict):
         files = [files]
@@ -278,10 +328,14 @@ def fetch_files(api_url: str, local_dir: str):
                 print(f"Converting {file_name} to markdown...")
                 convert_ipynb_to_md(file_path, md_file_path)
                 os.remove(file_path)  # Remove the .ipynb file after conversion
+                source_extensions[manifest_path(md_file_path, root_dir)] = ".ipynb"
+            else:
+                _, extension = os.path.splitext(file_name)
+                source_extensions[manifest_path(file_path, root_dir)] = extension
         elif file["type"] == "dir":
             subdir = os.path.join(local_dir, file["name"])
             os.makedirs(subdir, exist_ok=True)
-            fetch_files(file["url"], subdir)
+            fetch_files(file["url"], subdir, source_extensions, root_dir)
 
 
 def process_source(source: str):
@@ -295,6 +349,7 @@ def process_source(source: str):
     local_dir = config.get("local_dir", f"data/{config['repo']}_md_files")
     shutil.rmtree(local_dir, ignore_errors=True)
     os.makedirs(local_dir, exist_ok=True)
+    source_extensions: dict[str, str] = {}
 
     print(f"Processing source: {source}")
 
@@ -306,13 +361,14 @@ def process_source(source: str):
             )
             target_dir = os.path.join(local_dir, path_config.get("local_subdir", ""))
             os.makedirs(target_dir, exist_ok=True)
-            fetch_files(api_url, target_dir)
+            fetch_files(api_url, target_dir, source_extensions, local_dir)
     else:
         api_url = (
             f"https://api.github.com/repos/{config['owner']}/{config['repo']}/contents/{config['path']}"
         )
-        fetch_files(api_url, local_dir)
+        fetch_files(api_url, local_dir, source_extensions, local_dir)
 
+    save_source_extension_manifest(local_dir, source_extensions)
     print(f"Finished processing {source}")
 
 

data/scraping_scripts/llms_txt_to_markdown_docs.py

@@ -0,0 +1,194 @@
+"""
+Fetch Markdown documentation pages listed in llms.txt indexes.
+
+This downloader is for official docs sites that expose AI-friendly Markdown
+indexes, such as OpenAI's developer docs. It writes the linked Markdown pages
+into a local directory so the existing process_md_files.py pipeline can create
+JSONL, contextual nodes, and vector stores without needing an HTML crawler.
+
+Usage:
+    uv run -m data.scraping_scripts.llms_txt_to_markdown_docs openai_docs
+    uv run -m data.scraping_scripts.llms_txt_to_markdown_docs openai_docs --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+from datetime import datetime, timezone
+from typing import Iterable
+from urllib.parse import urlparse
+
+import requests
+from dotenv import load_dotenv
+
+try:
+    from data.scraping_scripts.source_registry import (
+        LLMS_TXT_SOURCE_KEYS,
+        SOURCE_CONFIGS,
+    )
+except ModuleNotFoundError:
+    from source_registry import LLMS_TXT_SOURCE_KEYS, SOURCE_CONFIGS
+
+load_dotenv()
+
+LINK_PATTERN = re.compile(r"\[[^\]]+\]\((https?://[^)\s]+)\)")
+SOURCE_EXTENSIONS_FILENAME = "_source_extensions.json"
+SOURCE_URLS_FILENAME = "_source_urls.json"
+LLMS_TXT_MANIFEST_FILENAME = "_llms_txt_manifest.json"
+REQUEST_TIMEOUT_SECONDS = 60
+
+
+class LLMSTxtDownloadError(RuntimeError):
+    """Raised when an llms.txt source cannot be downloaded safely."""
+
+
+def fetch_text(url: str) -> str:
+    response = requests.get(url, timeout=REQUEST_TIMEOUT_SECONDS)
+    response.raise_for_status()
+    return response.text
+
+
+def parse_markdown_links(index_text: str) -> list[str]:
+    """Extract HTTP(S) Markdown links from an llms.txt index."""
+    urls = []
+    seen = set()
+    for match in LINK_PATTERN.finditer(index_text):
+        url = match.group(1).strip()
+        if url not in seen:
+            urls.append(url)
+            seen.add(url)
+    return urls
+
+
+def should_download_url(url: str, include_prefixes: Iterable[str]) -> bool:
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"}:
+        return False
+    if not parsed.path.endswith((".md", ".mdx")):
+        return False
+
+    prefixes = tuple(include_prefixes)
+    return not prefixes or any(url.startswith(prefix) for prefix in prefixes)
+
+
+def local_relative_path(url: str) -> str:
+    parsed = urlparse(url)
+    relative_path = parsed.path.lstrip("/")
+    if not relative_path:
+        raise LLMSTxtDownloadError(f"Could not derive a local path from URL: {url}")
+    normalized = os.path.normpath(relative_path).replace(os.sep, "/")
+    if normalized == ".." or normalized.startswith("../"):
+        raise LLMSTxtDownloadError(f"Refusing unsafe local path for URL: {url}")
+    return normalized
+
+
+def write_text_file(path: str, content: str) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(content)
+        if content and not content.endswith("\n"):
+            f.write("\n")
+
+
+def write_json(path: str, payload: object) -> None:
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+        f.write("\n")
+
+
+def collect_doc_urls(config: dict) -> list[str]:
+    index_urls = config.get("llms_txt_urls") or []
+    if not index_urls:
+        raise LLMSTxtDownloadError("Source config is missing llms_txt_urls")
+
+    include_prefixes = config.get("llms_url_include_prefixes") or []
+    docs: list[str] = []
+    seen = set()
+
+    for index_url in index_urls:
+        print(f"Fetching index: {index_url}")
+        index_text = fetch_text(str(index_url))
+        for url in parse_markdown_links(index_text):
+            if not should_download_url(url, include_prefixes):
+                continue
+            if url in seen:
+                continue
+            docs.append(url)
+            seen.add(url)
+
+    return docs
+
+
+def process_source(source: str, *, dry_run: bool = False) -> None:
+    if source not in LLMS_TXT_SOURCE_KEYS:
+        available = ", ".join(LLMS_TXT_SOURCE_KEYS)
+        raise LLMSTxtDownloadError(
+            f"Unknown llms.txt source '{source}'. Available sources: {available}"
+        )
+
+    config = SOURCE_CONFIGS[source]
+    doc_urls = collect_doc_urls(config)
+
+    if not doc_urls:
+        raise LLMSTxtDownloadError(f"No Markdown docs found for source: {source}")
+
+    if dry_run:
+        print(f"Found {len(doc_urls)} Markdown docs for {source}")
+        return
+
+    local_dir = str(config["input_directory"])
+    shutil.rmtree(local_dir, ignore_errors=True)
+    os.makedirs(local_dir, exist_ok=True)
+
+    source_extensions: dict[str, str] = {}
+    source_urls: dict[str, str] = {}
+
+    print(f"Downloading {len(doc_urls)} Markdown docs for {source}")
+    for doc_url in doc_urls:
+        relative_path = local_relative_path(doc_url)
+        local_path = os.path.join(local_dir, relative_path)
+        print(f"Downloading {doc_url}")
+        content = fetch_text(doc_url)
+        write_text_file(local_path, content)
+        source_extensions[relative_path] = os.path.splitext(relative_path)[1]
+        source_urls[relative_path] = doc_url
+
+    write_json(os.path.join(local_dir, SOURCE_EXTENSIONS_FILENAME), source_extensions)
+    write_json(os.path.join(local_dir, SOURCE_URLS_FILENAME), source_urls)
+    write_json(
+        os.path.join(local_dir, LLMS_TXT_MANIFEST_FILENAME),
+        {
+            "downloadedAt": datetime.now(timezone.utc).isoformat(),
+            "indexUrls": config.get("llms_txt_urls") or [],
+            "documentCount": len(doc_urls),
+        },
+    )
+    print(f"Finished processing {source}")
+
+
+def main(sources: list[str], *, dry_run: bool = False) -> None:
+    for source in sources:
+        process_source(source, dry_run=dry_run)
+    print("All specified llms.txt sources have been processed.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "sources",
+        nargs="+",
+        choices=LLMS_TXT_SOURCE_KEYS,
+        help="Specify one or more llms.txt-backed sources to process",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Fetch indexes and report document counts without writing files",
+    )
+    args = parser.parse_args()
+
+    main(args.sources, dry_run=args.dry_run)

data/scraping_scripts/process_md_files.py

@@ -18,14 +18,15 @@ Key features:
 Usage:
     python process_md_files.py <source1> <source2> ...
 
-Where <source1>, <source2>, etc. are one or more of the predefined sources in SOURCE_CONFIGS
+Where <source1>, <source2>, etc. are one or more of the active sources in
+source_registry.py
 (e.g., 'transformers', 'llama_index', 'openai_cookbooks').
 
 The script processes all Markdown files in the specified input directories (and their subdirectories),
 applies the configured filters, and saves the results in JSONL files. Each line in the output
 files represents a single document with metadata and content.
 
-To add or modify sources, update the SOURCE_CONFIGS dictionary at the top of the script.
+To add, modify, or retire sources, update data/scraping_scripts/source_registry.py.
 """
 
 import argparse
@@ -38,170 +39,90 @@ from typing import Dict, List
 
 import tiktoken
 
+try:
+    from data.scraping_scripts.source_registry import SOURCE_CONFIGS
+except ModuleNotFoundError:
+    from source_registry import SOURCE_CONFIGS
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Configuration for different sources
-SOURCE_CONFIGS = {
-    "transformers": {
-        "base_url": "https://huggingface.co/docs/transformers/",
-        "input_directory": "data/transformers_md_files",
-        "output_file": "data/transformers_data.jsonl",
-        "source_name": "transformers",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": ["internal", "main_classes"],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "peft": {
-        "base_url": "https://huggingface.co/docs/peft/",
-        "input_directory": "data/peft_md_files",
-        "output_file": "data/peft_data.jsonl",
-        "source_name": "peft",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "trl": {
-        "base_url": "https://huggingface.co/docs/trl/",
-        "input_directory": "data/trl_md_files",
-        "output_file": "data/trl_data.jsonl",
-        "source_name": "trl",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "llama_index": {
-        "base_url": "https://docs.llamaindex.ai/en/stable/",
-        "input_directory": "data/llama_index_md_files",
-        "output_file": "data/llama_index_data.jsonl",
-        "source_name": "llama_index",
-        "use_include_list": True,
-        "included_dirs": [
-            "src/content/docs/framework/index.md",
-            "src/content/docs/framework/getting_started",
-            "src/content/docs/framework/understanding",
-            "src/content/docs/framework/use_cases",
-            "src/content/docs/framework/module_guides",
-            "src/content/docs/framework/optimizing",
-            "examples",
-        ],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "openai_cookbooks": {
-        "base_url": "https://github.com/openai/openai-cookbook/blob/main/examples/",
-        "input_directory": "data/openai-cookbook_md_files",
-        "output_file": "data/openai_cookbooks_data.jsonl",
-        "source_name": "openai_cookbooks",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": ".ipynb",
-    },
-    "langchain": {
-        "base_url": "https://docs.langchain.com/oss/python/",
-        "input_directory": "data/langchain_md_files",
-        "output_file": "data/langchain_data.jsonl",
-        "source_name": "langchain",
-        "use_include_list": True,
-        "included_dirs": [
-            "concepts",
-            "langchain",
-            "python/integrations",
-            "python/migrate",
-            "python/releases",
-        ],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [
-            "security-policy.mdx",
-            "release-policy.mdx",
-            "versioning.mdx",
-        ],
-        "url_extension": "",
-    },
-    "tai_blog": {
-        "base_url": "",
-        "input_directory": "",
-        "output_file": "data/tai_blog_data.jsonl",
-        "source_name": "tai_blog",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "full_stack_ai_engineering": {
-        "base_url": "",
-        "input_directory": "data/full_stack_ai_engineering",  # Path to the directory that contains the Markdown files
-        "output_file": "data/full_stack_ai_engineering_data.jsonl",
-        "source_name": "full_stack_ai_engineering",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "beginner_python_for_ai_engineering": {
-        "base_url": "",
-        "input_directory": "data/beginner_python_for_ai_engineering",  # Path to the directory that contains the Markdown files
-        "output_file": "data/beginner_python_for_ai_engineering_data.jsonl",
-        "source_name": "beginner_python_for_ai_engineering",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "master_ai_for_work": {
-        "base_url": "",
-        "input_directory": "data/master_ai_for_work",  # Path to the directory that contains the Markdown files
-        "output_file": "data/master_ai_for_work_data.jsonl",
-        "source_name": "master_ai_for_work",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-    "agentic_ai_engineering": {
-        "base_url": "",
-        "input_directory": "data/agentic_ai_engineering",  # Path to the directory that contains the Markdown files
-        "output_file": "data/agentic_ai_engineering_data.jsonl",  # Agentic AI Engineering
-        "source_name": "agentic_ai_engineering",
-        "use_include_list": False,
-        "included_dirs": [],
-        "excluded_dirs": [],
-        "excluded_root_files": [],
-        "included_root_files": [],
-        "url_extension": "",
-    },
-}
 
+def split_frontmatter(content: str) -> tuple[str | None, str]:
+    lines = content.splitlines()
+    if not lines or lines[0].strip() != "---":
+        return None, content
 
-def extract_title(content: str):
-    title_match = re.search(r"^#\s+(.+)$", content, re.MULTILINE)
-    if title_match:
-        return title_match.group(1).strip()
+    for index, line in enumerate(lines[1:], start=1):
+        if line.strip() in {"---", "..."}:
+            frontmatter = "\n".join(lines[1:index])
+            body = "\n".join(lines[index + 1 :])
+            return frontmatter, body
+
+    return None, content
+
+
+def clean_frontmatter_scalar(value: str) -> str | None:
+    value = value.strip()
+    if not value or value in {"|", ">"}:
+        return None
+
+    if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
+        quote = value[0]
+        value = value[1:-1]
+        if quote == '"':
+            value = value.replace(r"\"", '"').replace(r"\\", "\\")
+        else:
+            value = value.replace("''", "'")
+
+    value = value.strip()
+    return value or None
+
+
+def extract_frontmatter_title(frontmatter: str) -> str | None:
+    for key in ("title", "sidebarTitle"):
+        title_match = re.search(rf"(?m)^\s*{key}\s*:\s*(.+?)\s*$", frontmatter)
+        if title_match:
+            title = clean_frontmatter_scalar(title_match.group(1))
+            if title:
+                return title
+
+    return None
+
+
+def iter_markdown_lines_outside_code_fences(content: str):
+    in_code_fence = False
+    fence_char = None
+
+    for line in content.splitlines():
+        fence_match = re.match(r"^\s{0,3}(```+|~~~+)", line)
+        if fence_match:
+            current_fence_char = fence_match.group(1)[0]
+            if not in_code_fence:
+                in_code_fence = True
+                fence_char = current_fence_char
+            elif current_fence_char == fence_char:
+                in_code_fence = False
+                fence_char = None
+            continue
+
+        if not in_code_fence:
+            yield line
 
-    lines = content.split("\n")
+
+def extract_title(content: str):
+    frontmatter, body = split_frontmatter(content)
+    if frontmatter:
+        title = extract_frontmatter_title(frontmatter)
+        if title:
+            return title
+
+    for line in iter_markdown_lines_outside_code_fences(body):
+        title_match = re.match(r"^\s{0,3}#\s+(.+)$", line)
+        if title_match:
+            return title_match.group(1).strip()
+
+    lines = body.split("\n")
     for line in lines:
         if line.strip():
             return line.strip()
@@ -209,16 +130,70 @@ def extract_title(content: str):
     return None
 
 
-def generate_url(file_path: str, config: Dict) -> str:
+def load_source_extension_manifest(directory: str) -> Dict[str, str]:
+    manifest_path = os.path.join(directory, "_source_extensions.json")
+    if not os.path.exists(manifest_path):
+        return {}
+
+    try:
+        with open(manifest_path, "r", encoding="utf-8") as f:
+            manifest = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        logger.warning("Could not load source extension manifest: %s", manifest_path)
+        return {}
+
+    return {
+        str(path): str(extension)
+        for path, extension in manifest.items()
+        if isinstance(path, str) and isinstance(extension, str)
+    }
+
+
+def load_source_url_manifest(directory: str) -> Dict[str, str]:
+    manifest_path = os.path.join(directory, "_source_urls.json")
+    if not os.path.exists(manifest_path):
+        return {}
+
+    try:
+        with open(manifest_path, "r", encoding="utf-8") as f:
+            manifest = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        logger.warning("Could not load source URL manifest: %s", manifest_path)
+        return {}
+
+    return {
+        str(path): str(url)
+        for path, url in manifest.items()
+        if isinstance(path, str) and isinstance(url, str)
+    }
+
+
+def generate_url(
+    file_path: str,
+    config: Dict,
+    source_extension: str | None = None,
+    source_url: str | None = None,
+) -> str:
     """
     Return an empty string if base_url is empty;
     otherwise return the constructed URL as before.
     """
+    if source_url:
+        return source_url
+
     if not config["base_url"]:
         return ""
 
-    path_without_extension = os.path.splitext(file_path)[0]
     source_name = config["source_name"]
+    path_with_forward_slashes = file_path.replace("\\", "/")
+
+    if config.get("preserve_file_extension_in_url"):
+        if source_extension:
+            path_without_extension = os.path.splitext(path_with_forward_slashes)[0]
+            return config["base_url"] + path_without_extension + source_extension
+        return config["base_url"] + path_with_forward_slashes
+
+    path_without_extension = os.path.splitext(file_path)[0]
 
     if source_name == "llama_index":
         framework_prefix = "src/content/docs/framework/"
@@ -261,8 +236,36 @@ def remove_copyright_header(content: str) -> str:
     return cleaned_content.strip()
 
 
+def remove_inline_base64_images(content: str) -> str:
+    content = re.sub(
+        r"!\[[^\]]*\]\(\s*data:image/[^,\s)]+;base64,[A-Za-z0-9+/=]+(?:\s+\"[^\"]*\")?\s*\)",
+        "[inline image omitted]",
+        content,
+    )
+    content = re.sub(
+        r"<img\b[^>]*\bsrc=[\"']data:image/[^,\s\"']+;base64,[A-Za-z0-9+/=]+[\"'][^>]*>",
+        "[inline image omitted]",
+        content,
+        flags=re.IGNORECASE,
+    )
+    content = re.sub(
+        r"data:image/[^,\s)\"']+;base64,[A-Za-z0-9+/=]+",
+        "[inline image omitted]",
+        content,
+    )
+    return content
+
+
+def clean_document_content(content: str) -> str:
+    content = remove_copyright_header(content)
+    content = remove_inline_base64_images(content)
+    return content.strip()
+
+
 def process_md_files(directory: str, config: Dict) -> List[Dict]:
     jsonl_data = []
+    source_extension_manifest = load_source_extension_manifest(directory)
+    source_url_manifest = load_source_url_manifest(directory)
 
     for root, _, files in os.walk(directory):
         for file in files:
@@ -274,8 +277,9 @@ def process_md_files(directory: str, config: Dict) -> List[Dict]:
                     with open(file_path, "r", encoding="utf-8") as f:
                         content = f.read()
 
-                    title = extract_title(content)
-                    token_count = num_tokens_from_string(content, "cl100k_base")
+                    cleaned_content = clean_document_content(content)
+                    title = extract_title(cleaned_content)
+                    token_count = num_tokens_from_string(cleaned_content, "cl100k_base")
 
                     # Skip very small or extremely large files
                     if token_count < 100 or token_count > 200_000:
@@ -284,13 +288,18 @@ def process_md_files(directory: str, config: Dict) -> List[Dict]:
                         )
                         continue
 
-                    cleaned_content = remove_copyright_header(content)
-
                     json_object = {
                         "tokens": token_count,
                         "doc_id": str(uuid.uuid4()),
                         "name": (title if title else file),
-                        "url": generate_url(relative_path, config),
+                        "url": generate_url(
+                            relative_path,
+                            config,
+                            source_extension_manifest.get(
+                                relative_path.replace("\\", "/")
+                            ),
+                            source_url_manifest.get(relative_path.replace("\\", "/")),
+                        ),
                         "retrieve_doc": (token_count <= 8000),
                         "source": config["source_name"],
                         "content": cleaned_content,

data/scraping_scripts/retire_source_workflow.py

@@ -5,8 +5,9 @@ Retire one or more sources from the AI Tutor data pipeline.
 This removes the source from:
 1. data/all_sources_data.jsonl
 2. data/all_sources_contextual_nodes.pkl
-3. the rebuilt Chroma vector store
-4. the private Hugging Face data repository's per-source JSONL file
+3. data/scraping_scripts/source_registry.py
+4. the rebuilt Chroma vector store
+5. the private Hugging Face data repository's per-source JSONL file
 
 Example:
     uv run -m data.scraping_scripts.retire_source_workflow --sources 8-hour_primer --yes
@@ -15,9 +16,11 @@ Example:
 from __future__ import annotations
 
 import argparse
+import ast
 import json
 import os
 import pickle
+import re
 import shutil
 import subprocess
 import sys
@@ -33,7 +36,7 @@ from data.scraping_scripts.hf_auth import HuggingFaceAuthError, validate_hf_acce
 from scripts.chroma_rag import get_chunk_record_source
 
 try:
-    from data.scraping_scripts.process_md_files import SOURCE_CONFIGS
+    from data.scraping_scripts.source_registry import SOURCE_CONFIGS
 except Exception:
     SOURCE_CONFIGS = {}
 
@@ -43,6 +46,7 @@ DATA_REPO_ID = "towardsai-tutors/ai-tutor-data"
 VECTOR_REPO_ID = "towardsai-tutors/ai-tutor-vector-db"
 ALL_SOURCES_JSONL = Path("data/all_sources_data.jsonl")
 CONTEXTUAL_NODES = Path("data/all_sources_contextual_nodes.pkl")
+SOURCE_REGISTRY = Path("data/scraping_scripts/source_registry.py")
 
 
 def run_module(module_name: str, *module_args: str) -> subprocess.CompletedProcess:
@@ -168,6 +172,104 @@ def delete_local_source_files(sources: list[str], dry_run: bool) -> None:
         print(f"Deleted local source file {source_path} after backup to {backup_path}")
 
 
+def find_source_config_block(text: str, source: str) -> tuple[int, int] | None:
+    marker = f'    "{source}": {{'
+    start = text.find(marker)
+    if start == -1:
+        return None
+
+    body_start = text.find("{", start)
+    if body_start == -1:
+        return None
+
+    depth = 0
+    quote: str | None = None
+    escaped = False
+    for index in range(body_start, len(text)):
+        char = text[index]
+        if quote:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == quote:
+                quote = None
+            continue
+
+        if char in {"'", '"'}:
+            quote = char
+        elif char == "{":
+            depth += 1
+        elif char == "}":
+            depth -= 1
+            if depth == 0:
+                end = index + 1
+                if end < len(text) and text[end] == ",":
+                    end += 1
+                if end < len(text) and text[end] == "\n":
+                    end += 1
+                return start, end
+
+    return None
+
+
+def remove_source_from_registry_text(text: str, source: str) -> tuple[str, bool]:
+    original = text
+    block = find_source_config_block(text, source)
+    if block is not None:
+        start, end = block
+        text = text[:start] + text[end:]
+
+    source_pattern = re.escape(source)
+    text = re.sub(
+        rf'^[ \t]*"{source_pattern}",\n',
+        "",
+        text,
+        flags=re.MULTILINE,
+    )
+    text = re.sub(
+        rf'^[ \t]*"{source_pattern}":\s*"[^"\n]*",\n',
+        "",
+        text,
+        flags=re.MULTILINE,
+    )
+    return text, text != original
+
+
+def update_source_registry(sources: list[str], dry_run: bool) -> None:
+    if not SOURCE_REGISTRY.exists():
+        print(f"Warning: {SOURCE_REGISTRY} not found; registry was not updated.")
+        return
+
+    original = SOURCE_REGISTRY.read_text(encoding="utf-8")
+    updated = original
+    changed_sources: list[str] = []
+
+    for source in sources:
+        updated, changed = remove_source_from_registry_text(updated, source)
+        if changed:
+            changed_sources.append(source)
+
+    if not changed_sources:
+        print(f"No retired sources were found in {SOURCE_REGISTRY}.")
+        return
+
+    if dry_run:
+        print(
+            "Would remove from source_registry.py: "
+            + ", ".join(sorted(changed_sources))
+        )
+        return
+
+    ast.parse(updated, filename=str(SOURCE_REGISTRY))
+    backup_path = backup_file(SOURCE_REGISTRY)
+    SOURCE_REGISTRY.write_text(updated, encoding="utf-8")
+    print(
+        f"Removed {', '.join(sorted(changed_sources))} from {SOURCE_REGISTRY} "
+        f"after backup to {backup_path}"
+    )
+
+
 def rebuild_vector_store() -> None:
     if not os.getenv("COHERE_API_KEY"):
         raise SystemExit("COHERE_API_KEY is required to rebuild the Chroma vector store.")
@@ -246,7 +348,7 @@ def main() -> None:
         nargs="*",
         help=(
             "Optional exact per-source JSONL filenames to delete from the data repo. "
-            "Defaults to SOURCE_CONFIGS output files or <source>_data.jsonl."
+            "Defaults to source_registry.py output files or <source>_data.jsonl."
         ),
     )
     parser.add_argument(
@@ -279,6 +381,11 @@ def main() -> None:
         action="store_true",
         help="Delete matching local per-source JSONL files after backing them up.",
     )
+    parser.add_argument(
+        "--keep-source-registry",
+        action="store_true",
+        help="Do not remove retired sources from data/scraping_scripts/source_registry.py.",
+    )
     parser.add_argument(
         "--dry-run",
         action="store_true",
@@ -333,6 +440,9 @@ def main() -> None:
     if args.delete_local_source_files:
         delete_local_source_files(args.sources, args.dry_run)
 
+    if not args.keep_source_registry:
+        update_source_registry(args.sources, args.dry_run)
+
     if args.dry_run:
         print("Dry run complete. No files were changed.")
         return

data/scraping_scripts/source_registry.py

@@ -0,0 +1,358 @@
+"""Central source registry for the AI Tutor knowledge base.
+
+Sources listed here are active in the KB pipeline. To retire a source, run
+``retire_source_workflow.py``; confirmed retirements remove sources from this
+file automatically.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Iterable
+
+ALL_SOURCES_JSONL = "data/all_sources_data.jsonl"
+CONTEXTUAL_NODES_PKL = "data/all_sources_contextual_nodes.pkl"
+
+
+SOURCE_CONFIGS: dict[str, dict[str, Any]] = {
+    "transformers": {
+        "base_url": "https://huggingface.co/docs/transformers/",
+        "input_directory": "data/transformers_md_files",
+        "output_file": "data/transformers_data.jsonl",
+        "source_name": "transformers",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": ["internal"],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "peft": {
+        "base_url": "https://huggingface.co/docs/peft/",
+        "input_directory": "data/peft_md_files",
+        "output_file": "data/peft_data.jsonl",
+        "source_name": "peft",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "trl": {
+        "base_url": "https://huggingface.co/docs/trl/",
+        "input_directory": "data/trl_md_files",
+        "output_file": "data/trl_data.jsonl",
+        "source_name": "trl",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "llama_index": {
+        "base_url": "https://docs.llamaindex.ai/en/stable/",
+        "input_directory": "data/llama_index_md_files",
+        "output_file": "data/llama_index_data.jsonl",
+        "source_name": "llama_index",
+        "use_include_list": True,
+        "included_dirs": [
+            "src/content/docs/framework/index.md",
+            "src/content/docs/framework/getting_started",
+            "src/content/docs/framework/understanding",
+            "src/content/docs/framework/use_cases",
+            "src/content/docs/framework/module_guides",
+            "src/content/docs/framework/optimizing",
+            "src/content/docs/framework/community/faq",
+            "src/content/docs/framework/community/integrations",
+            "src/content/docs/framework/llama_cloud",
+            "examples",
+        ],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "langchain": {
+        "base_url": "https://docs.langchain.com/oss/python/",
+        "input_directory": "data/langchain_md_files",
+        "output_file": "data/langchain_data.jsonl",
+        "source_name": "langchain",
+        "use_include_list": True,
+        "included_dirs": [
+            "concepts",
+            "langchain",
+            "python/integrations/chat/",
+            "python/integrations/document_loaders/",
+            "python/integrations/document_transformers/",
+            "python/integrations/embeddings/",
+            "python/integrations/retrievers/",
+            "python/integrations/splitters/",
+            "python/integrations/stores/",
+            "python/integrations/tools/",
+            "python/integrations/vectorstores/",
+            "python/migrate",
+            "python/releases",
+        ],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [
+            "security-policy.mdx",
+            "release-policy.mdx",
+            "versioning.mdx",
+        ],
+        "url_extension": "",
+    },
+    "langgraph": {
+        "base_url": "https://docs.langchain.com/oss/python/langgraph/",
+        "input_directory": "data/langgraph_md_files",
+        "output_file": "data/langgraph_data.jsonl",
+        "source_name": "langgraph",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "deep_agents": {
+        "base_url": "https://docs.langchain.com/oss/python/deepagents/",
+        "input_directory": "data/deep_agents_md_files",
+        "output_file": "data/deep_agents_data.jsonl",
+        "source_name": "deep_agents",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "openai_docs": {
+        "base_url": "https://developers.openai.com/",
+        "input_directory": "data/openai_docs_md_files",
+        "output_file": "data/openai_docs_data.jsonl",
+        "source_name": "openai_docs",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+        "llms_txt_urls": [
+            "https://developers.openai.com/api/docs/llms.txt",
+            "https://developers.openai.com/codex/llms.txt",
+        ],
+        "llms_url_include_prefixes": [
+            "https://developers.openai.com/api/docs/",
+            "https://developers.openai.com/codex/",
+        ],
+    },
+    "claude_code_docs": {
+        "base_url": "https://code.claude.com/docs/",
+        "input_directory": "data/claude_code_docs_md_files",
+        "output_file": "data/claude_code_docs_data.jsonl",
+        "source_name": "claude_code_docs",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+        "llms_txt_urls": [
+            "https://code.claude.com/docs/llms.txt",
+        ],
+        "llms_url_include_prefixes": [
+            "https://code.claude.com/docs/en/",
+        ],
+    },
+    "full_stack_ai_engineering": {
+        "base_url": "",
+        "input_directory": "data/full_stack_ai_engineering",
+        "output_file": "data/full_stack_ai_engineering_data.jsonl",
+        "source_name": "full_stack_ai_engineering",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "beginner_python_for_ai_engineering": {
+        "base_url": "",
+        "input_directory": "data/beginner_python_for_ai_engineering",
+        "output_file": "data/beginner_python_for_ai_engineering_data.jsonl",
+        "source_name": "beginner_python_for_ai_engineering",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "master_ai_for_work": {
+        "base_url": "",
+        "input_directory": "data/master_ai_for_work",
+        "output_file": "data/master_ai_for_work_data.jsonl",
+        "source_name": "master_ai_for_work",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "agentic_ai_engineering": {
+        "base_url": "",
+        "input_directory": "data/agentic_ai_engineering",
+        "output_file": "data/agentic_ai_engineering_data.jsonl",
+        "source_name": "agentic_ai_engineering",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+}
+
+DOC_SOURCE_KEYS = (
+    "transformers",
+    "peft",
+    "trl",
+    "llama_index",
+    "langchain",
+    "langgraph",
+    "deep_agents",
+    "openai_docs",
+    "claude_code_docs",
+)
+GITHUB_SOURCE_KEYS = (
+    "transformers",
+    "peft",
+    "trl",
+    "llama_index",
+    "langchain",
+    "langgraph",
+    "deep_agents",
+)
+LLMS_TXT_SOURCE_KEYS = (
+    "openai_docs",
+    "claude_code_docs",
+)
+COURSE_SOURCE_KEYS = frozenset(
+    {
+        "full_stack_ai_engineering",
+        "beginner_python_for_ai_engineering",
+        "master_ai_for_work",
+        "agentic_ai_engineering",
+    }
+)
+
+ACTIVE_SOURCE_KEYS = frozenset(SOURCE_CONFIGS.keys())
+AVAILABLE_SOURCES = list(SOURCE_CONFIGS.keys())
+
+SOURCE_KEY_TO_LABEL = {
+    "transformers": "Transformers Docs",
+    "peft": "PEFT Docs",
+    "trl": "TRL Docs",
+    "llama_index": "LlamaIndex Docs",
+    "langchain": "LangChain Docs",
+    "langgraph": "LangGraph Docs",
+    "deep_agents": "Deep Agents Docs",
+    "openai_docs": "OpenAI Docs",
+    "claude_code_docs": "Claude Code Docs",
+    "full_stack_ai_engineering": "Full Stack AI Engineering",
+    "beginner_python_for_ai_engineering": "Beginner Python for AI Engineering",
+    "master_ai_for_work": "Master AI For Work",
+    "agentic_ai_engineering": "Agentic AI Engineering",
+}
+
+UI_SOURCE_KEYS = (
+    "openai_docs",
+    "claude_code_docs",
+    "langgraph",
+    "deep_agents",
+    "langchain",
+    "llama_index",
+    "transformers",
+    "peft",
+    "trl",
+    "agentic_ai_engineering",
+    "master_ai_for_work",
+    "full_stack_ai_engineering",
+    "beginner_python_for_ai_engineering",
+)
+SOURCE_UI_TO_KEY = {SOURCE_KEY_TO_LABEL[key]: key for key in UI_SOURCE_KEYS}
+AVAILABLE_SOURCES_UI = list(SOURCE_UI_TO_KEY.keys())
+
+DEFAULT_SELECTED_SOURCE_KEYS = (
+    "agentic_ai_engineering",
+    "master_ai_for_work",
+    "full_stack_ai_engineering",
+    "beginner_python_for_ai_engineering",
+    "openai_docs",
+    "claude_code_docs",
+    "langgraph",
+    "deep_agents",
+    "transformers",
+    "peft",
+    "trl",
+    "llama_index",
+    "langchain",
+)
+DEFAULT_SELECTED_SOURCES_UI = [
+    SOURCE_KEY_TO_LABEL[key] for key in DEFAULT_SELECTED_SOURCE_KEYS
+]
+
+
+def aggregate_data_files() -> dict[str, str]:
+    return {
+        ALL_SOURCES_JSONL: Path(ALL_SOURCES_JSONL).name,
+        CONTEXTUAL_NODES_PKL: Path(CONTEXTUAL_NODES_PKL).name,
+    }
+
+
+def source_data_files() -> dict[str, str]:
+    return {
+        str(config["output_file"]): Path(str(config["output_file"])).name
+        for config in SOURCE_CONFIGS.values()
+    }
+
+
+def source_output_files(sources: Iterable[str]) -> set[str]:
+    return {
+        str(SOURCE_CONFIGS[source]["output_file"])
+        for source in sources
+        if source in SOURCE_CONFIGS
+    }
+
+
+def required_data_files(*, include_source_files: bool = True) -> dict[str, str]:
+    files = aggregate_data_files()
+    if include_source_files:
+        files.update(source_data_files())
+    return files
+
+
+def upload_data_file_paths() -> list[str]:
+    return list(required_data_files().keys())
+
+
+def vector_store_source_configs() -> dict[str, dict[str, str]]:
+    configs = {
+        source: {
+            "input_file": str(config["output_file"]),
+            "db_name": f"chroma-db-{source}",
+            "document_dict_file": f"document_dict_{source}.pkl",
+            "bm25_index_file": f"bm25_index_{source}.pkl",
+        }
+        for source, config in SOURCE_CONFIGS.items()
+    }
+    configs["all_sources"] = {
+        "input_file": ALL_SOURCES_JSONL,
+        "db_name": "chroma-db-all_sources",
+        "document_dict_file": "document_dict_all_sources.pkl",
+        "bm25_index_file": "bm25_index_all_sources.pkl",
+    }
+    return configs

data/scraping_scripts/update_docs_workflow.py

@@ -2,14 +2,14 @@
 """
 AI Tutor App - Documentation Update Workflow
 
-This script automates the process of updating documentation from GitHub repositories:
-1. Download documentation from GitHub using the API
+This script automates the process of updating documentation sources:
+1. Download documentation from GitHub or official llms.txt indexes
 2. Process markdown files to create JSONL data
 3. Add contextual information to document nodes
 4. Create vector stores
 5. Upload databases to HuggingFace
 
-This workflow is specific to updating library documentation (Transformers, PEFT, LlamaIndex, etc.).
+This workflow is specific to updating library documentation (Transformers, PEFT, LlamaIndex, OpenAI docs, etc.).
 For adding courses, use the add_course_workflow.py script instead.
 
 Usage:
@@ -36,7 +36,17 @@ from typing import Dict, List, Set
 from dotenv import load_dotenv
 from huggingface_hub import hf_hub_download
 
+from data.scraping_scripts.contextual_node_pruning import (
+    prune_contextual_nodes_to_active_sources,
+)
 from data.scraping_scripts.hf_auth import HuggingFaceAuthError, validate_hf_access
+from data.scraping_scripts.source_registry import (
+    DOC_SOURCE_KEYS,
+    GITHUB_SOURCE_KEYS,
+    LLMS_TXT_SOURCE_KEYS,
+    required_data_files,
+    source_output_files,
+)
 from scripts.chroma_rag import get_chunk_record_doc_id
 
 # Load environment variables from .env file
@@ -62,27 +72,10 @@ def ensure_hf_access() -> None:
         sys.exit(1)
 
 
-def ensure_required_files_exist():
+def ensure_required_files_exist(sources_to_regenerate: List[str] | None = None):
     """Download required data files from HuggingFace if they don't exist locally."""
-    # List of files to check and download
-    required_files = {
-        # Critical files
-        "data/all_sources_data.jsonl": "all_sources_data.jsonl",
-        "data/all_sources_contextual_nodes.pkl": "all_sources_contextual_nodes.pkl",
-        # Documentation source files
-        "data/transformers_data.jsonl": "transformers_data.jsonl",
-        "data/peft_data.jsonl": "peft_data.jsonl",
-        "data/trl_data.jsonl": "trl_data.jsonl",
-        "data/llama_index_data.jsonl": "llama_index_data.jsonl",
-        "data/langchain_data.jsonl": "langchain_data.jsonl",
-        "data/openai_cookbooks_data.jsonl": "openai_cookbooks_data.jsonl",
-        # Course files
-        "data/tai_blog_data.jsonl": "tai_blog_data.jsonl",
-        "data/master_ai_for_work_data.jsonl": "master_ai_for_work_data.jsonl",
-        "data/agentic_ai_engineering_data.jsonl": "agentic_ai_engineering_data.jsonl",
-        "data/full_stack_ai_engineering_data.jsonl": "full_stack_ai_engineering_data.jsonl",
-        "data/beginner_python_for_ai_engineering_data.jsonl": "beginner_python_for_ai_engineering_data.jsonl",
-    }
+    required_files = required_data_files()
+    regenerated_source_files = source_output_files(sources_to_regenerate or [])
 
     # Critical files that must be downloaded
     critical_files = [
@@ -92,6 +85,14 @@ def ensure_required_files_exist():
 
     # Check and download each file
     for local_path, remote_filename in required_files.items():
+        if local_path in regenerated_source_files:
+            if not os.path.exists(local_path):
+                logger.info(
+                    "%s will be regenerated for this run; skipping HuggingFace download",
+                    remote_filename,
+                )
+            continue
+
         if not os.path.exists(local_path):
             logger.info(
                 f"{remote_filename} not found. Attempting to download from HuggingFace..."
@@ -129,15 +130,10 @@ def ensure_required_files_exist():
                         )
 
 
-# Documentation sources that can be updated via GitHub API
-GITHUB_SOURCES = [
-    "transformers",
-    "peft",
-    "trl",
-    "llama_index",
-    "openai_cookbooks",
-    "langchain",
-]
+# Documentation sources that can be updated automatically
+DOC_SOURCES = list(DOC_SOURCE_KEYS)
+GITHUB_SOURCES = list(GITHUB_SOURCE_KEYS)
+LLMS_TXT_SOURCES = list(LLMS_TXT_SOURCE_KEYS)
 
 
 def load_jsonl(file_path: str) -> List[Dict]:
@@ -178,6 +174,46 @@ def download_from_github(sources: List[str]) -> None:
         logger.info(f"Successfully downloaded {source} documentation")
 
 
+def download_from_llms_txt(sources: List[str]) -> None:
+    """Download documentation from llms.txt indexes."""
+    logger.info(
+        f"Downloading documentation from llms.txt indexes for sources: {sources}"
+    )
+
+    for source in sources:
+        if source not in LLMS_TXT_SOURCES:
+            logger.warning(
+                f"Source {source} is not an llms.txt source, skipping download"
+            )
+            continue
+
+        logger.info(f"Downloading {source} documentation")
+        result = run_module("data.scraping_scripts.llms_txt_to_markdown_docs", source)
+
+        if result.returncode != 0:
+            logger.error(
+                f"Error downloading {source} documentation. Stopping workflow to avoid overwriting source JSONL files with incomplete data."
+            )
+            sys.exit(1)
+
+        logger.info(f"Successfully downloaded {source} documentation")
+
+
+def download_documentation(sources: List[str]) -> None:
+    """Download docs with the right source-specific downloader."""
+    github_sources = [source for source in sources if source in GITHUB_SOURCES]
+    llms_txt_sources = [source for source in sources if source in LLMS_TXT_SOURCES]
+
+    if github_sources:
+        download_from_github(github_sources)
+    if llms_txt_sources:
+        download_from_llms_txt(llms_txt_sources)
+
+    unsupported = sorted(set(sources) - set(github_sources) - set(llms_txt_sources))
+    for source in unsupported:
+        logger.warning(f"Source {source} is not a downloadable docs source, skipping")
+
+
 def capture_source_versions(sources: List[str]) -> None:
     """Record latest release tag + SHA + indexed date per source."""
     logger.info(f"Capturing source versions for: {sources}")
@@ -202,6 +238,12 @@ def process_markdown_files(sources: List[str]) -> None:
         logger.error("Error processing markdown files - check output above")
         sys.exit(1)
 
+    if len(sources) == 1:
+        from data.scraping_scripts.process_md_files import combine_all_sources
+
+        logger.info("Rebuilding all_sources_data.jsonl after single-source update")
+        combine_all_sources(sources)
+
     logger.info("Successfully processed markdown files")
 
 
@@ -367,12 +409,12 @@ def main():
     parser.add_argument(
         "--sources",
         nargs="+",
-        choices=GITHUB_SOURCES,
-        default=GITHUB_SOURCES,
-        help="GitHub documentation sources to update",
+        choices=DOC_SOURCES,
+        default=DOC_SOURCES,
+        help="Documentation sources to update",
     )
     parser.add_argument(
-        "--skip-download", action="store_true", help="Skip downloading from GitHub"
+        "--skip-download", action="store_true", help="Skip downloading source docs"
     )
     parser.add_argument(
         "--skip-process", action="store_true", help="Skip processing markdown files"
@@ -403,12 +445,14 @@ def main():
 
     ensure_hf_access()
 
-    # Ensure required data files exist before proceeding
-    ensure_required_files_exist()
+    # Keep untouched source JSONLs by downloading them when needed, but don't
+    # require first-time sources that this run is about to regenerate.
+    sources_to_regenerate = [] if args.skip_process else args.sources
+    ensure_required_files_exist(sources_to_regenerate=sources_to_regenerate)
 
     # Execute the workflow steps
     if not args.skip_download:
-        download_from_github(args.sources)
+        download_documentation(args.sources)
         capture_source_versions(args.sources)
 
     if not args.skip_process:
@@ -417,6 +461,8 @@ def main():
     if not args.skip_context:
         add_context_to_nodes(not args.process_all_context)
 
+    prune_contextual_nodes_to_active_sources()
+
     if not args.skip_vectors:
         create_vector_stores()
 

data/scraping_scripts/upload_data_to_hf.py

@@ -22,33 +22,17 @@ from dotenv import load_dotenv
 
 try:
     from data.scraping_scripts.hf_auth import HuggingFaceAuthError, validate_hf_access
+    from data.scraping_scripts.source_registry import upload_data_file_paths
 except ModuleNotFoundError:
     from hf_auth import HuggingFaceAuthError, validate_hf_access
+    from source_registry import upload_data_file_paths
 
 load_dotenv()
 
 
 def upload_files_to_huggingface(repo_id="towardsai-tutors/ai-tutor-data"):
     """Upload data files to a private HuggingFace repository."""
-    # Main files to upload
-    files_to_upload = [
-        # Combined data and vector store
-        "data/all_sources_data.jsonl",
-        "data/all_sources_contextual_nodes.pkl",
-        # Individual source files
-        "data/transformers_data.jsonl",
-        "data/peft_data.jsonl",
-        "data/trl_data.jsonl",
-        "data/llama_index_data.jsonl",
-        "data/langchain_data.jsonl",
-        "data/openai_cookbooks_data.jsonl",
-        # Course files
-        "data/tai_blog_data.jsonl",
-        "data/master_ai_for_work_data.jsonl",
-        "data/agentic_ai_engineering_data.jsonl",
-        "data/full_stack_ai_engineering_data.jsonl",
-        "data/beginner_python_for_ai_engineering_data.jsonl",
-    ]
+    files_to_upload = upload_data_file_paths()
 
     # Filter to only include files that exist
     existing_files = []

frontend/lib/doc-metadata.ts

@@ -29,9 +29,29 @@ export const DOC_METADATA: Record<string, DocMetadata> = {
       "Framework for building LLM apps: chains, agents, tool-calling, and production observability.",
     docsUrl: "https://docs.langchain.com/oss/python/langchain/overview",
   },
+  langgraph: {
+    description:
+      "Graph-based runtime for reliable, stateful AI agents with persistence, streaming, human review, and deployment patterns.",
+    docsUrl: "https://docs.langchain.com/oss/python/langgraph/overview",
+  },
+  deep_agents: {
+    description:
+      "LangChain's deep agent harness for planning, delegation, filesystem context, and longer-running agent workflows.",
+    docsUrl: "https://docs.langchain.com/oss/python/deepagents/overview",
+  },
   openai_cookbooks: {
     description:
       "Example notebooks and recipes from OpenAI covering practical patterns for using their APIs.",
     docsUrl: "https://cookbook.openai.com",
   },
+  openai_docs: {
+    description:
+      "Official OpenAI API, Agents SDK, and Codex documentation from the developer docs Markdown index.",
+    docsUrl: "https://developers.openai.com",
+  },
+  claude_code_docs: {
+    description:
+      "Official Claude Code and Claude Agent SDK documentation from Anthropic's Markdown index.",
+    docsUrl: "https://code.claude.com/docs/en/overview",
+  },
 };

scripts/chat_service.py

@@ -19,6 +19,7 @@ from .chat_types import ChatEvent, ChatRequest, ChatTurn, SourceMatch
 from .chroma_rag import LocalChromaRetriever, format_tool_payload, parse_tool_payload
 from .prompts import build_system_prompt
 from .setup import (
+    BM25_INDEX_PATH,
     COURSE_SOURCE_KEYS,
     DOCUMENT_DICT_PATH,
     SOURCE_KEY_TO_LABEL,
@@ -51,6 +52,7 @@ def get_retriever() -> LocalChromaRetriever:
         db_path=VECTOR_DB_DIR,
         collection_name=VECTOR_COLLECTION_NAME,
         document_dict_path=DOCUMENT_DICT_PATH,
+        bm25_index_path=BM25_INDEX_PATH,
         cohere_api_key=cohere_api_key,
     )
 

scripts/chroma_rag.py

@@ -2,7 +2,13 @@ from __future__ import annotations
 
 import json
 import math
+import os
 import pickle
+import random
+import re
+import threading
+import time
+from collections import Counter, deque
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, Iterable
@@ -17,13 +23,100 @@ from tqdm.auto import tqdm
 
 DEFAULT_CHUNK_SIZE = 800
 DEFAULT_CHUNK_OVERLAP = 100
+DEFAULT_MAX_CHUNK_TOKENS = 1200
 DEFAULT_DENSE_TOP_K = 15
+DEFAULT_BM25_TOP_K = 30
+DEFAULT_FUSION_TOP_K = 30
 DEFAULT_RERANK_TOP_K = 5
+DEFAULT_RRF_K = 60
 DEFAULT_CONTEXT_TOKEN_BUDGET = 100_000
 DEFAULT_EMBED_MODEL = "embed-v4.0"
 DEFAULT_RERANK_MODEL = "rerank-v4.0-fast"
 DEFAULT_ENCODING = "cl100k_base"
 DEFAULT_OUTPUT_DIMENSION = 1024
+DEFAULT_COHERE_EMBED_BATCH_SIZE = 96
+DEFAULT_COHERE_EMBED_INPUTS_PER_MINUTE = 2_000
+DEFAULT_COHERE_EMBED_TPM_LIMIT = 0
+DEFAULT_COHERE_EMBED_RPM_LIMIT = 0
+DEFAULT_COHERE_EMBED_RATE_LIMIT_MARGIN = 0.8
+DEFAULT_COHERE_EMBED_WINDOW_SECONDS = 60.0
+DEFAULT_COHERE_EMBED_RETRY_ATTEMPTS = 8
+DEFAULT_SOURCE_VERSIONS_PATH = "data/source_versions.json"
+
+BM25_TOKEN_RE = re.compile(r"[A-Za-z][A-Za-z0-9_./:-]*|\d+(?:\.\d+)*")
+CAMEL_CASE_RE = re.compile(r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z]|$)|\d+")
+MARKDOWN_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*#*\s*$")
+MARKDOWN_FENCE_RE = re.compile(r"^\s*(`{3,}|~{3,})")
+BM25_STOP_WORDS = frozenset(
+    {
+        "a",
+        "an",
+        "and",
+        "are",
+        "as",
+        "at",
+        "be",
+        "by",
+        "for",
+        "from",
+        "how",
+        "i",
+        "in",
+        "is",
+        "it",
+        "of",
+        "on",
+        "or",
+        "that",
+        "the",
+        "this",
+        "to",
+        "use",
+        "what",
+        "when",
+        "where",
+        "which",
+        "with",
+        "you",
+        "your",
+    }
+)
+
+
+class SyncWindowLimiter:
+    def __init__(self, units_per_window: int, window_seconds: float) -> None:
+        self.units_per_window = max(1, units_per_window)
+        self.window_seconds = window_seconds
+        self._events: deque[tuple[float, int]] = deque()
+        self._used_units = 0
+        self._lock = threading.Lock()
+
+    def acquire(self, units: int) -> None:
+        units = min(max(1, units), self.units_per_window)
+
+        while True:
+            with self._lock:
+                now = time.monotonic()
+                self._prune(now)
+
+                if self._used_units + units <= self.units_per_window:
+                    self._events.append((now, units))
+                    self._used_units += units
+                    return
+
+                oldest_at, _ = self._events[0]
+                delay = max(0.1, self.window_seconds - (now - oldest_at))
+
+            time.sleep(delay + random.uniform(0.1, 0.75))
+
+    def _prune(self, now: float) -> None:
+        while self._events and now - self._events[0][0] >= self.window_seconds:
+            _, units = self._events.popleft()
+            self._used_units -= units
+
+
+_cohere_limiter_lock = threading.Lock()
+_cohere_limiters: dict[tuple[str, int, float], SyncWindowLimiter] = {}
 
 
 @dataclass(slots=True)
@@ -34,6 +127,20 @@ class ChunkRecord:
     metadata: dict[str, Any]
 
 
+@dataclass(slots=True)
+class MarkdownUnit:
+    text: str
+    heading_path: tuple[str, ...]
+    is_code: bool = False
+
+
+@dataclass(slots=True)
+class MarkdownChunk:
+    text: str
+    heading_path: tuple[str, ...]
+    tokens: int
+
+
 @dataclass(slots=True)
 class SearchResult:
     chunk_id: str
@@ -46,6 +153,108 @@ class SearchResult:
     score: float
     content: str
     chunk_content: str
+    heading_path: str = ""
+    retrieval_method: str = ""
+
+
+@dataclass(slots=True)
+class BM25Index:
+    records: list[ChunkRecord]
+    postings: dict[str, list[tuple[int, int]]]
+    document_frequencies: dict[str, int]
+    document_lengths: list[int]
+    average_document_length: float
+    k1: float = 1.5
+    b: float = 0.75
+
+    @classmethod
+    def build(
+        cls,
+        records: list[ChunkRecord],
+        *,
+        k1: float = 1.5,
+        b: float = 0.75,
+    ) -> "BM25Index":
+        postings: dict[str, list[tuple[int, int]]] = {}
+        document_frequencies: dict[str, int] = {}
+        document_lengths: list[int] = []
+
+        for doc_index, record in enumerate(records):
+            terms = tokenize_for_bm25(
+                format_chunk_for_retrieval(record.text, record.metadata)
+            )
+            counts = Counter(terms)
+            document_lengths.append(sum(counts.values()))
+
+            for term, term_frequency in counts.items():
+                postings.setdefault(term, []).append((doc_index, term_frequency))
+            for term in counts:
+                document_frequencies[term] = document_frequencies.get(term, 0) + 1
+
+        average_document_length = (
+            sum(document_lengths) / len(document_lengths) if document_lengths else 0.0
+        )
+        return cls(
+            records=records,
+            postings=postings,
+            document_frequencies=document_frequencies,
+            document_lengths=document_lengths,
+            average_document_length=average_document_length,
+            k1=k1,
+            b=b,
+        )
+
+    def search(
+        self,
+        query: str,
+        *,
+        allowed_sources: list[str] | None = None,
+        top_k: int = DEFAULT_BM25_TOP_K,
+    ) -> list[tuple[ChunkRecord, float]]:
+        query_terms = tokenize_for_bm25(query)
+        if not query_terms or not self.records:
+            return []
+
+        allowed = set(allowed_sources or [])
+        total_documents = len(self.records)
+        scores: dict[int, float] = {}
+
+        for term in set(query_terms):
+            postings = self.postings.get(term)
+            if not postings:
+                continue
+
+            document_frequency = self.document_frequencies.get(term, len(postings))
+            idf = math.log(
+                1.0
+                + (total_documents - document_frequency + 0.5)
+                / (document_frequency + 0.5)
+            )
+
+            for doc_index, term_frequency in postings:
+                record = self.records[doc_index]
+                if allowed and record.metadata.get("source") not in allowed:
+                    continue
+
+                document_length = self.document_lengths[doc_index]
+                if document_length <= 0 or self.average_document_length <= 0:
+                    continue
+
+                denominator = term_frequency + self.k1 * (
+                    1.0
+                    - self.b
+                    + self.b * document_length / self.average_document_length
+                )
+                scores[doc_index] = scores.get(doc_index, 0.0) + idf * (
+                    term_frequency * (self.k1 + 1.0) / denominator
+                )
+
+        ranked = sorted(scores.items(), key=lambda item: item[1], reverse=True)
+        return [
+            (self.records[doc_index], score)
+            for doc_index, score in ranked[: max(0, top_k)]
+            if score > 0.0
+        ]
 
 
 def batched(items: list[Any], size: int) -> Iterable[list[Any]]:
@@ -71,6 +280,133 @@ def get_token_encoding(model_name: str | None = None) -> tiktoken.Encoding:
     return tiktoken.get_encoding(DEFAULT_ENCODING)
 
 
+def load_source_versions(
+    path: str = DEFAULT_SOURCE_VERSIONS_PATH,
+) -> dict[str, dict[str, Any]]:
+    if not os.path.exists(path):
+        return {}
+    try:
+        with open(path, "r", encoding="utf-8") as handle:
+            data = json.load(handle)
+    except (OSError, json.JSONDecodeError):
+        return {}
+    if not isinstance(data, dict):
+        return {}
+    return {
+        str(source): dict(metadata)
+        for source, metadata in data.items()
+        if isinstance(metadata, dict)
+    }
+
+
+def source_version_for(
+    source: str,
+    source_versions: dict[str, dict[str, Any]] | None = None,
+) -> str:
+    metadata = (source_versions or {}).get(source, {})
+    for key in ("version", "sha", "indexedAt"):
+        value = metadata.get(key)
+        if value:
+            return str(value)
+    return ""
+
+
+def clean_heading_text(text: str) -> str:
+    return re.sub(r"\s+", " ", text.strip().strip("#").strip())
+
+
+def parse_markdown_units(
+    text: str,
+    *,
+    default_heading_path: tuple[str, ...] = (),
+) -> list[MarkdownUnit]:
+    lines = text.splitlines()
+    units: list[MarkdownUnit] = []
+    heading_stack: list[str] = []
+    text_buffer: list[str] = []
+
+    def active_heading_path() -> tuple[str, ...]:
+        return tuple(heading_stack) or default_heading_path
+
+    def flush_text_buffer() -> None:
+        nonlocal text_buffer
+        paragraph: list[str] = []
+        for buffered_line in text_buffer:
+            if buffered_line.strip():
+                paragraph.append(buffered_line)
+                continue
+            if paragraph:
+                units.append(
+                    MarkdownUnit(
+                        text="\n".join(paragraph).strip(),
+                        heading_path=active_heading_path(),
+                    )
+                )
+                paragraph = []
+        if paragraph:
+            units.append(
+                MarkdownUnit(
+                    text="\n".join(paragraph).strip(),
+                    heading_path=active_heading_path(),
+                )
+            )
+        text_buffer = []
+
+    index = 0
+    while index < len(lines):
+        line = lines[index]
+        fence_match = MARKDOWN_FENCE_RE.match(line)
+        if fence_match:
+            flush_text_buffer()
+            fence = fence_match.group(1)
+            fence_char = fence[0]
+            fence_len = len(fence)
+            code_lines = [line]
+            index += 1
+            while index < len(lines):
+                code_line = lines[index]
+                code_lines.append(code_line)
+                close_match = MARKDOWN_FENCE_RE.match(code_line)
+                if (
+                    close_match
+                    and close_match.group(1)[0] == fence_char
+                    and len(close_match.group(1)) >= fence_len
+                ):
+                    index += 1
+                    break
+                index += 1
+            units.append(
+                MarkdownUnit(
+                    text="\n".join(code_lines).strip(),
+                    heading_path=active_heading_path(),
+                    is_code=True,
+                )
+            )
+            continue
+
+        heading_match = MARKDOWN_HEADING_RE.match(line)
+        if heading_match:
+            flush_text_buffer()
+            level = len(heading_match.group(1))
+            title = clean_heading_text(heading_match.group(2))
+            heading_stack = heading_stack[: level - 1]
+            heading_stack.append(title)
+            units.append(
+                MarkdownUnit(
+                    text=line.strip(),
+                    heading_path=active_heading_path(),
+                )
+            )
+            index += 1
+            continue
+
+        text_buffer.append(line)
+        index += 1
+
+    flush_text_buffer()
+    return [unit for unit in units if unit.text.strip()]
+
+
 def token_window_chunks(
     text: str,
     chunk_size: int = DEFAULT_CHUNK_SIZE,
@@ -96,33 +432,166 @@ def token_window_chunks(
     return chunks
 
 
+def _heading_path_text(heading_path: tuple[str, ...] | list[str] | str | None) -> str:
+    if not heading_path:
+        return ""
+    if isinstance(heading_path, str):
+        return heading_path
+    return " > ".join(str(part) for part in heading_path if str(part).strip())
+
+
+def _split_large_unit(
+    unit: MarkdownUnit,
+    *,
+    encoding: tiktoken.Encoding,
+    chunk_size: int,
+    chunk_overlap: int,
+) -> list[MarkdownUnit]:
+    if unit.is_code:
+        return [unit]
+
+    token_count = len(encoding.encode(unit.text, disallowed_special=()))
+    if token_count <= chunk_size:
+        return [unit]
+    return [
+        MarkdownUnit(text=chunk, heading_path=unit.heading_path)
+        for chunk in token_window_chunks(
+            unit.text,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+    ]
+
+
+def heading_aware_markdown_chunks(
+    text: str,
+    *,
+    title: str = "",
+    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
+    encoding_name: str = DEFAULT_ENCODING,
+) -> list[MarkdownChunk]:
+    encoding = tiktoken.get_encoding(encoding_name)
+    default_heading_path = (title,) if title else ()
+    units = parse_markdown_units(text, default_heading_path=default_heading_path)
+    if not units:
+        return []
+
+    budget_size = min(max(chunk_size, 1), DEFAULT_MAX_CHUNK_TOKENS)
+    overlap_size = chunk_overlap
+
+    def unit_size(value: str) -> int:
+        return len(encoding.encode(value, disallowed_special=()))
+
+    chunks: list[MarkdownChunk] = []
+    current_parts: list[str] = []
+    current_heading_path: tuple[str, ...] = ()
+    current_size = 0
+
+    def flush_current() -> None:
+        nonlocal current_parts, current_heading_path, current_size
+        text_value = "\n\n".join(part for part in current_parts if part.strip()).strip()
+        if text_value:
+            chunks.append(
+                MarkdownChunk(
+                    text=text_value,
+                    heading_path=current_heading_path,
+                    tokens=len(encoding.encode(text_value, disallowed_special=())),
+                )
+            )
+        current_parts = []
+        current_heading_path = ()
+        current_size = 0
+
+    for unit in units:
+        split_units = _split_large_unit(
+            unit,
+            encoding=encoding,
+            chunk_size=budget_size,
+            chunk_overlap=overlap_size,
+        )
+        for split_unit in split_units:
+            size = unit_size(split_unit.text)
+            heading_changed = (
+                current_heading_path
+                and split_unit.heading_path != current_heading_path
+            )
+            would_exceed = current_parts and current_size + size > budget_size
+            if heading_changed or would_exceed:
+                flush_current()
+
+            if not current_parts:
+                current_heading_path = split_unit.heading_path
+            current_parts.append(split_unit.text)
+            current_size += size
+
+    flush_current()
+    return chunks
+
+
+def build_chunk_retrieval_header(metadata: dict[str, Any]) -> str:
+    lines: list[str] = []
+    title = _string_metadata_value(metadata.get("title"), metadata.get("name"))
+    source = _string_metadata_value(metadata.get("source"))
+    version = _string_metadata_value(metadata.get("source_version"))
+    heading_path = _string_metadata_value(metadata.get("heading_path"))
+
+    if title:
+        lines.append(f"Title: {title}")
+    if source:
+        lines.append(f"Source: {source}")
+    if version:
+        lines.append(f"Version: {version}")
+    if heading_path:
+        lines.append(f"Heading path: {heading_path}")
+    return "\n".join(lines)
+
+
+def format_chunk_for_retrieval(text: str, metadata: dict[str, Any]) -> str:
+    text = text.strip()
+    header = build_chunk_retrieval_header(metadata)
+    if not header:
+        return text
+    if not text:
+        return header
+    return f"{header}\n\n{text}"
+
+
 def build_chunk_records(
     documents: list[dict[str, Any]],
     chunk_size: int = DEFAULT_CHUNK_SIZE,
     chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
 ) -> list[ChunkRecord]:
     chunk_records: list[ChunkRecord] = []
+    source_versions = load_source_versions()
     for document in documents:
-        chunks = token_window_chunks(
+        chunks = heading_aware_markdown_chunks(
             document["content"],
+            title=str(document.get("name") or ""),
             chunk_size=chunk_size,
             chunk_overlap=chunk_overlap,
         )
-        for index, chunk_text in enumerate(chunks):
+        source = str(document["source"])
+        source_version = source_version_for(source, source_versions)
+        for index, chunk in enumerate(chunks):
+            heading_path = _heading_path_text(chunk.heading_path)
             metadata = {
                 "doc_id": document["doc_id"],
                 "title": document["name"],
                 "url": document["url"],
-                "source": document["source"],
+                "source": source,
+                "source_version": source_version,
                 "retrieve_doc": document["retrieve_doc"],
                 "tokens": document["tokens"],
+                "chunk_tokens": chunk.tokens,
                 "chunk_index": index,
+                "heading_path": heading_path,
             }
             chunk_records.append(
                 ChunkRecord(
                     chunk_id=f'{document["doc_id"]}:{index}',
                     doc_id=document["doc_id"],
-                    text=chunk_text,
+                    text=chunk.text,
                     metadata=metadata,
                 )
             )
@@ -255,6 +724,121 @@ def _title_from_url(url: str) -> str:
     return unquote(slug).replace("-", " ").replace("_", " ").strip()
 
 
+def _env_int(name: str, default: int) -> int:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    try:
+        return int(value)
+    except ValueError:
+        return default
+
+
+def _env_float(name: str, default: float) -> float:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    try:
+        return float(value)
+    except ValueError:
+        return default
+
+
+def _cohere_rate_limited_units(
+    limit: int | None,
+    env_name: str,
+    default: int,
+    margin: float,
+) -> int | None:
+    configured_limit = _env_int(env_name, default) if limit is None else limit
+    if configured_limit <= 0:
+        return None
+    return max(1, int(configured_limit * margin))
+
+
+def _get_cohere_limiter(
+    name: str,
+    units_per_window: int,
+    window_seconds: float,
+) -> SyncWindowLimiter:
+    key = (name, units_per_window, window_seconds)
+    with _cohere_limiter_lock:
+        limiter = _cohere_limiters.get(key)
+        if limiter is None:
+            limiter = SyncWindowLimiter(units_per_window, window_seconds)
+            _cohere_limiters[key] = limiter
+        return limiter
+
+
+def _count_embed_tokens(text: str, encoding: tiktoken.Encoding) -> int:
+    return len(encoding.encode(text, disallowed_special=()))
+
+
+def _iter_cohere_embed_batches(
+    texts: list[str],
+    token_counts: list[int],
+    *,
+    batch_size: int,
+    max_batch_tokens: int | None,
+) -> Iterable[tuple[list[str], int]]:
+    batch: list[str] = []
+    batch_tokens = 0
+
+    for text, token_count in zip(texts, token_counts, strict=True):
+        should_flush = len(batch) >= batch_size
+        if max_batch_tokens is not None and batch:
+            should_flush = should_flush or batch_tokens + token_count > max_batch_tokens
+
+        if should_flush:
+            yield batch, batch_tokens
+            batch = []
+            batch_tokens = 0
+
+        batch.append(text)
+        batch_tokens += max(1, token_count)
+
+    if batch:
+        yield batch, batch_tokens
+
+
+def _is_cohere_rate_limit_error(exc: BaseException) -> bool:
+    status_code = getattr(exc, "status_code", None)
+    if status_code == 429:
+        return True
+    if exc.__class__.__name__ == "TooManyRequestsError":
+        return True
+    return "rate limit" in str(exc).lower() and "429" in str(exc)
+
+
+def _cohere_retry_after_seconds(exc: BaseException) -> float | None:
+    headers = getattr(exc, "headers", None)
+    if not isinstance(headers, dict):
+        return None
+
+    retry_after = headers.get("retry-after") or headers.get("Retry-After")
+    if retry_after is None:
+        return None
+
+    try:
+        return float(retry_after)
+    except ValueError:
+        return None
+
+
+def _wait_for_cohere_retry(
+    exc: BaseException,
+    attempt: int,
+    window_seconds: float,
+) -> None:
+    retry_after = _cohere_retry_after_seconds(exc)
+    if retry_after is not None:
+        delay = retry_after
+    else:
+        delay = min(window_seconds, max(15.0, 2.0 ** attempt))
+
+    time.sleep(delay + random.uniform(0.5, 2.0))
+
+
 def _cohere_embeddings_list(response: Any) -> list[list[float]]:
     embeddings = getattr(response, "embeddings", None)
     if embeddings is None:
@@ -276,24 +860,96 @@ def embed_texts(
     input_type: str,
     model: str = DEFAULT_EMBED_MODEL,
     output_dimension: int = DEFAULT_OUTPUT_DIMENSION,
-    batch_size: int = 96,
+    batch_size: int = DEFAULT_COHERE_EMBED_BATCH_SIZE,
+    max_inputs_per_minute: int | None = None,
+    max_tokens_per_minute: int | None = None,
+    max_requests_per_minute: int | None = None,
     show_progress: bool = False,
     progress_desc: str = "Embedding",
 ) -> list[list[float]]:
+    batch_size = max(1, batch_size)
+    rate_limit_margin = _env_float(
+        "COHERE_EMBED_RATE_LIMIT_MARGIN",
+        DEFAULT_COHERE_EMBED_RATE_LIMIT_MARGIN,
+    )
+    window_seconds = _env_float(
+        "COHERE_EMBED_WINDOW_SECONDS",
+        DEFAULT_COHERE_EMBED_WINDOW_SECONDS,
+    )
+    token_window = _cohere_rate_limited_units(
+        max_tokens_per_minute,
+        "COHERE_EMBED_TPM_LIMIT",
+        DEFAULT_COHERE_EMBED_TPM_LIMIT,
+        rate_limit_margin,
+    )
+    input_window = _cohere_rate_limited_units(
+        max_inputs_per_minute,
+        "COHERE_EMBED_INPUTS_PER_MINUTE",
+        DEFAULT_COHERE_EMBED_INPUTS_PER_MINUTE,
+        rate_limit_margin,
+    )
+    request_window = _cohere_rate_limited_units(
+        max_requests_per_minute,
+        "COHERE_EMBED_RPM_LIMIT",
+        DEFAULT_COHERE_EMBED_RPM_LIMIT,
+        rate_limit_margin,
+    )
+    retry_attempts = max(
+        1,
+        _env_int("COHERE_EMBED_RETRY_ATTEMPTS", DEFAULT_COHERE_EMBED_RETRY_ATTEMPTS),
+    )
+    encoding = tiktoken.get_encoding(DEFAULT_ENCODING)
+    token_counts = [_count_embed_tokens(text, encoding) for text in texts]
+    token_limiter = (
+        _get_cohere_limiter("tokens", token_window, window_seconds)
+        if token_window is not None
+        else None
+    )
+    input_limiter = (
+        _get_cohere_limiter("inputs", input_window, window_seconds)
+        if input_window is not None
+        else None
+    )
+    request_limiter = (
+        _get_cohere_limiter("requests", request_window, window_seconds)
+        if request_window is not None
+        else None
+    )
+
     vectors: list[list[float]] = []
     progress = None
     if show_progress and texts:
         progress = tqdm(total=len(texts), desc=progress_desc, unit="chunk")
 
     try:
-        for batch in batched(texts, batch_size):
-            response = client.embed(
-                model=model,
-                input_type=input_type,
-                embedding_types=["float"],
-                output_dimension=output_dimension,
-                texts=batch,
-            )
+        for batch, batch_tokens in _iter_cohere_embed_batches(
+            texts,
+            token_counts,
+            batch_size=batch_size,
+            max_batch_tokens=token_window,
+        ):
+            for attempt in range(1, retry_attempts + 1):
+                if token_limiter is not None:
+                    token_limiter.acquire(batch_tokens)
+                if input_limiter is not None:
+                    input_limiter.acquire(len(batch))
+                if request_limiter is not None:
+                    request_limiter.acquire(1)
+
+                try:
+                    response = client.embed(
+                        model=model,
+                        input_type=input_type,
+                        embedding_types=["float"],
+                        output_dimension=output_dimension,
+                        texts=batch,
+                    )
+                    break
+                except Exception as exc:
+                    if not _is_cohere_rate_limit_error(exc) or attempt == retry_attempts:
+                        raise
+                    _wait_for_cohere_retry(exc, attempt, window_seconds)
+
             vectors.extend(_cohere_embeddings_list(response))
             if progress is not None:
                 progress.update(len(batch))
@@ -336,6 +992,8 @@ def rerank_results(
                 score=float(item.relevance_score),
                 content=result.content,
                 chunk_content=result.chunk_content,
+                heading_path=result.heading_path,
+                retrieval_method=result.retrieval_method,
             )
         )
     return reranked
@@ -363,6 +1021,113 @@ def _flatten_query_results(values: list[list[Any]] | None) -> list[Any]:
     return values[0]
 
 
+def _add_bm25_token(tokens: list[str], token: str) -> None:
+    token = token.lower().strip("._-/:-")
+    if not token:
+        return
+    if token in BM25_STOP_WORDS:
+        return
+    if len(token) == 1 and token not in {"c", "r"}:
+        return
+    tokens.append(token)
+
+
+def tokenize_for_bm25(text: str) -> list[str]:
+    tokens: list[str] = []
+    for match in BM25_TOKEN_RE.finditer(text):
+        raw_token = match.group(0)
+        _add_bm25_token(tokens, raw_token)
+
+        for part in re.split(r"[._/\-:]+", raw_token):
+            _add_bm25_token(tokens, part)
+            for camel_part in CAMEL_CASE_RE.findall(part):
+                _add_bm25_token(tokens, camel_part)
+
+    return tokens
+
+
+def save_bm25_index(index: BM25Index, output_file: str) -> None:
+    ensure_parent_dir(output_file)
+    with open(output_file, "wb") as handle:
+        pickle.dump(index, handle)
+
+
+def load_bm25_index(path: str) -> BM25Index | None:
+    if not os.path.exists(path):
+        return None
+    with open(path, "rb") as handle:
+        index = pickle.load(handle)
+    if isinstance(index, BM25Index):
+        return index
+    return None
+
+
+def default_bm25_index_path(document_dict_path: str) -> str:
+    path = Path(document_dict_path)
+    name = path.name
+    if name.startswith("document_dict_") and name.endswith(".pkl"):
+        source = name.removeprefix("document_dict_").removesuffix(".pkl")
+        return str(path.with_name(f"bm25_index_{source}.pkl"))
+    return str(path.with_name("bm25_index.pkl"))
+
+
+def result_dedupe_key(result: SearchResult) -> str:
+    if result.heading_path:
+        return f"{result.doc_id}::{result.heading_path}"
+    return result.doc_id or result.chunk_id
+
+
+def reciprocal_rank_fusion(
+    ranked_lists: list[list[SearchResult]],
+    *,
+    rrf_k: int = DEFAULT_RRF_K,
+    top_k: int = DEFAULT_FUSION_TOP_K,
+) -> list[SearchResult]:
+    fused_scores: dict[str, float] = {}
+    representatives: dict[str, SearchResult] = {}
+
+    for ranked_results in ranked_lists:
+        for rank, result in enumerate(ranked_results, start=1):
+            key = result_dedupe_key(result)
+            fused_scores[key] = fused_scores.get(key, 0.0) + 1.0 / (rrf_k + rank)
+
+            current = representatives.get(key)
+            if current is None:
+                representatives[key] = result
+            elif result.retrieval_method == "bm25" and current.retrieval_method != "bm25":
+                representatives[key] = result
+            elif (
+                result.retrieval_method == current.retrieval_method
+                and result.score > current.score
+            ):
+                representatives[key] = result
+
+    fused_results: list[SearchResult] = []
+    for key, score in fused_scores.items():
+        result = representatives[key]
+        fused_results.append(
+            SearchResult(
+                chunk_id=result.chunk_id,
+                doc_id=result.doc_id,
+                title=result.title,
+                url=result.url,
+                source=result.source,
+                retrieve_doc=result.retrieve_doc,
+                tokens=result.tokens,
+                score=score,
+                content=result.content,
+                chunk_content=result.chunk_content,
+                heading_path=result.heading_path,
+                retrieval_method="hybrid"
+                if len(ranked_lists) > 1
+                else result.retrieval_method,
+            )
+        )
+
+    fused_results.sort(key=lambda result: result.score, reverse=True)
+    return fused_results[: max(0, top_k)]
+
+
 class LocalChromaRetriever:
     def __init__(
         self,
@@ -374,7 +1139,11 @@ class LocalChromaRetriever:
         embed_model: str = DEFAULT_EMBED_MODEL,
         rerank_model: str = DEFAULT_RERANK_MODEL,
         dense_top_k: int = DEFAULT_DENSE_TOP_K,
+        bm25_top_k: int = DEFAULT_BM25_TOP_K,
+        fusion_top_k: int = DEFAULT_FUSION_TOP_K,
         rerank_top_k: int = DEFAULT_RERANK_TOP_K,
+        rrf_k: int = DEFAULT_RRF_K,
+        bm25_index_path: str | None = None,
         answer_model_name: str | None = None,
         token_budget: int = DEFAULT_CONTEXT_TOKEN_BUDGET,
     ) -> None:
@@ -382,17 +1151,24 @@ class LocalChromaRetriever:
         self._collection_name = collection_name
         self._document_dict_path = document_dict_path
         self._dense_top_k = dense_top_k
+        self._bm25_top_k = bm25_top_k
+        self._fusion_top_k = fusion_top_k
         self._rerank_top_k = rerank_top_k
+        self._rrf_k = rrf_k
         self._token_budget = token_budget
         self._embed_model = embed_model
         self._rerank_model = rerank_model
         self._encoding = get_token_encoding(answer_model_name)
+        self._bm25_index_path = bm25_index_path or default_bm25_index_path(
+            document_dict_path
+        )
 
         client = chromadb.PersistentClient(path=db_path)
         self._collection = client.get_or_create_collection(name=collection_name)
         with open(document_dict_path, "rb") as handle:
             self._document_dict: dict[str, dict[str, Any]] = pickle.load(handle)
 
+        self._bm25_index = load_bm25_index(self._bm25_index_path)
         self._cohere = cohere.ClientV2(api_key=cohere_api_key)
 
     def search(
@@ -400,6 +1176,31 @@ class LocalChromaRetriever:
         query: str,
         *,
         allowed_sources: list[str] | None = None,
+    ) -> list[SearchResult]:
+        dense_hits = self._dense_search(query, allowed_sources=allowed_sources)
+        bm25_hits = self._bm25_search(query, allowed_sources=allowed_sources)
+        fused_hits = reciprocal_rank_fusion(
+            [hits for hits in (dense_hits, bm25_hits) if hits],
+            rrf_k=self._rrf_k,
+            top_k=self._fusion_top_k,
+        )
+        if not fused_hits:
+            return []
+
+        reranked = rerank_results(
+            self._cohere,
+            query,
+            fused_hits,
+            model=self._rerank_model,
+            top_n=self._rerank_top_k,
+        )
+        return self._apply_token_budget(reranked)
+
+    def _dense_search(
+        self,
+        query: str,
+        *,
+        allowed_sources: list[str] | None = None,
     ) -> list[SearchResult]:
         query_embedding = embed_texts(
             self._cohere,
@@ -422,78 +1223,114 @@ class LocalChromaRetriever:
         distances = _flatten_query_results(raw_results.get("distances"))
 
         dense_hits: list[SearchResult] = []
-        seen_doc_ids: set[str] = set()
         for chunk_id, chunk_text, metadata, distance in zip(
             chunk_ids, documents, metadatas, distances, strict=False
         ):
             if metadata is None:
                 continue
 
-            doc_id = str(metadata["doc_id"])
-            if doc_id in seen_doc_ids:
-                continue
-            seen_doc_ids.add(doc_id)
-
-            full_doc = self._document_dict.get(doc_id)
-            if metadata.get("retrieve_doc") and full_doc is not None:
-                content = get_full_doc_content(full_doc)
-            else:
-                content = chunk_text
-
-            full_doc_name = full_doc.get("name") if isinstance(full_doc, dict) else None
-            url = _string_metadata_value(
-                metadata.get("url"),
-                full_doc.get("url") if isinstance(full_doc, dict) else None,
-            )
-            title = _string_metadata_value(
-                metadata.get("title"),
-                full_doc_name,
-                _title_from_url(url),
-                doc_id,
-            )
-            source = _string_metadata_value(
-                metadata.get("source"),
-                full_doc.get("source") if isinstance(full_doc, dict) else None,
-                default="unknown",
-            )
-            retrieve_doc = bool(
-                metadata.get("retrieve_doc")
-                if "retrieve_doc" in metadata
-                else (
-                    full_doc.get("retrieve_doc") if isinstance(full_doc, dict) else False
-                )
-            )
-            tokens_value = metadata.get("tokens")
-            if _is_missing_metadata_value(tokens_value) and isinstance(full_doc, dict):
-                tokens_value = full_doc.get("tokens")
-            try:
-                tokens = int(tokens_value)
-            except (TypeError, ValueError):
-                tokens = 0
-
             dense_hits.append(
-                SearchResult(
+                self._search_result_from_metadata(
                     chunk_id=str(chunk_id),
-                    doc_id=doc_id,
-                    title=title,
-                    url=url,
-                    source=source,
-                    retrieve_doc=retrieve_doc,
-                    tokens=tokens,
                     score=_distance_to_score(distance),
-                    content=content,
-                    chunk_content=str(chunk_text),
+                    chunk_text=str(chunk_text),
+                    metadata=dict(metadata),
+                    retrieval_method="dense",
                 )
             )
+        return dense_hits
 
-        reranked = rerank_results(
-            self._cohere,
-            query,
-            dense_hits,
-            model=self._rerank_model,
-            top_n=self._rerank_top_k,
+    def _bm25_search(
+        self,
+        query: str,
+        *,
+        allowed_sources: list[str] | None = None,
+    ) -> list[SearchResult]:
+        if self._bm25_index is None:
+            return []
+
+        return [
+            self._search_result_from_metadata(
+                chunk_id=record.chunk_id,
+                score=score,
+                chunk_text=format_chunk_for_retrieval(record.text, record.metadata),
+                metadata=record.metadata,
+                raw_chunk_text=record.text,
+                retrieval_method="bm25",
+            )
+            for record, score in self._bm25_index.search(
+                query,
+                allowed_sources=allowed_sources,
+                top_k=self._bm25_top_k,
+            )
+        ]
+
+    def _search_result_from_metadata(
+        self,
+        *,
+        chunk_id: str,
+        score: float,
+        chunk_text: str,
+        metadata: dict[str, Any],
+        retrieval_method: str,
+        raw_chunk_text: str | None = None,
+    ) -> SearchResult:
+        doc_id = str(metadata["doc_id"])
+        full_doc = self._document_dict.get(doc_id)
+        full_doc_name = full_doc.get("name") if isinstance(full_doc, dict) else None
+        url = _string_metadata_value(
+            metadata.get("url"),
+            full_doc.get("url") if isinstance(full_doc, dict) else None,
+        )
+        title = _string_metadata_value(
+            metadata.get("title"),
+            full_doc_name,
+            _title_from_url(url),
+            doc_id,
+        )
+        source = _string_metadata_value(
+            metadata.get("source"),
+            full_doc.get("source") if isinstance(full_doc, dict) else None,
+            default="unknown",
+        )
+        retrieve_doc = bool(
+            metadata.get("retrieve_doc")
+            if "retrieve_doc" in metadata
+            else (full_doc.get("retrieve_doc") if isinstance(full_doc, dict) else False)
+        )
+        tokens_value = metadata.get("tokens")
+        if _is_missing_metadata_value(tokens_value) and isinstance(full_doc, dict):
+            tokens_value = full_doc.get("tokens")
+        try:
+            tokens = int(tokens_value)
+        except (TypeError, ValueError):
+            tokens = 0
+
+        exact_chunk = _string_metadata_value(
+            raw_chunk_text,
+            metadata.get("raw_text"),
+            default=str(chunk_text),
+        )
+        chunk_for_context = format_chunk_for_retrieval(exact_chunk, metadata)
+        if retrieve_doc and full_doc is not None:
+            content = get_full_doc_content(full_doc)
+        else:
+            content = chunk_for_context
+
+        return SearchResult(
+            chunk_id=chunk_id,
+            doc_id=doc_id,
+            title=title,
+            url=url,
+            source=source,
+            retrieve_doc=retrieve_doc,
+            tokens=tokens,
+            score=score,
+            content=content,
+            chunk_content=exact_chunk,
+            heading_path=_string_metadata_value(metadata.get("heading_path")),
+            retrieval_method=retrieval_method,
         )
-        return self._apply_token_budget(reranked)
 
     def _apply_token_budget(self, results: list[SearchResult]) -> list[SearchResult]:
         filtered: list[SearchResult] = []

scripts/setup.py

@@ -3,6 +3,15 @@ import os
 import logfire
 from dotenv import load_dotenv
 
+from data.scraping_scripts.source_registry import (
+    AVAILABLE_SOURCES,
+    AVAILABLE_SOURCES_UI,
+    COURSE_SOURCE_KEYS,
+    DEFAULT_SELECTED_SOURCE_KEYS,
+    DEFAULT_SELECTED_SOURCES_UI,
+    SOURCE_KEY_TO_LABEL,
+    SOURCE_UI_TO_KEY,
+)
 from .utils import init_mongo_db
 
 load_dotenv(override=True)
@@ -14,6 +23,7 @@ except Exception:
 VECTOR_DB_DIR = "data/chroma-db-all_sources"
 VECTOR_COLLECTION_NAME = "chroma-db-all_sources"
 DOCUMENT_DICT_PATH = f"{VECTOR_DB_DIR}/document_dict_all_sources.pkl"
+BM25_INDEX_PATH = f"{VECTOR_DB_DIR}/bm25_index_all_sources.pkl"
 DEFAULT_MODEL_NAME = "google-genai:gemini-3-flash-preview"
 
 AVAILABLE_MODELS: tuple[dict[str, str], ...] = (
@@ -21,72 +31,6 @@ AVAILABLE_MODELS: tuple[dict[str, str], ...] = (
     {"id": "anthropic:claude-haiku-4-5", "label": "Claude Haiku 4.5"},
 )
 
-AVAILABLE_SOURCES_UI = [
-    "LangChain Docs",
-    "LlamaIndex Docs",
-    "Transformers Docs",
-    "PEFT Docs",
-    "TRL Docs",
-    "OpenAI Cookbooks",
-    "Agentic AI Engineering",
-    "Master AI For Work",
-    "Full Stack AI Engineering",
-    "Beginner Python for AI Engineering",
-]
-
-DEFAULT_SELECTED_SOURCES_UI = [
-    "Agentic AI Engineering",
-    "Master AI For Work",
-    "Full Stack AI Engineering",
-    "Beginner Python for AI Engineering",
-    "Transformers Docs",
-    "PEFT Docs",
-    "TRL Docs",
-    "LlamaIndex Docs",
-    "LangChain Docs",
-    "OpenAI Cookbooks",
-]
-
-AVAILABLE_SOURCES = [
-    "transformers",
-    "peft",
-    "trl",
-    "llama_index",
-    "langchain",
-    "openai_cookbooks",
-    "full_stack_ai_engineering",
-    "beginner_python_for_ai_engineering",
-    "master_ai_for_work",
-    "agentic_ai_engineering",
-]
-
-SOURCE_UI_TO_KEY = {
-    "Transformers Docs": "transformers",
-    "PEFT Docs": "peft",
-    "TRL Docs": "trl",
-    "LlamaIndex Docs": "llama_index",
-    "LangChain Docs": "langchain",
-    "OpenAI Cookbooks": "openai_cookbooks",
-    "Full Stack AI Engineering": "full_stack_ai_engineering",
-    "Beginner Python for AI Engineering": "beginner_python_for_ai_engineering",
-    "Master AI For Work": "master_ai_for_work",
-    "Agentic AI Engineering": "agentic_ai_engineering",
-}
-
-COURSE_SOURCE_KEYS = frozenset(
-    {
-        "full_stack_ai_engineering",
-        "beginner_python_for_ai_engineering",
-        "master_ai_for_work",
-        "agentic_ai_engineering",
-    }
-)
-
-SOURCE_KEY_TO_LABEL = {value: key for key, value in SOURCE_UI_TO_KEY.items()}
-DEFAULT_SELECTED_SOURCE_KEYS = tuple(
-    SOURCE_UI_TO_KEY[label] for label in DEFAULT_SELECTED_SOURCES_UI
-)
-
 CONCURRENCY_COUNT = int(os.getenv("CONCURRENCY_COUNT", 64))
 MONGODB_URI = os.getenv("MONGODB_URI")
 
@@ -122,6 +66,7 @@ __all__ = [
     "DEFAULT_SELECTED_SOURCES_UI",
     "CONCURRENCY_COUNT",
     "DEFAULT_MODEL_NAME",
+    "BM25_INDEX_PATH",
     "DOCUMENT_DICT_PATH",
     "SOURCE_KEY_TO_LABEL",
     "SOURCE_UI_TO_KEY",

tests/test_chroma_rag.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import asyncio
+import pickle
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from data.scraping_scripts.add_context_to_nodes import process
+from data.scraping_scripts.create_vector_stores import write_retrieval_artifacts
+from llama_index.core import Document
+from scripts.chroma_rag import (
+    BM25Index,
+    ChunkRecord,
+    build_chunk_records,
+    heading_aware_markdown_chunks,
+    reciprocal_rank_fusion,
+    load_bm25_index,
+    SearchResult,
+)
+
+
+class ChromaRagTestCase(unittest.TestCase):
+    def test_heading_aware_chunks_keep_code_blocks_intact(self) -> None:
+        code_lines = "\n".join(f"print({index})" for index in range(120))
+        markdown = f"""# Guide
+
+## Install
+
+Use `pip install`.
+
+## Example
+
+```python
+{code_lines}
+```
+
+After the example.
+"""
+
+        chunks = heading_aware_markdown_chunks(
+            markdown,
+            title="Guide",
+            chunk_size=80,
+        )
+
+        code_chunks = [chunk for chunk in chunks if "print(0)" in chunk.text]
+        self.assertEqual(len(code_chunks), 1)
+        self.assertIn("print(119)", code_chunks[0].text)
+        self.assertIn("Example", code_chunks[0].heading_path)
+
+    def test_build_chunk_records_adds_heading_metadata(self) -> None:
+        records = build_chunk_records(
+            [
+                {
+                    "doc_id": "doc-1",
+                    "name": "Guide",
+                    "url": "https://example.com/guide",
+                    "source": "transformers",
+                    "retrieve_doc": False,
+                    "tokens": 1000,
+                    "content": "# Guide\n\n## Install\n\nUse `AutoModel`.",
+                }
+            ]
+        )
+
+        self.assertEqual(records[0].metadata["heading_path"], "Guide")
+        self.assertEqual(records[1].metadata["heading_path"], "Guide > Install")
+        self.assertIn("source_version", records[0].metadata)
+
+    def test_bm25_search_finds_keywords_and_filters_sources(self) -> None:
+        records = [
+            ChunkRecord(
+                chunk_id="a",
+                doc_id="doc-a",
+                text="Use AutoModel.from_pretrained for model loading.",
+                metadata={"doc_id": "doc-a", "source": "transformers"},
+            ),
+            ChunkRecord(
+                chunk_id="b",
+                doc_id="doc-b",
+                text="Create a prompt template for chains.",
+                metadata={"doc_id": "doc-b", "source": "langchain"},
+            ),
+        ]
+        index = BM25Index.build(records)
+
+        hits = index.search("AutoModel.from_pretrained", allowed_sources=["transformers"])
+
+        self.assertEqual([record.chunk_id for record, _score in hits], ["a"])
+        self.assertEqual(index.search("AutoModel", allowed_sources=["langchain"]), [])
+
+    def test_retrieval_artifact_writer_persists_bm25_and_document_dict(self) -> None:
+        document_rows = [
+            {
+                "doc_id": "doc-1",
+                "name": "Transformers Loading",
+                "url": "https://example.com/loading",
+                "source": "transformers",
+                "retrieve_doc": False,
+                "tokens": 1200,
+                "content": "# Loading\n\n## AutoModel\n\nUse `AutoModel.from_pretrained`.",
+            }
+        ]
+
+        with tempfile.TemporaryDirectory(dir="/private/tmp") as temp_dir:
+            db_path = Path(temp_dir)
+            count = write_retrieval_artifacts(
+                config={
+                    "document_dict_file": "document_dict_test.pkl",
+                    "bm25_index_file": "bm25_index_test.pkl",
+                },
+                document_rows=document_rows,
+                db_path=str(db_path),
+            )
+
+            document_dict_path = db_path / "document_dict_test.pkl"
+            bm25_path = db_path / "bm25_index_test.pkl"
+
+            self.assertGreaterEqual(count, 1)
+            self.assertTrue(document_dict_path.exists())
+            self.assertTrue(bm25_path.exists())
+
+            with open(document_dict_path, "rb") as handle:
+                document_dict = pickle.load(handle)
+            self.assertEqual(document_dict["doc-1"]["name"], "Transformers Loading")
+
+            index = load_bm25_index(str(bm25_path))
+            self.assertIsNotNone(index)
+            assert index is not None
+            hits = index.search("AutoModel.from_pretrained")
+            self.assertEqual(hits[0][0].doc_id, "doc-1")
+            self.assertTrue(
+                any(record.metadata["heading_path"] for record in index.records)
+            )
+
+    def test_context_processing_uses_heading_chunks_and_raw_text_metadata(self) -> None:
+        async def fake_situate_context(_doc: str, chunk: str) -> str:
+            return f"Situated {chunk.splitlines()[0]}"
+
+        document = Document(
+            doc_id="doc-1",
+            text="# Guide\n\n## Setup\n\nUse `AutoModel.from_pretrained`.",
+            metadata={
+                "title": "Guide",
+                "url": "https://example.com/guide",
+                "tokens": 1000,
+                "retrieve_doc": False,
+                "source": "transformers",
+            },
+        )
+
+        with patch(
+            "data.scraping_scripts.add_context_to_nodes.situate_context",
+            fake_situate_context,
+        ):
+            records = asyncio.run(process([document], semaphore_limit=1))
+
+        self.assertGreaterEqual(len(records), 1)
+        setup_record = next(
+            record for record in records if record.metadata["heading_path"] == "Guide > Setup"
+        )
+        self.assertIn("raw_text", setup_record.metadata)
+        self.assertIn("Title: Guide", setup_record.text)
+        self.assertIn("Heading path: Guide > Setup", setup_record.text)
+        self.assertIn("Context: Situated", setup_record.text)
+
+    def test_rrf_prefers_overlap_across_ranked_lists(self) -> None:
+        dense_only = self._result("dense-only", 0.9, "dense")
+        overlap_dense = self._result("overlap", 0.7, "dense")
+        overlap_bm25 = self._result("overlap", 4.0, "bm25")
+        bm25_only = self._result("bm25-only", 5.0, "bm25")
+
+        fused = reciprocal_rank_fusion(
+            [[dense_only, overlap_dense], [bm25_only, overlap_bm25]],
+            top_k=4,
+        )
+
+        self.assertEqual(fused[0].chunk_id, "overlap")
+        self.assertEqual(fused[0].retrieval_method, "hybrid")
+
+    def _result(self, chunk_id: str, score: float, method: str) -> SearchResult:
+        return SearchResult(
+            chunk_id=chunk_id,
+            doc_id=chunk_id,
+            title=chunk_id,
+            url="",
+            source="test",
+            retrieve_doc=False,
+            tokens=10,
+            score=score,
+            content=chunk_id,
+            chunk_content=chunk_id,
+            heading_path="section",
+            retrieval_method=method,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_process_md_files.py

@@ -0,0 +1,58 @@
+from data.scraping_scripts.process_md_files import extract_title
+
+
+def test_extract_title_prefers_frontmatter_title() -> None:
+    content = """---
+title: "Dappier integration"
+description: "Integrate with the Dappier retriever using LangChain Python."
+---
+
+# DappierRetriever
+"""
+
+    assert extract_title(content) == "Dappier integration"
+
+
+def test_extract_title_skips_frontmatter_delimiters() -> None:
+    content = """---
+description: A page without an explicit title
+---
+
+# Real page title
+"""
+
+    assert extract_title(content) == "Real page title"
+
+
+def test_extract_title_uses_sidebar_title_when_title_missing() -> None:
+    content = """---
+sidebarTitle: Overview
+description: A page without an explicit title
+---
+
+# Longer body heading
+"""
+
+    assert extract_title(content) == "Overview"
+
+
+def test_extract_title_ignores_headings_inside_code_fences() -> None:
+    content = """```python
+# Not a page title
+```
+
+# Real page title
+"""
+
+    assert extract_title(content) == "Real page title"
+
+
+def test_extract_title_keeps_first_line_fallback() -> None:
+    content = """---
+description: A page without heading syntax
+---
+
+First paragraph fallback
+"""
+
+    assert extract_title(content) == "First paragraph fallback"

tests/test_retire_source_workflow.py

@@ -0,0 +1,54 @@
+import ast
+
+from data.scraping_scripts.retire_source_workflow import (
+    remove_source_from_registry_text,
+)
+
+
+def test_remove_source_from_registry_text_removes_active_source_entries():
+    registry = '''SOURCE_CONFIGS = {
+    "openai_cookbooks": {
+        "output_file": "data/openai_cookbooks_data.jsonl",
+        "nested": {"keep": "balanced"},
+    },
+    "langchain": {
+        "output_file": "data/langchain_data.jsonl",
+    },
+}
+
+DOC_SOURCE_KEYS = (
+    "openai_cookbooks",
+    "langchain",
+)
+GITHUB_SOURCE_KEYS = (
+    "openai_cookbooks",
+    "langchain",
+)
+SOURCE_KEY_TO_LABEL = {
+    "openai_cookbooks": "OpenAI Cookbooks",
+    "langchain": "LangChain Docs",
+}
+UI_SOURCE_KEYS = (
+    "openai_cookbooks",
+    "langchain",
+)
+'''
+
+    updated, changed = remove_source_from_registry_text(
+        registry,
+        "openai_cookbooks",
+    )
+
+    assert changed is True
+    assert "openai_cookbooks" not in updated
+    assert '"langchain"' in updated
+    ast.parse(updated)
+
+
+def test_remove_source_from_registry_text_ignores_missing_source():
+    registry = 'SOURCE_CONFIGS = {"langchain": {"output_file": "x"}}\n'
+
+    updated, changed = remove_source_from_registry_text(registry, "missing")
+
+    assert changed is False
+    assert updated == registry