app/indexing.py

"""FAISS indexing utilities (mirrors the Colab indexing notebook)."""

from __future__ import annotations

import json
import re
from pathlib import Path

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

from app.paths import docs_root, index_path, metadata_path

CHUNK_SIZE = 120
CHUNK_OVERLAP = 30
MIN_CHUNK_WORDS = 30


def clean_markdown(text: str) -> str:
    text = text.replace("\\#", "")
    text = text.replace("\\-", "")
    text = text.replace("\\", "")
    text = re.sub(r"[ \t]+", " ", text)
    return text.strip()


def chunk_text_sliding_window(
    text: str,
    chunk_size: int = CHUNK_SIZE,
    overlap: int = CHUNK_OVERLAP,
) -> list[str]:
    words = text.split()
    chunks: list[str] = []
    start = 0

    while start < len(words):
        end = start + chunk_size
        chunk = words[start:end]
        chunks.append(" ".join(chunk))
        start += chunk_size - overlap

    return chunks


def build_chunks_from_document(source: str, text: str) -> list[dict]:
    clean_text = clean_markdown(text)
    title = Path(source).stem
    chunks: list[dict] = []

    for idx, chunk in enumerate(chunk_text_sliding_window(clean_text)):
        if len(chunk.split()) > MIN_CHUNK_WORDS:
            chunks.append(
                {
                    "id": f"{source}_{idx}",
                    "text": chunk,
                    "source": source,
                    "title": title,
                }
            )

    return chunks


def existing_sources(metadata: list[dict]) -> set[str]:
    return {entry["source"] for entry in metadata}


def list_indexed_documents(metadata: list[dict], total_vectors: int) -> dict:
    """Summarize unique indexed sources and chunk counts from metadata."""
    by_source: dict[str, dict] = {}

    for entry in metadata:
        source = entry["source"]
        if source not in by_source:
            category, _, _ = source.partition("/")
            by_source[source] = {
                "source": source,
                "title": entry.get("title", Path(source).stem),
                "category": category,
                "chunks": 0,
            }
        by_source[source]["chunks"] += 1

    documents = sorted(by_source.values(), key=lambda item: item["source"])
    return {
        "total_documents": len(documents),
        "total_chunks": len(metadata),
        "total_vectors": total_vectors,
        "documents": documents,
    }


def save_raw_document(category: str, filename: str, content: str) -> Path:
    path = docs_root() / category / filename
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")
    return path


def append_chunks_to_index(
    embedder: SentenceTransformer,
    index: faiss.Index,
    metadata: list[dict],
    new_chunks: list[dict],
    *,
    persist: bool = True,
) -> int:
    """Embed pre-built chunks and append them to the FAISS index and metadata."""
    texts = [chunk["text"] for chunk in new_chunks]
    embeddings = embedder.encode(
        texts,
        convert_to_numpy=True,
        normalize_embeddings=True,
    )
    index.add(np.asarray(embeddings, dtype=np.float32))
    metadata.extend(new_chunks)

    if persist:
        faiss.write_index(index, str(index_path()))
        with open(metadata_path(), "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=2)

    return len(new_chunks)


def append_document_to_index(
    embedder: SentenceTransformer,
    index: faiss.Index,
    metadata: list[dict],
    *,
    source: str,
    text: str,
    persist: bool = True,
) -> tuple[list[dict], int]:
    """
    Chunk and embed a new document, append vectors to the FAISS index and metadata.

    Raises ValueError if the source is already indexed or no chunks are produced.
    """
    if source in existing_sources(metadata):
        raise ValueError(f"Document already indexed: {source}")

    new_chunks = build_chunks_from_document(source, text)
    if not new_chunks:
        raise ValueError(
            "No indexable chunks produced (document empty or shorter than minimum chunk size)"
        )

    vectors_added = append_chunks_to_index(
        embedder, index, metadata, new_chunks, persist=persist
    )
    return new_chunks, vectors_added