Coverage for src/asketmc_bot/index_builder.py: 0%
153 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1#!/usr/bin/env python3.10
2"""
3index_builder.py — embeddings configuration and LlamaIndex index build/load.
5Exported API:
6- build_index() -> VectorStoreIndex (async)
8Responsibilities:
9- configures LlamaIndex Settings.embed_model (BAAI/bge-m3) and SentenceSplitter
10- loads or rebuilds index based on document hashes
11- maintains file-lemma cache via asketmc_bot.lemma
12- annotates nodes with metadata["lemmas"]
13"""
15from __future__ import annotations
17# Stdlib
18import hashlib
19import json
20import logging
21import time
22from pathlib import Path
23from typing import Dict, List, Optional, Set, Tuple
25# Third-party
26try:
27 import torch # type: ignore
28except Exception: # ImportError + binary load errors
29 torch = None # type: ignore[assignment]
31from llama_index.core import Settings, StorageContext, VectorStoreIndex, load_index_from_storage
32from llama_index.core.node_parser import SentenceSplitter
33from llama_index.core.schema import Document
34from llama_index.embeddings.huggingface import HuggingFaceEmbedding
36# Local (src-layout safe)
37from asketmc_bot import config as cfg
38from asketmc_bot.lemma import (
39 extract_lemmas,
40 FILE_LEMMAS,
41 load_saved_lemmas,
42 update_file_lemmas_async,
43 save_chunk_lemma_cache,
44)
46log = logging.getLogger("asketmc.index")
48# ─────────────────────────────────────────────────────────
49# Embeddings configuration
50# ─────────────────────────────────────────────────────────
51def _require_torch() -> None:
52 if torch is None:
53 raise RuntimeError(
54 "Torch is not available. Install extra dependencies, e.g. "
55 "`pip install -e '.[gpu]'` (or move torch to an optional extra)."
56 )
59_require_torch()
60DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # type: ignore[union-attr]
61EMBED_LOG_EVERY = getattr(cfg, "EMBED_LOG_EVERY", 1000)
64class LoggingBGE(HuggingFaceEmbedding):
65 _counter: int = 0
66 __slots__ = ()
68 def _get_text_embedding(self, text: str):
69 t0 = time.perf_counter()
70 vec = super()._get_text_embedding(text)
71 type(self)._counter += 1
72 if type(self)._counter % EMBED_LOG_EVERY == 0:
73 logging.getLogger("asketmc.embed").info(
74 "EMB %d | %s | %s… | %.3fs",
75 type(self)._counter,
76 self._target_device,
77 text[:120].replace("\n", " "),
78 time.perf_counter() - t0,
79 )
80 return vec
83_SETTINGS_CONFIGURED = False
86def _configure_llamaindex_settings() -> None:
87 global _SETTINGS_CONFIGURED
88 if _SETTINGS_CONFIGURED:
89 return
91 Settings.embed_model = LoggingBGE("BAAI/bge-m3", normalize=True, device=DEVICE)
92 Settings.node_parser = SentenceSplitter(
93 chunk_size=getattr(cfg, "CHUNK_SIZE", 512),
94 chunk_overlap=getattr(cfg, "CHUNK_OVERLAP", 128),
95 include_metadata=False,
96 paragraph_separator="\n\n",
97 )
98 _SETTINGS_CONFIGURED = True
101# ─────────────────────────────────────────────────────────
102# Document helpers
103# ─────────────────────────────────────────────────────────
104def _make_document(fp: Path) -> Optional[Document]:
105 """
106 Create a Document with a stable id_ so we can delete/update incrementally.
107 Using filename as id_ assumes DOCS_PATH has unique filenames.
108 """
109 try:
110 text = fp.read_text("utf-8", "ignore")
111 except Exception as e:
112 log.warning("[build_index] Read error %s: %s", fp, e)
113 return None
114 return Document(text=text, metadata={"file_name": fp.name}, id_=fp.name)
117# ─────────────────────────────────────────────────────────
118# Helpers
119# ─────────────────────────────────────────────────────────
120def _doc_hash(fp: Path) -> str:
121 """SHA-256 of a file read by 1MB chunks."""
122 h = hashlib.sha256()
123 with fp.open("rb") as f:
124 for chunk in iter(lambda: f.read(1 << 20), b""):
125 h.update(chunk)
126 return h.hexdigest()
129def _load_stored_hashes() -> Dict[str, str]:
130 if not cfg.HASH_FILE.exists():
131 return {}
132 try:
133 return json.loads(cfg.HASH_FILE.read_text("utf-8"))
134 except Exception:
135 return {}
138def _compute_hashes_incremental(docs: List[Path], stored: Dict[str, str]) -> Dict[str, str]:
139 """
140 Reduce hashing cost:
141 - If HASH_FILE exists and is newer than doc mtime, reuse stored hash.
142 - Otherwise compute sha256.
143 """
144 hashes: Dict[str, str] = {}
145 hash_file_mtime_ns = cfg.HASH_FILE.stat().st_mtime_ns if cfg.HASH_FILE.exists() else 0
147 for fp in docs:
148 try:
149 mtime_ns = fp.stat().st_mtime_ns
150 except Exception:
151 mtime_ns = 0
153 if fp.name in stored and mtime_ns and mtime_ns <= hash_file_mtime_ns:
154 hashes[fp.name] = stored[fp.name]
155 else:
156 hashes[fp.name] = _doc_hash(fp)
158 return hashes
161def _plan_doc_changes(stored: Dict[str, str], current: Dict[str, str]) -> Tuple[Set[str], Set[str], Set[str]]:
162 """
163 Return (added, removed, modified) doc ids (filenames).
164 """
165 stored_keys = set(stored.keys())
166 current_keys = set(current.keys())
168 added = current_keys - stored_keys
169 removed = stored_keys - current_keys
170 modified = {k for k in (stored_keys & current_keys) if stored.get(k) != current.get(k)}
172 return added, removed, modified
175def _load_index_if_present() -> Optional[VectorStoreIndex]:
176 if not cfg.CACHE_PATH.exists():
177 return None
178 try:
179 return load_index_from_storage(StorageContext.from_defaults(persist_dir=str(cfg.CACHE_PATH)))
180 except Exception as e:
181 log.warning("[build_index] Failed to load cached index: %s", e)
182 return None
185# ─────────────────────────────────────────────────────────
186# Build / Load index
187# ─────────────────────────────────────────────────────────
188async def build_index() -> VectorStoreIndex:
189 """
190 Build or load index of documents from cfg.DOCS_PATH.
191 - Validates hashes (cfg.HASH_FILE)
192 - Loads FILE_LEMMAS cache and updates it for changed files
193 - Annotates nodes with metadata["lemmas"]
194 - Persists index, hashes, and chunk lemma cache
195 """
196 _configure_llamaindex_settings()
198 docs: List[Path] = [p for p in cfg.DOCS_PATH.glob("*") if p.is_file()]
199 stored: Dict[str, str] = _load_stored_hashes()
200 hashes: Dict[str, str] = _compute_hashes_incremental(docs, stored)
202 # 1) File-lemma cache
203 load_saved_lemmas()
204 try:
205 await update_file_lemmas_async(docs, stored, hashes)
206 except Exception as e:
207 log.warning("[build_index] update_file_lemmas_async failed: %s", e)
209 # 2) Index: load if present, then apply incremental changes
210 idx: Optional[VectorStoreIndex] = _load_index_if_present()
211 if idx is not None:
212 log.info("[build_index] Loaded index from cache: %s", cfg.CACHE_PATH)
214 added, removed, modified = _plan_doc_changes(stored, hashes)
215 if added or removed or modified:
216 log.info(
217 "[build_index] Incremental update | added=%d removed=%d modified=%d",
218 len(added),
219 len(removed),
220 len(modified),
221 )
223 # removals first
224 for doc_id in sorted(removed):
225 try:
226 idx.delete_ref_doc(doc_id, delete_from_docstore=True)
227 except Exception as e:
228 log.warning("[build_index] delete_ref_doc failed (%s): %s", doc_id, e)
230 # updates
231 for doc_id in sorted(modified):
232 fp = cfg.DOCS_PATH / doc_id
233 doc = _make_document(fp)
234 if doc is None:
235 continue
236 try:
237 idx.update_ref_doc(doc)
238 except Exception as e:
239 log.warning("[build_index] update_ref_doc failed (%s): %s", doc_id, e)
241 # inserts
242 for doc_id in sorted(added):
243 fp = cfg.DOCS_PATH / doc_id
244 doc = _make_document(fp)
245 if doc is None:
246 continue
247 try:
248 idx.insert(doc)
249 except Exception as e:
250 log.warning("[build_index] insert failed (%s): %s", doc_id, e)
252 # cache missing or load failed -> build from scratch
253 if idx is None:
254 ll_docs: List[Document] = []
255 for fp in docs:
256 doc = _make_document(fp)
257 if doc is not None:
258 ll_docs.append(doc)
260 idx = VectorStoreIndex.from_documents(ll_docs)
261 log.info("[build_index] Built new index from %d docs", len(ll_docs))
263 # 3) Annotate lemmas in node metadata
264 for node in idx.docstore.docs.values():
265 fname = node.metadata.get("file_name")
266 if not fname:
267 continue
268 lem = FILE_LEMMAS.get(fname)
269 if lem is None:
270 lem = extract_lemmas(node.get_content())
271 node.metadata["lemmas"] = list(lem)
273 # 4) Persist index, hashes, chunk-lemma-cache
274 try:
275 idx.storage_context.persist(str(cfg.CACHE_PATH))
276 cfg.HASH_FILE.parent.mkdir(parents=True, exist_ok=True)
277 cfg.HASH_FILE.write_text(json.dumps(hashes, ensure_ascii=False, indent=2), encoding="utf-8")
278 save_chunk_lemma_cache()
279 log.info("[build_index] Persisted index and hashes.")
280 except Exception as e:
281 log.warning("[build_index] Persist failed: %s", e)
283 return idx