studyhub/core/loader.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from ..ingest.index import load_faiss
from ..ingest.models import Chunk
from ..ingest.store import read_chunks, read_json


@dataclass
class RetrievalIndex:
    chunks: list[Chunk]
    faiss_index: object
    concept_index: dict[str, list[str]]
    by_id: dict[str, Chunk]
    row_by_id: dict[str, int]


def load_index(index_dir) -> RetrievalIndex:
    index_dir = Path(index_dir)
    chunks = read_chunks(index_dir / "chunks.jsonl")
    faiss_index = load_faiss(index_dir / "embeddings.faiss")
    concept_index = read_json(index_dir / "concept_index.json")
    by_id = {c.id: c for c in chunks}
    row_by_id = {c.id: i for i, c in enumerate(chunks)}
    return RetrievalIndex(chunks, faiss_index, concept_index, by_id, row_by_id)