Iterative Retrieval - Progressive Query Refinement

from typing import List, Dict, Tuple

def sufficiency_check(query: str, context_chunks: List[Dict], llm) -> Tuple[bool, str]:
    """
    Ask the LLM whether the current context is sufficient to answer the query.
    Returns (is_sufficient, gap_description).
    """
    context_text = "\n\n".join([c['content'] for c in context_chunks])

    prompt = f"""You are evaluating whether a set of retrieved documents provides 
enough information to fully answer a question.

Original question: {query}

Retrieved context so far:
{context_text}

Assessment:
1. Can this question be fully answered from the context above? (yes/no)
2. If no, what specific information is still missing?

Respond as JSON: {{"sufficient": true/false, "gap": "description of missing info or empty string"}}"""

    result = llm.invoke(prompt, schema={"sufficient": bool, "gap": str})
    return result['sufficient'], result['gap']


def refine_query(original_query: str, gap: str, context_chunks: List[Dict], llm) -> str:
    """
    Generate a more specific follow-up query targeting the identified gap.
    Uses retrieved chunks as context so the new query builds on what we know.
    """
    context_summary = "\n".join([f"- {c['content'][:200]}" for c in context_chunks[:3]])

    prompt = f"""You are refining a search query to fill a specific knowledge gap.

Original question: {original_query}

What we already know (retrieved context summary):
{context_summary}

What is still missing: {gap}

Write a concise, targeted search query to retrieve documents that fill this gap.
Return only the query, nothing else."""

    return llm.invoke(prompt).strip()


def iterative_rag(
    query: str,
    retriever,
    llm,
    top_k: int = 5,
    max_iterations: int = 4,
    sufficiency_threshold: bool = True,
) -> Dict:
    """
    Iterative RAG pipeline.
    Runs retrieve → assess → refine → retrieve until context is sufficient
    or max_iterations is reached.
    """
    accumulated_context: List[Dict] = []
    seen_doc_ids = set()
    query_history = [query]
    current_query = query

    for iteration in range(max_iterations):
        # Step 1: Retrieve with the current query
        results = retriever.search(current_query, top_k=top_k)

        # Deduplicate — don't re-add chunks we already have
        new_chunks = [r for r in results if r['id'] not in seen_doc_ids]
        for chunk in new_chunks:
            seen_doc_ids.add(chunk['id'])
        accumulated_context.extend(new_chunks)

        # Step 2: Check if accumulated context is now sufficient
        is_sufficient, gap = sufficiency_check(query, accumulated_context, llm)

        if is_sufficient:
            break

        # Step 3: If not sufficient and iterations remain, refine the query
        if iteration < max_iterations - 1:
            current_query = refine_query(query, gap, accumulated_context, llm)
            query_history.append(current_query)

    # Step 4: Generate final answer from accumulated context
    context_text = "\n\n".join([c['content'] for c in accumulated_context])
    answer = llm.invoke(
        f"Context:\n{context_text}\n\nQuestion: {query}\n\nAnswer:"
    )

    return {
        'answer': answer,
        'iterations': len(query_history),
        'query_history': query_history,
        'context_chunks_used': len(accumulated_context),
        'sufficient': is_sufficient,
    }