Spaces:
Paused
Paused
| """Split source files into retrieval chunks. | |
| Chunks stay close to the natural unit of each format: a documentation section, | |
| one ROS interface definition, or a window of an MJCF file. Every chunk carries | |
| enough header text ("MuJoCo XML reference > geom") that it reads on its own | |
| when pasted into a prompt. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import re | |
| from dataclasses import dataclass, field | |
| from pathlib import PurePosixPath | |
| MAX_CHARS = 1800 | |
| OVERLAP = 200 | |
| class Chunk: | |
| text: str | |
| source: str # e.g. "mujoco_docs:doc/XMLreference.rst#geom" | |
| meta: dict = field(default_factory=dict) | |
| def id(self) -> str: | |
| return hashlib.sha1(f"{self.source}\n{self.text}".encode()).hexdigest() | |
| def _windows(text: str, max_chars: int = MAX_CHARS, overlap: int = OVERLAP) -> list[str]: | |
| """Split on paragraph boundaries into pieces of at most ~max_chars.""" | |
| text = text.strip() | |
| if len(text) <= max_chars: | |
| return [text] if text else [] | |
| paragraphs = re.split(r"\n\s*\n", text) | |
| out: list[str] = [] | |
| buf = "" | |
| for para in paragraphs: | |
| while len(para) > max_chars: # a single huge paragraph (e.g. a table) | |
| if buf: | |
| out.append(buf) | |
| buf = "" | |
| out.append(para[:max_chars]) | |
| para = para[max_chars - overlap :] | |
| if buf and len(buf) + len(para) + 2 > max_chars: | |
| out.append(buf) | |
| buf = buf[-overlap:] + "\n\n" + para if overlap else para | |
| else: | |
| buf = f"{buf}\n\n{para}" if buf else para | |
| if buf.strip(): | |
| out.append(buf) | |
| return out | |
| # RST section titles are underlined (and optionally overlined) with punctuation. | |
| _RST_TITLE = re.compile(r"^(?P<title>[^\n]+)\n(?P<rule>[=\-~^\"'`#*+:.]{3,})[ \t]*$", re.MULTILINE) | |
| _MD_TITLE = re.compile(r"^(#{1,4})\s+(.+)$", re.MULTILINE) | |
| def _sections(text: str, markdown: bool) -> list[tuple[str, str]]: | |
| """(title, body) pairs in document order; text before the first title gets title ''.""" | |
| pattern = _MD_TITLE if markdown else _RST_TITLE | |
| matches = list(pattern.finditer(text)) | |
| if not matches: | |
| return [("", text)] | |
| sections = [] | |
| if matches[0].start() > 0: | |
| sections.append(("", text[: matches[0].start()])) | |
| for i, m in enumerate(matches): | |
| title = (m.group(2) if markdown else m.group("title")).strip() | |
| end = matches[i + 1].start() if i + 1 < len(matches) else len(text) | |
| sections.append((title, text[m.end() : end])) | |
| return sections | |
| def chunk_document(text: str, source: str, meta: dict, doc_title: str) -> list[Chunk]: | |
| markdown = source.endswith(".md") | |
| chunks = [] | |
| for title, body in _sections(text, markdown): | |
| header = f"{doc_title} > {title}" if title else doc_title | |
| for piece in _windows(body): | |
| if len(piece) < 40: | |
| continue | |
| anchor = re.sub(r"[^a-z0-9]+", "-", title.lower()).strip("-") | |
| chunks.append(Chunk(f"[{header}]\n{piece}", f"{source}#{anchor}" if anchor else source, dict(meta))) | |
| return chunks | |
| def chunk_code(text: str, source: str, meta: dict, label: str) -> list[Chunk]: | |
| return [Chunk(f"[{label}]\n{piece}", source, dict(meta)) for piece in _windows(text, overlap=0)] | |
| def interface_name(path: str, package: str | None = None) -> str | None: | |
| """'std_msgs/msg/String.msg' -> 'std_msgs/msg/String' (package can be given for flat repos).""" | |
| p = PurePosixPath(path) | |
| if p.suffix not in (".msg", ".srv", ".action") or p.parent.name not in ("msg", "srv", "action"): | |
| return None | |
| pkg = package or (p.parent.parent.name if p.parent.parent.name else None) | |
| if not pkg: | |
| return None | |
| return f"{pkg}/{p.parent.name}/{p.stem}" | |
| def interface_kind(name: str) -> str: | |
| return "Service" if "/srv/" in name else "Action" if "/action/" in name else "Message" | |
| def chunk_interface(text: str, source: str, meta: dict, name: str) -> list[Chunk]: | |
| kind = interface_kind(name) | |
| distro = meta.get("distro", "") | |
| body = text.strip() | |
| return [Chunk(f"[ROS 2 {distro} {kind} {name}]\n{body}", source, {**meta, "interface": name})] | |
| def chunk_file(path: str, text: str, source_name: str, meta: dict, package: str | None = None) -> list[Chunk]: | |
| source = f"{source_name}:{path}" | |
| suffix = PurePosixPath(path).suffix | |
| if suffix in (".msg", ".srv", ".action"): | |
| name = interface_name(path, package) | |
| return chunk_interface(text, source, meta, name) if name else [] | |
| if path.endswith("XMLschema.rst"): | |
| return [schema_chunk(k, e, source, meta) for k, e in parse_mjcf_schema(text).items()] | |
| if suffix in (".rst", ".md"): | |
| return chunk_document(text, source, meta, PurePosixPath(path).stem.replace("-", " ")) | |
| if suffix == ".xml" and "launch" not in path: | |
| return chunk_code(text, source, meta, f"MJCF example {path}") | |
| return chunk_code(text, source, meta, f"Example file {path}") | |
| # --------------------------------------------------------------------------- MJCF schema | |
| _DROPDOWN = re.compile(r"\.\. dropdown:: :ref:`([^<`]+)<([\w-]+)>`") | |
| _ATTR_REF = re.compile(r"^\s+:ref:`(\w+)<([\w-]+)>`\s*$") | |
| def parse_mjcf_schema(text: str) -> dict[str, dict]: | |
| """Parse MuJoCo's doc/XMLschema.rst into {ref_key: {"element", "path", "attributes"}}. | |
| ref_key is the doc anchor, e.g. "body-geom"; path is the parent chain, e.g. "body/geom". | |
| """ | |
| schema: dict[str, dict] = {} | |
| key = None | |
| for line in text.splitlines(): | |
| m = _DROPDOWN.search(line) | |
| if m: | |
| element, key = m.groups() | |
| element = element.replace("(world)", "") | |
| schema.setdefault(key, {"element": element, "path": key.replace("-", "/"), "attributes": []}) | |
| continue | |
| a = _ATTR_REF.match(line) | |
| if a and key is not None and a.group(2) == f"{key}-{a.group(1)}": | |
| schema[key]["attributes"].append(a.group(1)) | |
| return schema | |
| def schema_chunk(key: str, entry: dict, source: str, meta: dict) -> Chunk: | |
| attrs = ", ".join(entry["attributes"]) or "(no attributes)" | |
| text = f"[MJCF schema] <{entry['element']}> (path {entry['path']}): allowed attributes: {attrs}" | |
| return Chunk(text, f"{source}#{key}", {**meta, "kind": "schema", "element": entry["element"], "schema_key": key}) | |