Recursive Summarization - Hierarchical Document Compression

SUMMARIZE_PROMPT = """Summarize the following text concisely.
Preserve all key facts, figures, named entities, and conclusions.
Do not add information not present in the source.

Text:
{text}

Summary:"""

def summarize_chunk(text, llm):
    """Produce an information-dense summary of a single text block"""
    prompt = SUMMARIZE_PROMPT.format(text=text)
    return llm.invoke(prompt)

def recursive_summarize(chunks, llm, group_size=3):
    """
    Build the summary tree bottom-up.
    Returns a list of (level, node_id, parent_id, content) tuples.
    """
    nodes = []
    current_level = [
        {"level": 0, "id": f"L0_{i}", "parent_id": None, "content": c}
        for i, c in enumerate(chunks)
    ]
    nodes.extend(current_level)

    level = 1
    while len(current_level) > 1:
        next_level = []
        # Group current-level nodes and summarize each group
        for i in range(0, len(current_level), group_size):
            group = current_level[i : i + group_size]
            combined = "\n\n".join(n["content"] for n in group)
            summary = summarize_chunk(combined, llm)
            parent_id = f"L{level}_{i // group_size}"

            # Link children to this parent
            for child in group:
                child["parent_id"] = parent_id

            node = {
                "level": level,
                "id": parent_id,
                "parent_id": None,
                "content": summary,
                "children": [n["id"] for n in group],
            }
            next_level.append(node)
            nodes.append(node)

        current_level = next_level
        level += 1

    return nodes

def index_tree(nodes, vectorstore):
    """Embed and index every node in the summary tree"""
    for node in nodes:
        vectorstore.add({
            "id": node["id"],
            "content": node["content"],
            "level": node["level"],
            "metadata": {"level": node["level"], "node_id": node["id"]},
        })

def hierarchical_retrieve(query, vectorstore, top_k=5):
    """
    Search all tree levels and return the most relevant nodes.
    The vector index naturally surfaces the right level:
    abstract queries match high-level summaries,
    specific queries match leaf chunks.
    """
    results = vectorstore.search(query, top_k=top_k)
    # Results already span all levels — return as-is
    return results

def adaptive_level_retrieve(query, vectorstore,
                             detail_keywords=None, top_k=5):
    """
    Optionally bias toward a specific tree level based on query intent.
    'give me the gist' → prefer higher levels
    'find the specific figure' → prefer level 0
    """
    all_results = vectorstore.search(query, top_k=top_k * 2)

    if detail_keywords and any(kw in query.lower() for kw in detail_keywords):
        # Bias toward leaf nodes for detail-seeking queries
        filtered = [r for r in all_results if r["level"] == 0]
        return filtered[:top_k] or all_results[:top_k]

    return all_results[:top_k]

DETAIL_KEYWORDS = ["specific", "exact", "figure", "number", "quote", "line"]

def recursive_summarization_rag(query, vectorstore, llm,
                                 documents=None, top_k=5):
    # Build tree once at index time (skip if already indexed)
    if documents:
        chunks = chunk_documents(documents, chunk_size=512)
        nodes = recursive_summarize(chunks, llm, group_size=3)
        index_tree(nodes, vectorstore)

    # Retrieve from tree (all levels searchable)
    results = adaptive_level_retrieve(
        query, vectorstore, DETAIL_KEYWORDS, top_k=top_k
    )

    # Generate
    context = "\n\n".join([r["content"] for r in results])
    response = llm.generate(
        f"Context: {context}\n\nQuestion: {query}"
    )

    return response, {
        "retrieved_levels": [r["level"] for r in results],
        "num_results": len(results),
    }