"""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) @staticmethod 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"}, )