DRAG con KNEE para RAG: poda dinámica

DRAG con KNEE: Segmentación Dinámica para RAG de Alto Rendimiento

DRAG con KNEE es un algoritmo avanzado de generación aumentada por recuperación que utiliza la eliminación de puntos de codo para optimizar la selección de contexto. Construye árboles vectoriales jerárquicos en Qdrant y aprovecha una búsqueda híbrida con Fusión de Clasificación Recíproca (RRF) para recuperar de forma adaptativa el contexto relevante, superando las limitaciones de los métodos estáticos tipo top_k, como el desbordamiento de contexto o la pérdida de información.

Creación de Índices Jerárquicos

Los documentos se dividen en páginas. Para cada página, un modelo de lenguaje (LLM) genera un resumen conciso y términos clave mediante salida estructurada; estos se convierten en nodos hoja del árbol.

Las hojas se agrupan en nodos padres (3–5 por grupo). Cada nodo padre hereda los términos de sus hijos y genera su propio resumen al combinar el contenido de los nodos descendientes. Este proceso se repite hasta alcanzar el nodo raíz (parent_id = -1).

En Qdrant, los datos se almacenan como una lista plana de puntos con metadatos para filtrar. Esto permite manipular flexiblemente la estructura jerárquica a través de campos de carga útil (payload).

Búsqueda Híbrida y RRF

La comparación de puntos combina embeddings densos y dispersos usando Fusión de Clasificación Recíproca:

$$RRFscore(d \in D) = \sum_{r \in R} \frac{1}{k + r(d)}$$

donde $D$ es el conjunto de documentos, $R$ representa los rankings (densos, dispersos), y $k=60$.

Algoritmos de Búsqueda

Búsqueda por Ramas

Encuentra las raíces de nivel superior, luego compara recursivamente los padres con sus hijos. Si los hijos obtienen mejor puntuación, se profundiza más; si no, se mantiene el nodo padre.

def parent_vs_children(query: str, parent: ScoredPoint) -> list[ScoredPoint]:
    file_name = parent.payload["file_name"]
    parent_id = parent.payload["id"]
    child_ids = parent.payload["child_ids"]
    all_ids = [parent_id] + child_ids

    if len(child_ids) == 0:
        console.print(
            f"Parent with id {parent_id} is leave!", style="bold underline violet"
        )
        return []

    console.print()
    console.print("#" * 20, style="red")
    console.print(
        f"Children (ids: {child_ids}) [underline]VS[/underline] Parent (id: {parent_id}) [ FILE: {file_name} ]\nFor query: {query[:20]}...",
        style="bold violet",
    )

    with RAGalicClient() as client: # artifacts from a previous life :)
        dense_model: str = client.client.embedding_model_name  # type: ignore
        sparse_model: str = client.client.sparse_embedding_model_name  # type: ignore
        children_and_parent_filter = Filter(
            must=FieldCondition(key="id", match=MatchAny(any=all_ids))
        )
        sorted_points = client.client.query_points(
            collection_name="ragalic",
            prefetch=[
                Prefetch(
                    query=Document(text=query, model=dense_model),
                    using="dense",
                    limit=len(all_ids),
                    filter=children_and_parent_filter,
                ),
                Prefetch(
                    query=Document(text=query, model=sparse_model),
                    using="sparse",
                    limit=len(all_ids),
                    filter=children_and_parent_filter,
                ),
            ],
            query=FusionQuery(fusion=Fusion.RRF),
            query_filter=children_and_parent_filter,
            limit=len(all_ids),
        ).points

    children_better_then_parent = []
    for point in sorted_points:
        if point.payload["id"] == parent_id:
            break
        children_better_then_parent.append(point)

    ids = [child.payload["id"] for child in children_better_then_parent]
    console.print(
        f"Child ids which is better than parent: {ids if ids else 'N/A'}",
        style="italic purple",
    )
    console.print("#" * 20, style="red")
    console.print()

    return children_better_then_parent


def branch_search(query: str, num_roots: int = 3) -> list[Chunk]:
    roots = find_roots(query=query, num_to_find=num_roots)

    final_points = []
    points_to_process = roots
    while len(points_to_process) > 0:
        new_point_to_process = []
        for point in points_to_process:
            new_points = parent_vs_children(query=query, parent=point)
            if len(new_points) == 0:
                console.print(
                    f"--- NEW FINAL POINT (id: {point.payload['id']} | file: {point.payload['file_name']} | pages: {point.payload['page_start']} - {point.payload['page_end']}) ---",
                    style="bold green",
                )
                final_points.append(point)
            else:
                new_point_to_process.extend(new_points)
        points_to_process = new_point_to_process

    console.print(
        f"--- WAS FOUND {len(final_points)} POINTS FOR QUERY: '{query[:20]}...' ---",
        style="bold green on white",
    )

    return prepare_chunks(final_points)

Uso:

from rag_lib.search import branch_search

results = branch_search(query="your_query", num_roots=3)
for chunk in results:
    print(f"{chunk.file_name} (pages {chunk.page_start}-{chunk.page_end}):
{chunk.text[:200]}...")

Ventajas: Alta cobertura. Desventajas: Puede incluir ramas irrelevantes.

Búsqueda por Haz (Fijo)

• Se toma la generación actual de padres y sus hijos.
• Los ganadores entre padres eliminan a sus hijos; los ganadores entre hijos eliminan a sus padres.
• Los supervivientes compiten por un ancho de haz fijo.

Reduce el ruido, pero puede cortar contenido relevante o incluir el irrelevante.

Haz Adaptativo con Punto de Codo

Anchura dinámica usando el método del punto de codo. Identifica el punto de inflexión en la curva RRF (convexidad descendente).

Eliminación dinámica de ruido: en cada iteración, se aplica primero la poda por punto de codo, luego se genera el siguiente haz. Descarta aproximadamente el 62% de los documentos irrelevantes.

Conclusiones Clave

• Árboles jerárquicos en Qdrant con términos heredados.
• Búsqueda híbrida + RRF para comparación eficaz de nodos.
• Poda por punto de codo para segmentación dinámica basada en relevancia.
• Búsqueda por ramas para completitud, haz adaptativo para eficiencia.
• Salida estructurada del LLM minimiza falsas creaciones (hallucinations).

— Editorial Team