Spaces:
Sleeping
Sleeping
| """Local knowledge base: manifest-driven ingest, ruri-small embedding, ChromaDB Top-K.""" | |
| from __future__ import annotations | |
| import re | |
| from pathlib import Path | |
| from typing import Any, cast | |
| import chromadb | |
| import numpy as np | |
| import yaml | |
| from chromadb.api import ClientAPI | |
| from chromadb.api.types import Embedding, Metadata | |
| from sentence_transformers import SentenceTransformer | |
| from src.models import RAGConfig, RetrievedChunk | |
| ROOT = Path(__file__).resolve().parents[1] | |
| SAMPLE_DOCS_DIR = ROOT / "sample_docs" | |
| MANIFEST_PATH = SAMPLE_DOCS_DIR / "manifest.yaml" | |
| # ruri models expect task prefixes for best retrieval quality. | |
| _QUERY_PREFIX = "クエリ: " | |
| _PASSAGE_PREFIX = "文章: " | |
| _FRONTMATTER_RE = re.compile(r"^---\n.*?\n---\n", re.DOTALL) | |
| def _strip_frontmatter(text: str) -> str: | |
| return _FRONTMATTER_RE.sub("", text, count=1) | |
| def _chunk_text(text: str, chunk_size: int, overlap: int) -> list[str]: | |
| """Recursive-ish character chunking with overlap (Cookbook ~200-500 chars).""" | |
| text = text.strip() | |
| if not text: | |
| return [] | |
| chunks: list[str] = [] | |
| start = 0 | |
| length = len(text) | |
| while start < length: | |
| end = min(start + chunk_size, length) | |
| chunks.append(text[start:end]) | |
| if end == length: | |
| break | |
| start = end - overlap | |
| return chunks | |
| class LocalKnowledgeBase: | |
| """Ingest sample_docs into ChromaDB and serve Top-K retrieval, 100% locally.""" | |
| def __init__(self, rag: RAGConfig) -> None: | |
| self._rag = rag | |
| self._embedder = SentenceTransformer(rag.embedding_model, trust_remote_code=True) | |
| self._client: ClientAPI = chromadb.PersistentClient(path=str(ROOT / rag.persist_dir)) | |
| self._collection = self._client.get_or_create_collection( | |
| name=rag.collection_name, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| def count(self) -> int: | |
| return self._collection.count() | |
| def _embed_passages(self, texts: list[str]) -> list[list[float]]: | |
| vectors = self._embedder.encode( # pyright: ignore[reportUnknownMemberType] | |
| [_PASSAGE_PREFIX + t for t in texts], | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| ) | |
| return cast(list[list[float]], vectors.tolist()) | |
| def _embed_query(self, text: str) -> list[float]: | |
| vector = self._embedder.encode( # pyright: ignore[reportUnknownMemberType] | |
| [_QUERY_PREFIX + text], | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| ) | |
| return cast(list[float], vector[0].tolist()) | |
| def load_manifest(self) -> dict[str, Any]: | |
| with MANIFEST_PATH.open("r", encoding="utf-8") as fh: | |
| return yaml.safe_load(fh) | |
| def ingest_manifest(self, reset: bool = True) -> int: | |
| """Ingest all status=ready documents listed in manifest.yaml. Returns chunk count.""" | |
| manifest = self.load_manifest() | |
| if reset: | |
| self._reset_collection() | |
| total = 0 | |
| for doc in manifest.get("documents", []): | |
| if doc.get("status") != "ready": | |
| continue | |
| path = SAMPLE_DOCS_DIR / doc["path"] | |
| if not path.exists(): | |
| continue | |
| total += self._ingest_file( | |
| path=path, | |
| doc_id=doc["doc_id"], | |
| tier=doc.get("tier", "unknown"), | |
| synthetic=bool(doc.get("synthetic", False)), | |
| source_url=doc.get("source_url"), | |
| ) | |
| return total | |
| def ingest_text(self, text: str, doc_id: str, tier: str = "uploaded", synthetic: bool = False) -> int: | |
| """Ingest raw text (e.g. uploaded file) into the collection.""" | |
| return self._add_chunks(text, doc_id, tier, synthetic, None) | |
| def _ingest_file(self, path: Path, doc_id: str, tier: str, synthetic: bool, source_url: str | None) -> int: | |
| raw = path.read_text(encoding="utf-8") | |
| body = _strip_frontmatter(raw) if path.suffix == ".md" else raw | |
| return self._add_chunks(body, doc_id, tier, synthetic, source_url) | |
| def _add_chunks(self, body: str, doc_id: str, tier: str, synthetic: bool, source_url: str | None) -> int: | |
| chunks = _chunk_text(body, self._rag.chunk_size, self._rag.chunk_overlap) | |
| if not chunks: | |
| return 0 | |
| embeddings = self._embed_passages(chunks) | |
| ids = [f"{doc_id}::{i}" for i in range(len(chunks))] | |
| metadatas: list[dict[str, Any]] = [ | |
| { | |
| "doc_id": doc_id, | |
| "tier": tier, | |
| "synthetic": synthetic, | |
| "source_url": source_url or "", | |
| } | |
| for _ in chunks | |
| ] | |
| self._collection.upsert( | |
| ids=ids, | |
| documents=chunks, | |
| embeddings=cast(list[Embedding], embeddings), | |
| metadatas=cast(list[Metadata], metadatas), | |
| ) | |
| return len(chunks) | |
| def retrieve( | |
| self, | |
| query: str, | |
| top_k: int | None = None, | |
| allowed_tiers: set[str] | None = None, | |
| use_mmr: bool = True, | |
| mmr_lambda: float = 0.6, | |
| fetch_multiplier: int = 4, | |
| ) -> list[RetrievedChunk]: | |
| """Multi-stage retrieval: over-fetch -> tier filter -> MMR re-rank -> Top-K. | |
| Args: | |
| allowed_tiers: if set, drop candidates whose tier is not in this set. | |
| This is the data-egress guard (e.g. SAFE route -> no confidential). | |
| use_mmr: apply Maximal Marginal Relevance for relevance/diversity balance. | |
| mmr_lambda: 1.0 = pure relevance, 0.0 = pure diversity. | |
| """ | |
| if self._collection.count() == 0: | |
| return [] | |
| k = top_k or self._rag.top_k | |
| fetch_k = max(k * fetch_multiplier, k) | |
| query_vec = np.asarray(self._embed_query(query), dtype=np.float32) | |
| result = self._collection.query( | |
| query_embeddings=[query_vec.tolist()], | |
| n_results=fetch_k, | |
| include=["documents", "metadatas", "distances", "embeddings"], | |
| ) | |
| doc_rows = result.get("documents") | |
| meta_rows = result.get("metadatas") | |
| dist_rows = result.get("distances") | |
| emb_rows = result.get("embeddings") | |
| if not doc_rows or not meta_rows or not dist_rows or not emb_rows: | |
| return [] | |
| docs = doc_rows[0] | |
| metas = meta_rows[0] | |
| dists = dist_rows[0] | |
| embs = emb_rows[0] | |
| candidates: list[tuple[RetrievedChunk, np.ndarray]] = [] | |
| for text, meta, dist, emb in zip(docs, metas, dists, embs, strict=False): | |
| tier = str(meta.get("tier", "")) | |
| if allowed_tiers is not None and tier not in allowed_tiers: | |
| continue | |
| chunk = RetrievedChunk( | |
| text=text, | |
| score=round(1.0 - float(dist), 4), | |
| doc_id=str(meta.get("doc_id", "")), | |
| tier=tier, | |
| synthetic=bool(meta.get("synthetic", False)), | |
| source_url=str(meta.get("source_url") or "") or None, | |
| ) | |
| candidates.append((chunk, np.asarray(emb, dtype=np.float32))) | |
| if not candidates: | |
| return [] | |
| if not use_mmr: | |
| return [c for c, _ in candidates[:k]] | |
| return self._mmr_select(query_vec, candidates, k, mmr_lambda) | |
| def _mmr_select( | |
| query_vec: np.ndarray, | |
| candidates: list[tuple[RetrievedChunk, np.ndarray]], | |
| k: int, | |
| lam: float, | |
| ) -> list[RetrievedChunk]: | |
| """Greedy MMR selection over normalized candidate embeddings.""" | |
| chunks = [c for c, _ in candidates] | |
| mat = np.vstack([e for _, e in candidates]) | |
| rel = mat @ query_vec # embeddings are L2-normalized -> cosine | |
| selected: list[int] = [] | |
| remaining = set(range(len(chunks))) | |
| while remaining and len(selected) < k: | |
| best_idx, best_score = None, -np.inf | |
| for i in remaining: | |
| diversity = max((float(mat[i] @ mat[j]) for j in selected), default=0.0) | |
| score = lam * float(rel[i]) - (1.0 - lam) * diversity | |
| if score > best_score: | |
| best_idx, best_score = i, score | |
| assert best_idx is not None | |
| selected.append(best_idx) | |
| remaining.discard(best_idx) | |
| return [chunks[i] for i in selected] | |
| def _reset_collection(self) -> None: | |
| self._client.delete_collection(self._rag.collection_name) | |
| self._collection = self._client.get_or_create_collection( | |
| name=self._rag.collection_name, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |