src/rag/vector_store.py

"""
vector_store.py
-----------------
A small, real, in-memory vector store over the ResearchPilot corpus.
No external vector DB service is required, which keeps the project
deployable on a free Hugging Face Space with no paid infrastructure.

Documents are chunked (a paper's abstract text is split into overlapping
windows) so retrieval returns passage-level evidence rather than whole
documents, matching the "evidence passages" requirement.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import List

import numpy as np

from src.rag.embeddings import get_embedding_backend


@dataclass
class Chunk:
    chunk_id: str
    doc_id: str
    title: str
    authors: str
    year: int
    url: str
    topic: str
    text: str
    embedding: np.ndarray = field(default=None, repr=False)


def _split_into_chunks(text: str, max_words: int = 60, overlap: int = 15) -> List[str]:
    words = text.split()
    if len(words) <= max_words:
        return [text]
    chunks = []
    start = 0
    while start < len(words):
        end = min(start + max_words, len(words))
        chunks.append(" ".join(words[start:end]))
        if end == len(words):
            break
        start = end - overlap
    return chunks


class VectorStore:
    def __init__(self, corpus_path: str):
        self.corpus_path = corpus_path
        self.chunks: List[Chunk] = []
        self.backend = get_embedding_backend()
        self._build()

    def _build(self) -> None:
        with open(self.corpus_path, "r", encoding="utf-8") as f:
            docs = json.load(f)

        raw_chunks: List[Chunk] = []
        for doc in docs:
            pieces = _split_into_chunks(doc["text"])
            for i, piece in enumerate(pieces):
                raw_chunks.append(
                    Chunk(
                        chunk_id=f"{doc['id']}::{i}",
                        doc_id=doc["id"],
                        title=doc["title"],
                        authors=doc.get("authors", ""),
                        year=doc.get("year", 0),
                        url=doc.get("url", ""),
                        topic=doc.get("topic", ""),
                        text=piece,
                    )
                )

        texts = [c.text for c in raw_chunks]
        self.backend.fit_corpus(texts)
        embeddings = self.backend.encode(texts)
        for chunk, emb in zip(raw_chunks, embeddings):
            chunk.embedding = emb
        self.chunks = raw_chunks

    def num_documents(self) -> int:
        return len({c.doc_id for c in self.chunks})

    def num_chunks(self) -> int:
        return len(self.chunks)

    def embedding_backend_name(self) -> str:
        return self.backend.backend_name

    def all_chunks(self) -> List[Chunk]:
        return self.chunks