Spaces:
Sleeping
Sleeping
File size: 6,373 Bytes
4b81334 6be14c8 4b81334 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """RAG engine: PyMuPDF4LLM -> chunks -> HF-API embeddings -> LanceDB.
Designed for HF Spaces Free CPU Basic:
* No local embedding/LLM weights are downloaded.
* LanceDB lives on ephemeral disk locally, but its files are mirrored into
the linked HF Dataset so cold starts can restore the vector cache instead
of rebuilding from PDFs every time.
"""
from __future__ import annotations
import logging
from pathlib import Path
from threading import Lock
from typing import Iterable, List, Optional
import pymupdf4llm
from llama_index.core import Document, StorageContext, VectorStoreIndex, Settings
from llama_index.core.node_parser import MarkdownNodeParser
from llama_index.core.schema import MetadataMode, NodeWithScore, TextNode
from llama_index.vector_stores.lancedb import LanceDBVectorStore
from config import EMBED_MODEL, HF_TOKEN, LANCEDB_PATH
from hf_embedding import SyncHuggingFaceInferenceEmbedding
logger = logging.getLogger(__name__)
TABLE_NAME = "documents"
# Filename is stored on every node so we can delete by source.
FILENAME_KEY = "source_filename"
class RagEngine:
"""Thread-safe singleton wrapper around the LanceDB-backed index."""
def __init__(self) -> None:
self._lock = Lock()
self._index: Optional[VectorStoreIndex] = None
self._vector_store: Optional[LanceDBVectorStore] = None
if not HF_TOKEN:
logger.warning(
"HF_TOKEN is not set — embedding calls will fail. "
"Set it in .env before uploading or querying documents."
)
Settings.embed_model = None # type: ignore[assignment]
else:
Settings.embed_model = SyncHuggingFaceInferenceEmbedding(
model_name=EMBED_MODEL,
token=HF_TOKEN,
)
Settings.llm = None # We call the LLM ourselves in main.py for streaming.
Settings.node_parser = MarkdownNodeParser()
Path(LANCEDB_PATH).mkdir(parents=True, exist_ok=True)
self._vector_store = LanceDBVectorStore(
uri=LANCEDB_PATH,
table_name=TABLE_NAME,
mode="overwrite" if not self._table_exists() else "append",
)
# ---------- internals ----------
def _table_exists(self) -> bool:
import lancedb
try:
db = lancedb.connect(LANCEDB_PATH)
return TABLE_NAME in db.table_names()
except Exception: # noqa: BLE001
return False
def _ensure_index(self) -> VectorStoreIndex:
if self._index is None:
storage = StorageContext.from_defaults(vector_store=self._vector_store)
if self._table_exists():
self._index = VectorStoreIndex.from_vector_store(
vector_store=self._vector_store,
storage_context=storage,
)
else:
self._index = VectorStoreIndex.from_documents(
[], storage_context=storage
)
return self._index
@staticmethod
def _pdf_to_documents(pdf_path: Path, filename: str) -> List[Document]:
try:
md_text = pymupdf4llm.to_markdown(str(pdf_path))
except Exception as exc: # noqa: BLE001
logger.warning("Failed to parse %s: %s", filename, exc)
return []
if not md_text.strip():
return []
return [
Document(
text=md_text,
metadata={FILENAME_KEY: filename},
excluded_llm_metadata_keys=[FILENAME_KEY],
excluded_embed_metadata_keys=[FILENAME_KEY],
)
]
# ---------- public API ----------
def index_pdf(self, pdf_path: Path, filename: str) -> int:
"""Parse a single PDF and insert its chunks. Returns # nodes added."""
with self._lock:
# First, drop any existing nodes for this filename (re-uploads).
self.delete_by_filename(filename, _locked=True)
docs = self._pdf_to_documents(pdf_path, filename)
if not docs:
return 0
index = self._ensure_index()
nodes = Settings.node_parser.get_nodes_from_documents(docs)
for n in nodes:
n.metadata[FILENAME_KEY] = filename
index.insert_nodes(nodes)
return len(nodes)
def index_many(self, items: Iterable[tuple[Path, str]]) -> int:
total = 0
for path, name in items:
total += self.index_pdf(path, name)
return total
def delete_by_filename(self, filename: str, *, _locked: bool = False) -> int:
"""Remove all nodes with metadata.source_filename == filename."""
def _do() -> int:
import lancedb
if not self._table_exists():
return 0
db = lancedb.connect(LANCEDB_PATH)
tbl = db.open_table(TABLE_NAME)
# LanceDB stores metadata as a struct column called "metadata".
# Filter accesses nested fields with dot syntax.
try:
before = tbl.count_rows()
tbl.delete(f"metadata.{FILENAME_KEY} = '{filename}'")
removed = before - tbl.count_rows()
except Exception as exc: # noqa: BLE001
logger.warning("delete_by_filename failed: %s", exc)
removed = 0
# Force the index to be rebuilt next query.
self._index = None
return removed
if _locked:
return _do()
with self._lock:
return _do()
def retrieve(self, query: str, top_k: int = 4) -> List[NodeWithScore]:
with self._lock:
if not self._table_exists():
return []
index = self._ensure_index()
retriever = index.as_retriever(similarity_top_k=top_k)
return retriever.retrieve(query)
@staticmethod
def format_context(nodes: List[NodeWithScore]) -> str:
if not nodes:
return ""
chunks = []
for i, n in enumerate(nodes, 1):
src = n.node.metadata.get(FILENAME_KEY, "unknown")
text = n.node.get_content(metadata_mode=MetadataMode.NONE).strip()
chunks.append(f"[{i}] (source: {src})\n{text}")
return "\n\n---\n\n".join(chunks)
|