| from __future__ import annotations |
|
|
| import json |
| import math |
| from collections import Counter |
| from pathlib import Path |
|
|
| from rag.schemas import DocumentChunk, RetrievalHit |
|
|
|
|
| class SparseKeywordIndex: |
| def __init__(self, index_path: str) -> None: |
| self.index_path = Path(index_path) |
| self.index_path.parent.mkdir(parents=True, exist_ok=True) |
| self._records: dict[str, dict] = self._load() |
|
|
| def upsert(self, chunks: list[DocumentChunk]) -> None: |
| for chunk in chunks: |
| term_counts = Counter(self._tokenize(chunk.text)) |
| self._records[chunk.chunk_id] = { |
| "document_id": chunk.document_id, |
| "text": chunk.text, |
| "metadata": chunk.metadata, |
| "term_counts": dict(term_counts), |
| "length": sum(term_counts.values()), |
| } |
| self._save() |
|
|
| def search(self, query: str, top_k: int = 5) -> list[RetrievalHit]: |
| query_terms = self._tokenize(query) |
| if not query_terms or not self._records: |
| return [] |
|
|
| doc_freq = Counter() |
| for record in self._records.values(): |
| for token in set(record["term_counts"]): |
| doc_freq[token] += 1 |
|
|
| total_docs = max(len(self._records), 1) |
| avg_doc_length = sum(record["length"] for record in self._records.values()) / total_docs |
| scores: list[tuple[float, str, dict]] = [] |
|
|
| for chunk_id, record in self._records.items(): |
| score = 0.0 |
| for term in query_terms: |
| tf = record["term_counts"].get(term, 0) |
| if not tf: |
| continue |
| idf = math.log(1 + (total_docs - doc_freq[term] + 0.5) / (doc_freq[term] + 0.5)) |
| score += idf * ((tf * 2.2) / (tf + 1.2 * (1 - 0.75 + 0.75 * record["length"] / max(avg_doc_length, 1)))) |
| if score > 0: |
| scores.append((score, chunk_id, record)) |
|
|
| scores.sort(key=lambda item: item[0], reverse=True) |
| hits = [] |
| for score, chunk_id, record in scores[:top_k]: |
| hits.append( |
| RetrievalHit( |
| chunk_id=chunk_id, |
| text=record["text"], |
| score=score, |
| metadata=record["metadata"], |
| ) |
| ) |
| return hits |
|
|
| def list_documents(self) -> list[dict]: |
| seen: dict[str, dict] = {} |
| for record in self._records.values(): |
| document_id = record.get("document_id", "") |
| metadata = record.get("metadata", {}) |
| if document_id and document_id not in seen: |
| seen[document_id] = { |
| "document_id": document_id, |
| "file_name": metadata.get("file_name", ""), |
| } |
| return sorted(seen.values(), key=lambda item: item["file_name"]) |
|
|
| def _load(self) -> dict[str, dict]: |
| if not self.index_path.exists(): |
| return {} |
| try: |
| return json.loads(self.index_path.read_text(encoding="utf-8")) |
| except Exception: |
| return {} |
|
|
| def _save(self) -> None: |
| self.index_path.write_text(json.dumps(self._records, ensure_ascii=True, indent=2), encoding="utf-8") |
|
|
| @staticmethod |
| def _tokenize(text: str) -> list[str]: |
| return [token.strip(".,:;!?()[]{}\"'").lower() for token in text.split() if token.strip()] |
|
|
