Spaces:
Paused
Paused
Chandra Kiran
Harden simulation runs: official Foxglove embed, WebM video, CI image build
04ad412 unverified | """Local Chroma vector store with the two collections Stage 1 needs.""" | |
| from __future__ import annotations | |
| from collections.abc import Iterable, Sequence | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import chromadb | |
| try: # Chroma's local embedding model runs on ONNX Runtime; keep it from sending usage telemetry. | |
| import onnxruntime | |
| onnxruntime.disable_telemetry_events() | |
| except (ImportError, AttributeError): | |
| pass | |
| from .chunking import Chunk | |
| COLLECTIONS = ("mujoco", "ros2") | |
| class Hit: | |
| text: str | |
| source: str | |
| distance: float | |
| meta: dict | |
| class RagStore: | |
| """Thin wrapper around a persistent Chroma client. | |
| The default embedding is Chroma's bundled all-MiniLM-L6-v2 (ONNX, runs on | |
| CPU, downloaded once on first use). Tests pass their own ``embedding_function``. | |
| """ | |
| def __init__(self, path: str | Path, embedding_function=None): | |
| self.path = Path(path) | |
| self.path.mkdir(parents=True, exist_ok=True) | |
| self._client = chromadb.PersistentClient( | |
| path=str(self.path), settings=chromadb.config.Settings(anonymized_telemetry=False) | |
| ) | |
| self._ef = embedding_function | |
| def collection(self, name: str): | |
| if name not in COLLECTIONS: | |
| raise ValueError(f"unknown collection {name!r}; expected one of {COLLECTIONS}") | |
| kwargs = {"name": name, "metadata": {"hnsw:space": "cosine"}} | |
| if self._ef is not None: | |
| kwargs["embedding_function"] = self._ef | |
| return self._client.get_or_create_collection(**kwargs) | |
| def reset(self, name: str) -> None: | |
| try: | |
| self._client.delete_collection(name) | |
| except Exception: # collection did not exist | |
| pass | |
| def add(self, name: str, chunks: Iterable[Chunk], batch: int = 128) -> int: | |
| col = self.collection(name) | |
| unique = {c.id: c for c in chunks}.values() | |
| items = list(unique) | |
| for i in range(0, len(items), batch): | |
| part = items[i : i + batch] | |
| col.upsert( | |
| ids=[c.id for c in part], | |
| documents=[c.text for c in part], | |
| metadatas=[{**{k: str(v) for k, v in c.meta.items()}, "source": c.source} for c in part], | |
| ) | |
| return len(items) | |
| def count(self, name: str) -> int: | |
| return self.collection(name).count() | |
| def query(self, name: str, text: str, k: int = 6, where: dict | None = None) -> list[Hit]: | |
| col = self.collection(name) | |
| if col.count() == 0: | |
| return [] | |
| res = col.query(query_texts=[text], n_results=min(k, col.count()), where=where or None) | |
| hits = [] | |
| for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0], strict=True): | |
| hits.append(Hit(text=doc, source=meta.get("source", ""), distance=float(dist), meta=dict(meta))) | |
| return hits | |
| def format_context(hits: Sequence[Hit], max_chars: int = 14000) -> str: | |
| """Numbered context block for a prompt, truncated to a character budget.""" | |
| parts, used = [], 0 | |
| for i, h in enumerate(hits, 1): | |
| block = f'<doc id={i} source="{h.source}">\n{h.text}\n</doc>' | |
| if used + len(block) > max_chars: | |
| break | |
| parts.append(block) | |
| used += len(block) | |
| return "\n\n".join(parts) | |