Fix Lost in the Middle Syndrome in RAG Retrieval

# position_test.py
import os
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage

def test_position_sensitivity(context_chunks: list[str], query: str) -> dict:
    """
    Place the answer in first, middle, and last position.
    If accuracy drops in the middle, you have the syndrome.
    """
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    results = {}

    for position in ["first", "middle", "last"]:
        chunks = context_chunks.copy()
        # Move the known-answer chunk to test position
        answer_chunk = chunks.pop(0)

        if position == "first":
            ordered = [answer_chunk] + chunks
        elif position == "last":
            ordered = chunks + [answer_chunk]
        else:
            mid = len(chunks) // 2
            ordered = chunks[:mid] + [answer_chunk] + chunks[mid:]

        context = "\n\n---\n\n".join(ordered)
        prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer concisely."
        response = llm.invoke([HumanMessage(content=prompt)])
        results[position] = response.content

    return results

# reranker.py
from sentence_transformers import CrossEncoder
from dataclasses import dataclass

@dataclass
class RankedChunk:
    text: str
    original_rank: int
    rerank_score: float

def rerank_chunks(
    query: str,
    chunks: list[str],
    model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
    top_k: int = 5
) -> list[RankedChunk]:
    """
    Score each chunk against the query with a cross-encoder.
    Returns top_k chunks sorted by relevance score descending.
    """
    encoder = CrossEncoder(model_name)

    # Cross-encoder scores query+chunk pairs directly
    pairs = [(query, chunk) for chunk in chunks]
    scores = encoder.predict(pairs)

    ranked = [
        RankedChunk(text=chunk, original_rank=i, rerank_score=float(score))
        for i, (chunk, score) in enumerate(zip(chunks, scores))
    ]

    # Sort by score, take top_k
    ranked.sort(key=lambda x: x.rerank_score, reverse=True)
    return ranked[:top_k]

pip install sentence-transformers

# rag_pipeline.py
def ushape_order(ranked_chunks: list[RankedChunk]) -> list[str]:
    """
    Place highest-scored chunks at the edges of the context.
    Lowest-scored chunks land in the middle.

    Example with 5 chunks ranked [A, B, C, D, E] by score:
    Result order: [A, C, E, D, B]
    Positions:    [0, 2, 4, 3, 1]  ← A and B at edges
    """
    texts = [c.text for c in ranked_chunks]
    result = []
    left, right = 0, len(texts) - 1

    # Alternate placing from highest-scored at left/right edges
    toggle = True
    for chunk in texts:
        if toggle:
            result.insert(0, chunk)   # Push to front
        else:
            result.append(chunk)      # Push to back
        toggle = not toggle

    return result


def build_rag_prompt(query: str, chunks: list[str]) -> str:
    reranked = rerank_chunks(query, chunks, top_k=5)
    ordered = ushape_order(reranked)
    context = "\n\n---\n\n".join(ordered)

    return (
        f"Use only the provided context to answer the question.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {query}\n\n"
        f"Answer:"
    )

# full_pipeline.py
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.schema import HumanMessage

def rag_query(query: str, vectorstore: Chroma, llm: ChatOpenAI) -> str:
    # Retrieve more than you need — reranker will trim
    raw_docs = vectorstore.similarity_search(query, k=10)
    raw_chunks = [doc.page_content for doc in raw_docs]

    # Rerank and reorder
    prompt = build_rag_prompt(query, raw_chunks)

    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content


# Usage
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(embedding_function=embeddings, persist_directory="./chroma_db")
llm = ChatOpenAI(model="gpt-4o", temperature=0)

answer = rag_query("What were Q3 revenue figures?", vectorstore, llm)
print(answer)

python position_test.py

# Before fix: measure how often the correct answer came from a middle-position chunk
# After fix: that number should drop — correct chunks are now at edges

def score_position_hits(results: list[dict]) -> float:
    """Returns fraction of correct answers NOT from middle-position chunks."""
    edge_hits = sum(1 for r in results if r["correct_position"] in ("first", "last"))
    return edge_hits / len(results)