| """RAG knowledge base: ingest, search, no hallucination.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import logging |
| import uuid |
| from pathlib import Path |
| from typing import Any |
|
|
| from bs4 import BeautifulSoup |
| from pypdf import PdfReader |
|
|
| from app.config import get_settings |
| from app.services.embeddings import ( |
| chunk_text, |
| cosine_similarity, |
| embed_texts, |
| token_overlap_score, |
| ) |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class KnowledgeStore: |
| """File-backed vector store (Chroma when available).""" |
|
|
| def __init__(self) -> None: |
| self.settings = get_settings() |
| self._collection = None |
| self._fallback_path = ( |
| Path(self.settings.chroma_path) / "fallback.json" |
| ) |
| self._fallback: list[dict[str, Any]] = [] |
| self._init_store() |
|
|
| def _init_store(self) -> None: |
| path = Path(self.settings.chroma_path) |
| path.mkdir(parents=True, exist_ok=True) |
| if self.settings.use_json_store: |
| self._collection = None |
| self._load_fallback() |
| logger.info("JSON knowledge store ready") |
| return |
| try: |
| import chromadb |
|
|
| client = chromadb.PersistentClient(path=str(path)) |
| self._collection = client.get_or_create_collection( |
| name="replyflow_kb", |
| metadata={"hnsw:space": "cosine"}, |
| ) |
| logger.info("Chroma knowledge store ready") |
| except Exception as exc: |
| logger.warning( |
| "Chroma unavailable (%s); using JSON fallback", |
| exc, |
| ) |
| self._collection = None |
| self._load_fallback() |
|
|
| def _load_fallback(self) -> None: |
| if self._fallback_path.exists(): |
| self._fallback = json.loads( |
| self._fallback_path.read_text(encoding="utf-8") |
| ) |
| else: |
| self._fallback = [] |
|
|
| def _save_fallback(self) -> None: |
| self._fallback_path.parent.mkdir( |
| parents=True, exist_ok=True |
| ) |
| self._fallback_path.write_text( |
| json.dumps(self._fallback, ensure_ascii=False), |
| encoding="utf-8", |
| ) |
|
|
| @property |
| def chunk_count(self) -> int: |
| if self._collection is not None: |
| return int(self._collection.count()) |
| return len(self._fallback) |
|
|
| async def add_chunks( |
| self, |
| chunks: list[str], |
| metadata: dict[str, Any], |
| ) -> int: |
| """Embed and store text chunks.""" |
| if not chunks: |
| return 0 |
| embeddings = await embed_texts( |
| chunks, |
| force_local=self.settings.seed_local_embeddings, |
| ) |
| ids = [str(uuid.uuid4()) for _ in chunks] |
| metadatas = [ |
| {**metadata, "chunk_index": i} |
| for i in range(len(chunks)) |
| ] |
|
|
| if self._collection is not None: |
| self._collection.add( |
| ids=ids, |
| documents=chunks, |
| embeddings=embeddings, |
| metadatas=metadatas, |
| ) |
| else: |
| for i, chunk in enumerate(chunks): |
| self._fallback.append( |
| { |
| "id": ids[i], |
| "document": chunk, |
| "embedding": embeddings[i], |
| "metadata": metadatas[i], |
| } |
| ) |
| self._save_fallback() |
| return len(chunks) |
|
|
| async def search( |
| self, |
| query: str, |
| top_k: int = 4, |
| min_score: float = 0.12, |
| ) -> list[dict[str, Any]]: |
| """Return relevant chunks above similarity threshold.""" |
| if self.chunk_count == 0: |
| return [] |
|
|
| query_emb = ( |
| await embed_texts( |
| [query], |
| force_local=self.settings.seed_local_embeddings, |
| ) |
| )[0] |
| use_openai = ( |
| get_settings().has_openai |
| and not self.settings.seed_local_embeddings |
| ) |
|
|
| if self._collection is not None: |
| result = self._collection.query( |
| query_embeddings=[query_emb], |
| n_results=min(top_k, max(1, self.chunk_count)), |
| ) |
| docs = result.get("documents", [[]])[0] |
| metas = result.get("metadatas", [[]])[0] |
| dists = result.get("distances", [[]])[0] |
| hits: list[dict[str, Any]] = [] |
| for doc, meta, dist in zip(docs, metas, dists): |
| vector_score = 1.0 - float(dist) |
| keyword = token_overlap_score(query, doc) |
| score = ( |
| vector_score |
| if use_openai |
| else 0.45 * vector_score + 0.55 * keyword |
| ) |
| if score >= min_score: |
| hits.append( |
| { |
| "text": doc, |
| "metadata": meta or {}, |
| "score": score, |
| } |
| ) |
| hits.sort(key=lambda x: x["score"], reverse=True) |
| return hits[:top_k] |
|
|
| scored = [] |
| for item in self._fallback: |
| vector_score = cosine_similarity( |
| query_emb, item["embedding"] |
| ) |
| keyword = token_overlap_score(query, item["document"]) |
| score = 0.35 * vector_score + 0.65 * keyword |
| if score >= min_score: |
| scored.append( |
| { |
| "text": item["document"], |
| "metadata": item["metadata"], |
| "score": score, |
| } |
| ) |
| scored.sort(key=lambda x: x["score"], reverse=True) |
| return scored[:top_k] |
|
|
|
|
| _store: KnowledgeStore | None = None |
|
|
|
|
| def get_knowledge_store() -> KnowledgeStore: |
| """Singleton knowledge store.""" |
| global _store |
| if _store is None: |
| _store = KnowledgeStore() |
| return _store |
|
|
|
|
| def extract_text_from_file(path: Path) -> str: |
| """Extract text from pdf, html, md, or txt.""" |
| suffix = path.suffix.lower() |
| raw = path.read_bytes() |
| if suffix == ".pdf": |
| reader = PdfReader(str(path)) |
| parts = [] |
| for page in reader.pages: |
| parts.append(page.extract_text() or "") |
| return "\n".join(parts) |
| text = raw.decode("utf-8", errors="ignore") |
| if suffix in {".html", ".htm"}: |
| soup = BeautifulSoup(text, "html.parser") |
| return soup.get_text(separator="\n") |
| return text |
|
|
|
|
| async def ingest_path(path: Path, source_type: str) -> int: |
| """Ingest a single file into the knowledge store.""" |
| text = extract_text_from_file(path) |
| chunks = chunk_text(text) |
| store = get_knowledge_store() |
| return await store.add_chunks( |
| chunks, |
| { |
| "filename": path.name, |
| "source_type": source_type, |
| }, |
| ) |
|
|
|
|
| async def ingest_directory(directory: Path) -> int: |
| """Ingest all supported files from a directory.""" |
| total = 0 |
| patterns = ("*.md", "*.txt", "*.pdf", "*.html", "*.htm") |
| files: list[Path] = [] |
| for pattern in patterns: |
| files.extend(directory.glob(pattern)) |
| for path in sorted(set(files)): |
| source = path.suffix.lower().lstrip(".") or "text" |
| count = await ingest_path(path, source) |
| total += count |
| logger.info("Ingested %s (%s chunks)", path.name, count) |
| return total |
|
|
