ssearch/build_store.py

# build_store.py
#
# Build or update the vector store from journal entries in ./data.
#
# Default mode (incremental): loads the existing index and adds only
# new or modified files.  Use --rebuild for a full rebuild from scratch.
#
# January 2026
# E. M. Furst
# Used Sonnet 4.5 to suggest changes; Opus 4.6 for incremental update

# Environment vars must be set before importing huggingface/transformers
# libraries, because huggingface_hub.constants evaluates HF_HUB_OFFLINE
# at import time. Matches the query scripts.
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_CACHE"] = "./models"
os.environ["HF_HUB_OFFLINE"] = "1"

from llama_index.core import (
    Document,
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
    Settings,
)
from pathlib import Path
from llama_index.core.readers.file.base import default_file_metadata_func
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.node_parser import SentenceSplitter
import argparse
import datetime
import re
import time

# Shared constants
DATA_DIR = Path("./data")
PERSIST_DIR = "./store"
EMBED_MODEL_NAME = "BAAI/bge-large-en-v1.5"
CHUNK_SIZE = 256
CHUNK_OVERLAP = 25

# Journal files are named YYYY-MM-DD.txt; the entry date is the filename.
# Filesystem timestamps are NOT the entry date (the 2000-07-29 entry carries
# creation_date 2001-05-25), so the date has to be parsed from the name.
ENTRY_DATE_RE = re.compile(r"^(\d{4})-(\d{2})-(\d{2})\.txt$")

# --- Metadata policy -------------------------------------------------------
#
# Every key below stays in the node's metadata dict and comes back on every
# retrieval result, whatever these lists say. The lists only control which
# keys get concatenated into the *text* that the embedder and the LLM see.
#
# Previously nothing was excluded from the embed text, so all 5210 chunks were
# embedded with a 36-token absolute path glued to the front -- 14% of a
# 256-token chunk, and the same constant string in every vector, which pulled
# the mean pairwise cosine up from 0.561 to 0.627 and made similarity cutoffs
# useless.
#
# Note BM25Retriever also indexes MetadataMode.EMBED, so the embed text is the
# only place it can see dates. That is why entry_date stays in: it keeps BM25
# able to match "2022-08-14" or "2022" at ~8 tokens instead of 36.
EXCLUDED_EMBED_KEYS = [
    "file_path", "file_name", "file_type", "file_size",
    "creation_date", "last_modified_date", "last_accessed_date",
    "year", "month",
]
# The LLM gets file_name instead -- it is the same date, and the synthesis
# prompt asks the model to cite file names.
EXCLUDED_LLM_KEYS = [
    "file_path", "file_type", "file_size",
    "creation_date", "last_modified_date", "last_accessed_date",
    "entry_date", "year", "month",
]


def journal_file_metadata(file_path: str) -> dict:
    """Default file metadata, plus the entry date parsed from the filename."""
    meta = default_file_metadata_func(file_path)
    m = ENTRY_DATE_RE.match(Path(file_path).name)
    if m:
        year, month, day = m.groups()
        meta["entry_date"] = f"{year}-{month}-{day}"
        meta["year"] = int(year)
        meta["month"] = int(month)
    return meta


def apply_metadata_policy(documents: list[Document]) -> list[Document]:
    """Set embed/LLM metadata exclusions. Nodes inherit these from the parent."""
    for doc in documents:
        doc.excluded_embed_metadata_keys = list(EXCLUDED_EMBED_KEYS)
        doc.excluded_llm_metadata_keys = list(EXCLUDED_LLM_KEYS)
    return documents


def load_documents(input_files=None) -> list[Document]:
    """Load journal documents with the shared metadata policy applied."""
    if input_files is not None:
        reader = SimpleDirectoryReader(
            input_files=input_files, file_metadata=journal_file_metadata
        )
    else:
        reader = SimpleDirectoryReader(
            str(DATA_DIR), file_metadata=journal_file_metadata
        )
    return apply_metadata_policy(reader.load_data())


_EMBED_TOKENIZER = None


def get_embed_tokenizer():
    """The embedding model's own tokenizer, loaded once.

    SentenceSplitter otherwise counts with tiktoken (cl100k_base) while the
    embedder counts in BERT WordPiece. They agree for ordinary prose (median
    ratio 1.01x) but diverge badly on unusual text -- up to 2.96x here. That
    let 94 chunks exceed CHUNK_SIZE and pushed 4 past bge-large's 512-token
    max_seq_length, where their tails were silently truncated at embed time and
    never made it into the vector. Counting in the embedder's own units makes
    chunk_size mean what it says.
    """
    global _EMBED_TOKENIZER
    if _EMBED_TOKENIZER is None:
        from transformers import AutoTokenizer
        _EMBED_TOKENIZER = AutoTokenizer.from_pretrained(
            EMBED_MODEL_NAME, cache_dir="./models", local_files_only=True
        )
    return _EMBED_TOKENIZER


def get_text_splitter():
    tok = get_embed_tokenizer()
    return SentenceSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        paragraph_separator="\n\n",
        tokenizer=lambda text: tok.encode(text, add_special_tokens=False),
    )


def rebuild():
    """Full rebuild: delete and recreate the vector store from scratch."""
    if not DATA_DIR.exists():
        raise FileNotFoundError(f"Data directory not found: {DATA_DIR.absolute()}")

    print(f"Loading documents from {DATA_DIR.absolute()}...")
    documents = load_documents()

    if not documents:
        raise ValueError("No documents found in data directory")

    print(f"Loaded {len(documents)} document(s)")

    print("Building vector index...")
    index = VectorStoreIndex.from_documents(
        documents,
        transformations=[get_text_splitter()],
        show_progress=True,
    )

    index.storage_context.persist(persist_dir=PERSIST_DIR)
    print(f"Index built and saved to {PERSIST_DIR}")


def update():
    """Incremental update: add new files, re-index modified files, remove deleted files."""
    if not DATA_DIR.exists():
        raise FileNotFoundError(f"Data directory not found: {DATA_DIR.absolute()}")

    # Set transformations BEFORE constructing the index.
    #
    # index.insert() chunks with self._transformations, which BaseIndex captures
    # once at construction time (transformations or Settings.transformations).
    # Setting Settings.transformations after load_index_from_storage() is too
    # late and is silently ignored -- which is what happened here until
    # 2026-09-20: every incrementally added file was chunked at the 1024-token
    # default instead of CHUNK_SIZE. Pass it explicitly as well, so this does
    # not depend on global state at all.
    Settings.transformations = [get_text_splitter()]

    # Load existing index
    print(f"Loading existing index from {PERSIST_DIR}...")
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(
        storage_context, transformations=[get_text_splitter()]
    )

    # Build lookup of indexed files: file_name -> (ref_doc_id, metadata)
    all_ref_docs = index.docstore.get_all_ref_doc_info()
    indexed = {}
    for ref_id, info in all_ref_docs.items():
        fname = info.metadata.get("file_name")
        if fname:
            indexed[fname] = (ref_id, info.metadata)

    print(f"Index contains {len(indexed)} documents")

    # Scan current files on disk.
    #
    # Skip hidden files. pathlib's glob("*.txt") matches dotfiles (unlike shell
    # globbing), but rebuild() goes through SimpleDirectoryReader, which skips
    # them. Without this filter the two modes disagree: a stray .2026-02-21.txt
    # was classified "new" after every rebuild and re-inserted via input_files=,
    # which bypasses the hidden-file filter.
    disk_files = {
        f.name: f for f in sorted(DATA_DIR.glob("*.txt"))
        if not f.name.startswith(".")
    }
    print(f"Data directory contains {len(disk_files)} files")

    # Classify files
    new_files = []
    modified_files = []
    deleted_files = []
    unchanged = 0

    for fname, fpath in disk_files.items():
        if fname not in indexed:
            new_files.append(fpath)
        else:
            ref_id, meta = indexed[fname]
            # Compare file size and modification date
            stat = fpath.stat()
            disk_size = stat.st_size
            # Must use UTC to match SimpleDirectoryReader's date format
            disk_mdate = datetime.datetime.fromtimestamp(
                stat.st_mtime, tz=datetime.timezone.utc
            ).strftime("%Y-%m-%d")

            stored_size = meta.get("file_size")
            stored_mdate = meta.get("last_modified_date")

            if disk_size != stored_size or disk_mdate != stored_mdate:
                modified_files.append((fpath, ref_id))
            else:
                unchanged += 1

    for fname, (ref_id, meta) in indexed.items():
        if fname not in disk_files:
            deleted_files.append((fname, ref_id))

    # Report
    print(f"\n  New:       {len(new_files)}")
    print(f"  Modified:  {len(modified_files)}")
    print(f"  Deleted:   {len(deleted_files)}")
    print(f"  Unchanged: {unchanged}")

    if not new_files and not modified_files and not deleted_files:
        print("\nNothing to do.")
        return

    # Process deletions (including modified files that need re-indexing)
    for fname, ref_id in deleted_files:
        print(f"  Removing {fname}")
        index.delete_ref_doc(ref_id, delete_from_docstore=True)

    for fpath, ref_id in modified_files:
        print(f"  Re-indexing {fpath.name} (modified)")
        index.delete_ref_doc(ref_id, delete_from_docstore=True)

    # Process additions (new files + modified files)
    files_to_add = new_files + [fpath for fpath, _ in modified_files]
    if files_to_add:
        print(f"\nIndexing {len(files_to_add)} file(s)...")
        # Use "./" prefix to match paths from full build (pathlib strips it)
        docs = load_documents(input_files=[f"./{f}" for f in files_to_add])
        for doc in docs:
            index.insert(doc)

    # Persist
    index.storage_context.persist(persist_dir=PERSIST_DIR)
    print(f"\nIndex updated and saved to {PERSIST_DIR}")


def main():
    parser = argparse.ArgumentParser(
        description="Build or update the vector store from journal entries."
    )
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Full rebuild from scratch (default: incremental update)",
    )
    args = parser.parse_args()

    # Configure embedding model
    # cache_folder + local_files_only match the query scripts, so the build
    # and the queries are guaranteed to use identical weights.
    embed_model = HuggingFaceEmbedding(
        model_name=EMBED_MODEL_NAME,
        cache_folder="./models",
        local_files_only=True,
    )
    Settings.embed_model = embed_model

    start = time.time()

    if args.rebuild:
        print("Mode: full rebuild")
        rebuild()
    else:
        print("Mode: incremental update")
        if not Path(PERSIST_DIR).exists():
            print(f"No existing index at {PERSIST_DIR}, doing full rebuild.")
            rebuild()
        else:
            update()

    elapsed = time.time() - start
    print(f"Done in {elapsed:.1f}s")


if __name__ == "__main__":
    main()