Spaces:
Paused
Paused
| """Build the vector store from rag/sources.yaml. | |
| python -m rag.ingest # all sources | |
| python -m rag.ingest --only mujoco # one collection | |
| python -m rag.ingest --reset # drop and rebuild collections | |
| Clones each source repository (shallow, sparse) into <data>/sources/, chunks | |
| the included files, and upserts them into Chroma at <data>/chroma/. Also | |
| writes <data>/ros2_interfaces.json, the catalogue of message/service types | |
| per distro that /code/generate validates against. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import fnmatch | |
| import json | |
| import subprocess | |
| import sys | |
| from pathlib import Path | |
| import yaml | |
| from api.config import get_settings | |
| from .chunking import Chunk, chunk_file, interface_name, parse_mjcf_schema | |
| from .store import COLLECTIONS, RagStore | |
| SOURCES_FILE = Path(__file__).with_name("sources.yaml") | |
| def load_sources(path: Path = SOURCES_FILE) -> list[dict]: | |
| return yaml.safe_load(path.read_text())["sources"] | |
| def _git(*args: str, cwd: Path | None = None) -> str: | |
| return subprocess.run(["git", *args], cwd=cwd, check=True, capture_output=True, text=True).stdout | |
| def checkout(source: dict, sources_dir: Path) -> Path: | |
| """Shallow sparse clone (or update) of a source repo; returns its directory.""" | |
| dest = sources_dir / source["name"] | |
| ref = source.get("ref", "main") | |
| if not (dest / ".git").exists(): | |
| dest.parent.mkdir(parents=True, exist_ok=True) | |
| _git( | |
| "clone", | |
| "--quiet", | |
| "--depth", | |
| "1", | |
| "--filter=blob:none", | |
| "--no-checkout", | |
| "-b", | |
| ref, | |
| source["repo"], | |
| str(dest), | |
| ) | |
| else: | |
| _git("fetch", "--quiet", "--depth", "1", "origin", ref, cwd=dest) | |
| _git("reset", "--quiet", "--hard", "FETCH_HEAD", cwd=dest) | |
| patterns = list(source.get("checkout", [])) + list(source["include"]) | |
| _git("sparse-checkout", "set", "--no-cone", *patterns, cwd=dest) | |
| _git("checkout", "--quiet", cwd=dest) | |
| return dest | |
| def included_files(root: Path, patterns: list[str]) -> list[str]: | |
| files = [] | |
| for p in sorted(root.rglob("*")): | |
| if p.is_file() and ".git" not in p.parts: | |
| rel = p.relative_to(root).as_posix() | |
| if any(fnmatch.fnmatch(rel, pat) for pat in patterns): | |
| files.append(rel) | |
| return files | |
| def ingest(only: str | None = None, reset: bool = False, log=print) -> dict: | |
| settings = get_settings() | |
| sources_dir = settings.data_dir / "sources" | |
| store = RagStore(settings.data_dir / "chroma") | |
| sources = [s for s in load_sources() if only in (None, s["collection"])] | |
| if reset: | |
| for name in {s["collection"] for s in sources}: | |
| store.reset(name) | |
| interfaces: dict[str, dict[str, str]] = {} | |
| catalog_path = settings.data_dir / "ros2_interfaces.json" | |
| if catalog_path.exists() and not reset: | |
| interfaces = json.loads(catalog_path.read_text()) | |
| counts: dict[str, int] = {} | |
| for source in sources: | |
| log(f"[{source['name']}] fetching {source['repo']}@{source.get('ref', 'main')}") | |
| root = checkout(source, sources_dir) | |
| files = included_files(root, source["include"]) | |
| if not files: | |
| log(f"[{source['name']}] warning: no files matched {source['include']}") | |
| chunks: list[Chunk] = [] | |
| meta = dict(source.get("meta", {})) | |
| for rel in files: | |
| text = (root / rel).read_text(errors="replace") | |
| chunks += chunk_file(rel, text, source["name"], meta, source.get("package")) | |
| if rel.endswith("XMLschema.rst"): | |
| schema_path = settings.data_dir / "mjcf_schema.json" | |
| schema_path.write_text(json.dumps(parse_mjcf_schema(text), indent=1)) | |
| log(f"[{source['name']}] wrote {schema_path}") | |
| name = interface_name(rel, source.get("package")) | |
| if name and meta.get("distro"): | |
| interfaces.setdefault(meta["distro"], {})[name] = text | |
| n = store.add(source["collection"], chunks) | |
| counts[source["name"]] = n | |
| log(f"[{source['name']}] {len(files)} files -> {n} chunks") | |
| if interfaces: | |
| catalog_path.parent.mkdir(parents=True, exist_ok=True) | |
| catalog_path.write_text(json.dumps(interfaces, indent=1, sort_keys=True)) | |
| log(f"interface catalogue: {', '.join(f'{d}={len(v)}' for d, v in sorted(interfaces.items()))}") | |
| for name in COLLECTIONS: | |
| log(f"collection {name}: {store.count(name)} chunks") | |
| return counts | |
| def main(argv: list[str] | None = None) -> int: | |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| parser.add_argument("--only", choices=COLLECTIONS, help="ingest one collection") | |
| parser.add_argument("--reset", action="store_true", help="drop the collections first") | |
| args = parser.parse_args(argv) | |
| try: | |
| ingest(args.only, args.reset) | |
| except subprocess.CalledProcessError as e: | |
| print(f"error: {' '.join(e.cmd)} failed:\n{e.stderr}", file=sys.stderr) | |
| return 1 | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |