111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Per-scope BM25 index (rebuild on change; bm25s has no cheap incremental insert).
|
|
|
|
Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from loggers import get_logger
|
|
from utils.paths.storage_roots import ensure_dir, rag_bm25_root
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_load_lock = threading.Lock()
|
|
_cache: dict[str, tuple[Any, list[str]]] = {}
|
|
|
|
|
|
def _scope_dir(scope: str) -> Path:
|
|
return rag_bm25_root() / scope
|
|
|
|
|
|
def _ids_path(scope: str) -> Path:
|
|
return _scope_dir(scope) / "ids.json"
|
|
|
|
|
|
def _has_index(scope: str) -> bool:
|
|
return _ids_path(scope).is_file()
|
|
|
|
|
|
def _evict(scope: str) -> None:
|
|
_cache.pop(scope, None)
|
|
|
|
|
|
def rebuild_index(scope: str, chunks: list[dict]) -> None:
|
|
"""Rebuild scope's BM25 from full chunk list. Empty list deletes the index."""
|
|
import bm25s
|
|
|
|
base = _scope_dir(scope)
|
|
if not chunks:
|
|
delete_scope(scope)
|
|
return
|
|
texts = [c["text"] for c in chunks]
|
|
ids = [c["id"] for c in chunks]
|
|
tokens = bm25s.tokenize(texts, show_progress = False)
|
|
retriever = bm25s.BM25()
|
|
retriever.index(tokens, show_progress = False)
|
|
# bm25s.BM25.save does not unlink stale files; clear the dir first.
|
|
delete_scope(scope)
|
|
ensure_dir(base)
|
|
retriever.save(str(base))
|
|
_ids_path(scope).write_text(json.dumps(ids))
|
|
with _load_lock:
|
|
_cache[scope] = (retriever, ids)
|
|
|
|
|
|
def _load(scope: str) -> tuple[Any, list[str]] | None:
|
|
if not _has_index(scope):
|
|
return None
|
|
with _load_lock:
|
|
if scope in _cache:
|
|
return _cache[scope]
|
|
import bm25s
|
|
|
|
try:
|
|
retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False)
|
|
ids = json.loads(_ids_path(scope).read_text())
|
|
except (FileNotFoundError, OSError, json.JSONDecodeError, ValueError) as exc:
|
|
# Corrupt/partial index: treat as missing so re-ingest rebuilds cleanly.
|
|
logger.warning(
|
|
"bm25 index unreadable for scope %s (%s: %s); treating as missing",
|
|
scope,
|
|
type(exc).__name__,
|
|
exc,
|
|
)
|
|
return None
|
|
_cache[scope] = (retriever, ids)
|
|
return _cache[scope]
|
|
|
|
|
|
def search(scope: str, query: str, k: int) -> list[tuple[str, float]]:
|
|
import bm25s
|
|
|
|
loaded = _load(scope)
|
|
if loaded is None:
|
|
return []
|
|
retriever, ids = loaded
|
|
if not ids:
|
|
return []
|
|
k_actual = min(k, len(ids))
|
|
q_tokens = bm25s.tokenize([query], show_progress = False)
|
|
indices, scores = retriever.retrieve(q_tokens, k = k_actual, show_progress = False)
|
|
out: list[tuple[str, float]] = []
|
|
for pos in range(indices.shape[1]):
|
|
idx = int(indices[0][pos])
|
|
out.append((ids[idx], float(scores[0][pos])))
|
|
return out
|
|
|
|
|
|
def delete_scope(scope: str) -> None:
|
|
base = _scope_dir(scope)
|
|
if base.exists():
|
|
shutil.rmtree(base, ignore_errors = True)
|
|
_evict(scope)
|