Pybricks RAG Assistant

#!/usr/bin/env python3 """ Build a beginner-friendly Pybricks RAG corpus into ChromaDB. - Crawls/loads selected Pybricks docs/tutorials + key GitHub changelogs. - Normalizes & chunks (HTML/Markdown-aware; preserves code blocks). - Creates multi-view docs (full text + headings/signatures + identifiers). - Indexes into ChromaDB collections: pybricks_docs and pybricks_snippets. Requires: requests, beautifulsoup4, lxml, markdownify, tiktoken, chromadb, sentence-transformers pip install requests beautifulsoup4 lxml markdownify tiktoken chromadb sentence-transformers """ import re import time import json import hashlib from urllib.parse import urljoin, urlparse import requests from bs4 import BeautifulSoup from markdownify import markdownify as md import tiktoken import chromadb from chromadb.utils import embedding_functions ############################ # Config ############################ SEED_URLS = [ # API docs (latest & versioned) "https://docs.pybricks.com/en/latest/robotics.html", # DriveBase, etc. "https://docs.pybricks.com/en/latest/parameters/", # Port, Direction, Stop, etc. "https://docs.pybricks.com/", # root (to discover other key pages) "https://docs.pybricks.com/en/v3.5.0/robotics.html", # older ref (version contrast) "https://docs.pybricks.com/en/v3.3.0/robotics.html", # Tutorials / Getting started "https://pybricks.com/learn/getting-started/pybricks-environment/", "https://pybricks.com/learn/getting-started/install-pybricks/", "https://pybricks.com/projects/tutorials/wireless/hub-to-device/pc-communication/", "https://pybricks.com/projects/tutorials/wireless/hub-to-hub/broadcast/", "https://pybricks.com/projects/tutorials/wireless/remote-control/button-basics/", "https://pybricks.com/projects/sets/mindstorms-robot-inventor/other-models/quick-reference/", "https://code.pybricks.com/", # Changelogs (version awareness) "https://github.com/pybricks/pybricks-micropython/blob/master/CHANGELOG.md", "https://github.com/pybricks/pybricks-code/blob/master/CHANGELOG.md", # pybricksdev (optional power tool docs) "https://docs.pybricks.com/projects/pybricksdev/en/latest/api/", ] ALLOWED_DOMAINS = { "docs.pybricks.com", "pybricks.com", "code.pybricks.com", "github.com", } MAX_PAGES = 80 # keep crawl polite; bump when needed CRAWL_TIMEOUT = 15 USER_AGENT = "pybricks-rag-ingestor/1.0 (+for educational/FLL use)" # Chunking defaults DOC_CHUNK_TOKENS = 600 DOC_OVERLAP_TOKENS = 100 # Chroma config CHROMA_DIR = "./chroma_pybricks" COLL_DOCS = "pybricks_docs" COLL_SNIPPETS = "pybricks_snippets" # short task-oriented code templates/snips # Embeddings: sentence-transformers (local) EMBED_MODEL = "all-MiniLM-L6-v2" ############################ # Helpers ############################ tok = tiktoken.get_encoding("cl100k_base") def count_tokens(text: str) -> int: return len(tok.encode(text)) def hash_id(s: str) -> str: return hashlib.sha256(s.encode("utf-8")).hexdigest()[:32] def within_domain(url: str) -> bool: try: host = urlparse(url).netloc return any(host.endswith(d) for d in ALLOWED_DOMAINS) except Exception: return False def get(url: str) -> requests.Response: return requests.get(url, headers={"User-Agent": USER_AGENT}, timeout=CRAWL_TIMEOUT) def extract_links(url: str, soup: BeautifulSoup): for a in soup.select("a[href]"): href = a["href"] if href.startswith("#"): continue u = urljoin(url, href) if within_domain(u): yield u.split("#")[0] def html_to_markdown_keep_code(html: str) -> str: # Convert HTML to Markdown but keep code blocks and headings clean # md() handles most; we'll fix common issues text = md(html, heading_style="ATX", strip=["nav","footer","script","style","noscript"]) # compact extra blank lines text = re.sub(r"\n{3,}", "\n\n", text).strip() return text def normalize_markdown(text: str) -> str: # de-dup whitespace, make headings spaced text = re.sub(r"[ \t]+\n", "\n", text) return text.strip() ############################ # Chunkers ############################ HEADING_PATTERN = re.compile(r"^(#{1,3})\s+(.*)", re.MULTILINE) def split_by_headings(markdown: str): """ Split markdown by H1–H3. Keep fenced code blocks with their section. """ lines = markdown.splitlines() sections = [] current = {"title": None, "content": []} def push(): if current["content"]: sections.append({ "title": current["title"] or "", "text": "\n".join(current["content"]).strip() }) for i, line in enumerate(lines): if HEADING_PATTERN.match(line): # new section if current["content"]: push() current = {"title": None, "content": []} current["title"] = HEADING_PATTERN.match(line).group(2) current["content"].append(line) else: current["content"].append(line) push() return [s for s in sections if s["text"]] def window_tokens(text: str, size_tokens: int, overlap_tokens: int): ids = tok.encode(text) n = len(ids) if n <= size_tokens: yield text return start = 0 while start < n: end = min(n, start + size_tokens) chunk_ids = ids[start:end] yield tok.decode(chunk_ids) if end == n: break start = max(0, end - overlap_tokens) def identifiers_view(text: str) -> str: # Pull likely “identifier-ish” tokens: backticked names, UPPER_CASE, CamelCase, Port.X, Direction.X, drivebase methods, HTML ids ids = re.findall(r"`([^`]+)`|([A-Z_]{2,})|([A-Za-z_][A-Za-z0-9_]+)|id=\"([^\"]+)\"", text) flat = [] for g in ids: flat.extend([x for x in g if x]) # de-noise trivial words flat = [w for w in flat if len(w) > 2] return "\n".join(sorted(set(flat))) ############################ # Crawl ############################ def crawl(): seen = set() queue = list(SEED_URLS) pages = [] while queue and len(pages) < MAX_PAGES: url = queue.pop(0) if url in seen or not within_domain(url): continue seen.add(url) try: r = get(url) if r.status_code != 200: continue ct = r.headers.get("Content-Type","") html = r.text soup = BeautifulSoup(html, "lxml") # strip obvious boilerplate for tag in soup(["nav","footer","script","style","noscript"]): tag.decompose() body = soup.body or soup main = body.select_one("main") or body # Many docs are Sphinx → good structure + headings # Convert to markdown for easier heading split markdown = html_to_markdown_keep_code(str(main)) markdown = normalize_markdown(markdown) pages.append({ "url": url, "markdown": markdown }) # discover more links from key roots (docs/tutorials) if any(url.startswith(prefix) for prefix in [ "https://docs.pybricks.com/", "https://pybricks.com/learn/", "https://pybricks.com/projects/", ]): for u in set(extract_links(url, soup)): if u not in seen and within_domain(u): # keep within same area if len(pages) + len(queue) < MAX_PAGES: queue.append(u) time.sleep(0.3) # be polite except Exception: continue return pages ############################ # Build records & index ############################ def build_records(pages): records = [] for p in pages: url = p["url"] mdtext = p["markdown"] # version tags from URL version = "latest" if "/en/latest/" in url else ( re.search(r"/en/v?(\d+\.\d+\.\d+)/", url).group(1) if re.search(r"/en/v?(\d+\.\d+\.\d+)/", url) else "stable" ) # sectionize by H1–H3 sections = split_by_headings(mdtext) or [{"title":"","text":mdtext}] for sec in sections: # windowing to stay within token budget for win in window_tokens(sec["text"], DOC_CHUNK_TOKENS, DOC_OVERLAP_TOKENS): rec_id = hash_id(url + "|" + sec["title"] + "|" + win[:200]) records.append({ "id": rec_id, "text": win, "meta": { "url": url, "title": sec["title"], "version": version, "source": urlparse(url).netloc, "breadcrumbs": sec["title"], "type": "doc" }, "identifiers": identifiers_view(win) }) return records ############################ # ChromaDB ############################ def upsert_chroma(records): client = chromadb.PersistentClient(path=CHROMA_DIR) ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBED_MODEL) docs = client.get_or_create_collection(name=COLL_DOCS, embedding_function=ef, metadata={"hnsw:space":"cosine"}) ids = [r["id"] for r in records] texts = [r["text"] for r in records] metas = [r["meta"] for r in records] # store a second view as metadata field for regex/BM25 on your side, but we can also make a tiny aux collection: docs.add(ids=ids, documents=texts, metadatas=metas) # Optional: create a lightweight identifiers collection to help rerank exact symbol hits ident = client.get_or_create_collection(name=f"{COLL_DOCS}_identifiers", embedding_function=ef, metadata={"hnsw:space":"cosine"}) ident.add( ids=[f"I_{i}" for i in ids], documents=[r["identifiers"] for r in records], metadatas=[{**r["meta"], "type":"identifiers"} for r in records] ) # Starter snippet/templates (task-oriented) snippets = [ { "id": "snip_drivebase_basic", "text": """# Drive straight and turn from pybricks.hubs import PrimeHub from pybricks.pupdevices import Motor from pybricks.parameters import Port from pybricks.robotics import DriveBase hub = PrimeHub() left = Motor(Port.A) right = Motor(Port.B) bot = DriveBase(left, right, wheel_diameter=56, axle_track=114) bot.straight(300) # mm bot.turn(90) # + is clockwise/right """, "meta": {"topic":"drivebase","version":"stable","url":"https://docs.pybricks.com/en/latest/robotics.html"} }, ] sn = client.get_or_create_collection(name=COLL_SNIPPETS, embedding_function=ef, metadata={"hnsw:space":"cosine"}) sn.add( ids=[s["id"] for s in snippets], documents=[s["text"] for s in snippets], metadatas=[s["meta"] for s in snippets] ) return client def main(): print("Crawling…") pages = crawl() print(f"Crawled {len(pages)} pages") records = build_records(pages) print(f"Built {len(records)} records") client = upsert_chroma(records) print("ChromaDB ready at", CHROMA_DIR) if __name__ == "__main__": main()

#!/usr/bin/env python3 """ Minimal MCP-style tool server exposing Pybricks RAG functions over ChromaDB. Run: uvicorn mcp_server:app --reload --port 8765 Tools: - /tools/docs.search - /tools/code.search - /tools/snippet.suggest - /tools/compare.versions """ import re from typing import List, Optional, Dict, Any from fastapi import FastAPI, Body, Query from pydantic import BaseModel import chromadb from chromadb.utils import embedding_functions CHROMA_DIR = "./chroma_pybricks" COLL_DOCS = "pybricks_docs" COLL_DOCS_IDENT = f"{COLL_DOCS}_identifiers" COLL_SNIPPETS = "pybricks_snippets" EMBED_MODEL = "all-MiniLM-L6-v2" client = chromadb.PersistentClient(path=CHROMA_DIR) ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBED_MODEL) docs = client.get_or_create_collection(name=COLL_DOCS, embedding_function=ef) ident = client.get_or_create_collection(name=COLL_DOCS_IDENT, embedding_function=ef) snips = client.get_or_create_collection(name=COLL_SNIPPETS, embedding_function=ef) app = FastAPI(title="Pybricks MCP Tools") class SearchRequest(BaseModel): query: str k: int = 8 filters: Optional[Dict[str, Any]] = None # e.g., {"version":"latest"} or {"source":"docs.pybricks.com"} class SearchResponse(BaseModel): results: List[Dict[str, Any]] def apply_filters(metadatas: List[Dict[str, Any]], filters: Dict[str, Any]) -> List[int]: idxs = [] for i, m in enumerate(metadatas): ok = True for k, v in filters.items(): if str(m.get(k)) != str(v): ok = False break if ok: idxs.append(i) return idxs @app.post("/tools/docs.search", response_model=SearchResponse) def docs_search(req: SearchRequest): # dense search res = docs.query(query_texts=[req.query], n_results=max(req.k, 8), include=["documents","metadatas","distances"]) docs_idxs = list(range(len(res["documents"][0]))) # optional hard filter by simple metadata if req.filters: keep = apply_filters(res["metadatas"][0], req.filters) docs_idxs = [i for i in docs_idxs if i in keep] results = [] for i in docs_idxs[:req.k]: results.append({ "text": res["documents"][0][i], "score": 1 - res["distances"][0][i], "meta": res["metadatas"][0][i], }) return {"results": results} class CodeSearchRequest(BaseModel): symbol: str k: int = 8 version: Optional[str] = None @app.post("/tools/code.search", response_model=SearchResponse) def code_search(req: CodeSearchRequest): # prioritize identifier view collection; fall back to docs res = ident.query(query_texts=[req.symbol], n_results=max(req.k, 8), include=["documents","metadatas","distances"]) out = [] for i in range(len(res["documents"][0])): m = res["metadatas"][0][i] if req.version and m.get("version") != req.version: continue out.append({ "text": res["documents"][0][i], "score": 1 - res["distances"][0][i], "meta": m, }) if not out: return docs_search(SearchRequest(query=req.symbol, k=req.k, filters={"version": req.version} if req.version else None)) return {"results": out[:req.k]} class SnippetSuggestRequest(BaseModel): goal: str hardware: Optional[str] = None k: int = 3 @app.post("/tools/snippet.suggest", response_model=SearchResponse) def snippet_suggest(req: SnippetSuggestRequest): query = req.goal + (" " + req.hardware if req.hardware else "") res = snips.query(query_texts=[query], n_results=req.k, include=["documents","metadatas","distances"]) results = [{ "text": res["documents"][0][i], "score": 1 - res["distances"][0][i], "meta": res["metadatas"][0][i], } for i in range(len(res["documents"][0]))] return {"results": results} class CompareVersionsRequest(BaseModel): feature: str k: int = 6 @app.post("/tools/compare.versions", response_model=SearchResponse) def compare_versions(req: CompareVersionsRequest): # Look for the feature in CHANGELOG pages inside docs collection res = docs.query(query_texts=[f"CHANGELOG {req.feature}"], n_results=20, include=["documents","metadatas","distances"]) hits = [] for i in range(len(res["documents"][0])): m = res["metadatas"][0][i] if "github.com/pybricks" in m.get("url",""): hits.append({ "text": res["documents"][0][i], "score": 1 - res["distances"][0][i], "meta": m }) return {"results": hits[:req.k]}

python ingest_pybricks.py

pip install fastapi uvicorn uvicorn mcp_server:app --reload --port 8765

from pybricks.hubs import PrimeHub from pybricks.pupdevices import Motor from pybricks.parameters import Port from pybricks.robotics import DriveBase hub = PrimeHub() left = Motor(Port.A) right = Motor(Port.B) bot = DriveBase(left, right, wheel_diameter=56, axle_track=114) bot.straight(300) # millimeters bot.turn(90) # + is clockwise (right)