Sentence Window Retrieval - Precise Match, Rich Context

from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.core import VectorStoreIndex

# --- Index time ---
# Parse documents into sentence nodes; surrounding window stored in metadata
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=2,                         # N=2: 5-sentence window
    window_metadata_key="window",
    original_text_metadata_key="original_text",
)

nodes = node_parser.get_nodes_from_documents(documents)
index = VectorStoreIndex(nodes)

# --- Query time ---
# MetadataReplacementPostProcessor swaps the matched sentence → full window
postprocessor = MetadataReplacementPostProcessor(
    target_metadata_key="window"
)

query_engine = index.as_query_engine(
    similarity_top_k=5,
    node_postprocessors=[postprocessor],
)

response = query_engine.query("What triggers apoptosis in T-cells?")

import re
from typing import List, Dict
import numpy as np


def split_into_sentences(text: str) -> List[str]:
    """Split text into individual sentences."""
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s for s in sentences if len(s) > 10]


def build_sentence_index(
    documents: List[Dict],
    embed_fn,
    window_size: int = 2,
):
    """
    Index each sentence independently, storing position metadata.
    Returns a list of sentence records and their embeddings.
    """
    records = []
    for doc in documents:
        sentences = split_into_sentences(doc["content"])
        for i, sentence in enumerate(sentences):
            records.append({
                "doc_id": doc["id"],
                "sentence_index": i,
                "sentence": sentence,
                "all_sentences": sentences,  # stored for window expansion
                "source": doc["source"],
            })

    texts = [r["sentence"] for r in records]
    embeddings = embed_fn(texts)            # shape: (N, dim)
    return records, np.array(embeddings)


def expand_to_window(record: Dict, window_size: int) -> str:
    """Expand the matched sentence to its surrounding N-sentence window."""
    sentences = record["all_sentences"]
    i = record["sentence_index"]
    start = max(0, i - window_size)
    end = min(len(sentences), i + window_size + 1)
    return " ".join(sentences[start:end])


def sentence_window_retrieve(
    query: str,
    records: List[Dict],
    embeddings: np.ndarray,
    embed_fn,
    window_size: int = 2,
    top_k: int = 5,
) -> List[Dict]:
    """Retrieve top-k sentences and expand each to its context window."""
    query_vec = embed_fn([query])[0]
    scores = embeddings @ query_vec / (
        np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_vec) + 1e-9
    )
    top_indices = np.argsort(scores)[::-1][:top_k]

    results = []
    for idx in top_indices:
        record = records[idx]
        window_text = expand_to_window(record, window_size)
        results.append({
            "matched_sentence": record["sentence"],
            "window_context": window_text,
            "score": float(scores[idx]),
            "source": record["source"],
        })
    return results


def sentence_window_rag(
    query: str,
    records: List[Dict],
    embeddings: np.ndarray,
    embed_fn,
    llm,
    window_size: int = 2,
) -> Dict:
    # Retrieve with automatic window expansion
    retrieved = sentence_window_retrieve(
        query, records, embeddings, embed_fn, window_size
    )

    # Build generation context from expanded windows (not raw sentences)
    context = "\n\n".join([r["window_context"] for r in retrieved])

    response = llm.invoke(f"Context:\n{context}\n\nQuestion: {query}")
    return response, retrieved