Advanced RAG: Sentence-Window Retrieval for Precise LLM Context

# requirements.txt
# unstructured[pdf]
# nltk
# sentence-transformers
# pinecone-client
# tqdm
# python-dotenv

import os
import re
import nltk
import pinecone
from unstructured.partition.pdf import partition_pdf
from sentence_transformers import SentenceTransformer
from tqdm.auto import tqdm
from dotenv import load_dotenv

# --- Configuration and Initialization ---
load_dotenv()
nltk.download('punkt', quiet=True)

PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT")
INDEX_NAME = "sentence-window-index"

# Use a high-quality sentence-level model
EMBEDDING_MODEL = SentenceTransformer('all-MiniLM-L6-v2')
EMBEDDING_DIMENSION = EMBEDDING_MODEL.get_sentence_embedding_dimension()

def initialize_pinecone():
    """Initializes and returns the Pinecone index."""
    pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
    if INDEX_NAME not in pinecone.list_indexes():
        pinecone.create_index(
            name=INDEX_NAME,
            dimension=EMBEDDING_DIMENSION,
            metric='cosine' # Cosine similarity is standard for sentence transformers
        )
    return pinecone.Index(INDEX_NAME)

# --- Document Processing and Sentence Splitting ---

def clean_text(text):
    """A simple text cleaner to remove excessive whitespace and artifacts."""
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def process_document(file_path, doc_id):
    """Processes a PDF, splits it into sentences, and prepares for indexing."""
    print(f"Processing document: {doc_id}")
    elements = partition_pdf(filename=file_path)
    full_text = "\n\n".join([e.text for e in elements])
    
    # NLTK is generally more robust for sentence splitting than simple regex
    sentences = nltk.sent_tokenize(full_text)
    
    # Clean and filter out short/empty sentences
    cleaned_sentences = [clean_text(s) for s in sentences if len(s.split()) > 3]
    
    print(f"  - Extracted {len(cleaned_sentences)} sentences.")
    return cleaned_sentences

# --- Indexing Logic ---

def index_sentences(index, sentences, doc_id, batch_size=100):
    """Embeds sentences and upserts them into Pinecone with metadata."""
    print(f"Embedding and indexing sentences for {doc_id}...")
    for i in tqdm(range(0, len(sentences), batch_size)):
        batch_sentences = sentences[i:i+batch_size]
        batch_indices = range(i, i + len(batch_sentences))
        
        # Create embeddings
        embeddings = EMBEDDING_MODEL.encode(batch_sentences).tolist()
        
        # Prepare vectors for upsert
        vectors_to_upsert = []
        for j, (sentence, embedding) in enumerate(zip(batch_sentences, embeddings)):
            sentence_index = batch_indices[j]
            vector_id = f"{doc_id}-sent{sentence_index}"
            metadata = {
                'document_id': doc_id,
                'sentence_index': sentence_index,
                'text': sentence
            }
            vectors_to_upsert.append((vector_id, embedding, metadata))
            
        # Upsert the batch
        index.upsert(vectors=vectors_to_upsert)
    print(f"Finished indexing for {doc_id}.")

# --- Main Execution ---
if __name__ == '__main__':
    # This assumes you have a PDF file named 'attention-is-all-you-need.pdf'
    # in a 'data' directory.
    doc_path = "data/attention-is-all-you-need.pdf"
    doc_id = "attention-paper-v1"
    
    # 1. Initialize Pinecone Index
    pinecone_index = initialize_pinecone()
    
    # 2. Process the document
    all_sentences = process_document(doc_path, doc_id)
    
    # For the retrieval step, we need the full list of sentences easily accessible.
    # In a production system, this would be stored in a more robust cache 
    # like Redis or a document database (e.g., MongoDB, DynamoDB).
    # For this example, we'll just save it to a simple dictionary.
    document_sentence_store = {doc_id: all_sentences}
    
    # 3. Index the sentences
    index_sentences(pinecone_index, all_sentences, doc_id)
    
    # You can now query this index. The `document_sentence_store` is critical for the next step.
    print("\nIndexing complete. The system is ready for retrieval.")

# This code builds upon the previous section.
# Assume `pinecone_index` and `document_sentence_store` are already populated.

# --- Retrieval and Context Expansion ---

def retrieve_and_expand_context(query, index, doc_store, window_size=2, top_k=5):
    """
    Retrieves top_k sentences and expands their context using a window.
    Handles overlapping windows by merging them.
    
    Args:
        query (str): The user's query.
        index (pinecone.Index): The initialized Pinecone index.
        doc_store (dict): A dictionary mapping doc_id to a list of its sentences.
        window_size (int): Number of sentences to include before and after the retrieved sentence.
        top_k (int): The number of top sentences to retrieve.
        
    Returns:
        str: A single string containing the merged, expanded context.
    """
    print(f"Executing query: '{query}'")
    
    # 1. Embed the query
    query_embedding = EMBEDDING_MODEL.encode(query).tolist()
    
    # 2. Retrieve top_k relevant sentences
    retrieval_results = index.query(
        vector=query_embedding,
        top_k=top_k,
        include_metadata=True
    )
    
    # 3. Extract sentence indices and document IDs
    retrieved_indices = []
    for match in retrieval_results['matches']:
        metadata = match['metadata']
        retrieved_indices.append(
            (metadata['document_id'], metadata['sentence_index'])
        )
    
    print(f"  - Retrieved sentence indices: {retrieved_indices}")
    
    # 4. Expand and Merge Windows
    # We need to handle windows that might overlap.
    # A simple way is to collect all sentence indices we need to fetch,
    # put them in a set to remove duplicates, and then sort them.
    
    final_indices_to_fetch = set()
    for doc_id, sent_idx in retrieved_indices:
        start_index = max(0, sent_idx - window_size)
        end_index = sent_idx + window_size + 1 # +1 because slice is exclusive
        
        # Add all indices in the window to the set
        for i in range(start_index, end_index):
            final_indices_to_fetch.add((doc_id, i))
            
    # Sort the indices to reconstruct the text in the correct order
    sorted_indices = sorted(list(final_indices_to_fetch), key=lambda x: (x[0], x[1]))
    
    # 5. Reconstruct the final context
    context_parts = []
    current_doc_id = None
    for doc_id, sent_idx in sorted_indices:
        if doc_id != current_doc_id:
            # Add a separator for context from different documents if necessary
            if current_doc_id is not None:
                context_parts.append("\n---\n")
            current_doc_id = doc_id
        
        # Fetch the sentence from our document store
        sentences_for_doc = doc_store.get(doc_id, [])
        if sent_idx < len(sentences_for_doc):
            context_parts.append(sentences_for_doc[sent_idx])

    final_context = " ".join(context_parts)
    print("  - Final constructed context sent to LLM:")
    print(final_context)
    return final_context

# --- Example Usage (Continuing from Part 1) ---
if __name__ == '__main__':
    # This is a continuation. Ensure the indexing script has been run.
    pinecone_index = initialize_pinecone()
    
    # In a real app, you'd load this from your persistent store (Redis, etc.)
    doc_id = "attention-paper-v1"
    # Re-create the sentence store for this example run
    doc_path = "data/attention-is-all-you-need.pdf"
    all_sentences = process_document(doc_path, doc_id)
    document_sentence_store = {doc_id: all_sentences}
    
    
    # --- Query Examples ---
    print("\n--- Query 1: What is the Transformer architecture? ---")
    query1 = "What is the core architecture of the Transformer model?"
    context1 = retrieve_and_expand_context(query1, pinecone_index, document_sentence_store, window_size=2, top_k=3)
    
    print("\n--- Query 2: How does self-attention work? ---")
    query2 = "Explain the mechanism of self-attention."
    context2 = retrieve_and_expand_context(query2, pinecone_index, document_sentence_store, window_size=3, top_k=5)

# Conceptual code for re-ranking
from sentence_transformers.cross_encoder import CrossEncoder

# ... after retrieving and expanding windows into a list called `expanded_windows`

cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

# Create pairs of [query, window_text] for scoring
query_window_pairs = [[query, window['text']] for window in expanded_windows]

# Score all pairs
scores = cross_encoder.predict(query_window_pairs)

# Combine scores with the windows and sort
for i in range(len(expanded_windows)):
    expanded_windows[i]['rerank_score'] = scores[i]

sorted_windows = sorted(expanded_windows, key=lambda x: x['rerank_score'], reverse=True)

# Select the top N (e.g., 3) re-ranked windows to pass to the LLM
final_top_windows = sorted_windows[:3]
final_context = "\n---\n".join([window['text'] for window in final_top_windows])