pilotcore/retrieval/runtime.py

from pilotcore.retrieval.retriever import (
    retrieve_chunks,
)

from pilotcore.retrieval.vector_store import (
    search_vectors,
)

from pilotcore.tracing.spans import (
    start_span,
    end_span,
)
from pilotcore.retrieval.reranker import rerank_chunks

from pilotcore.retrieval.multi_query import (
    generate_queries,
)


def deduplicate_chunks(chunks):

    seen = set()
    unique = []

    for chunk in chunks:

        text = chunk.chunk.text.strip()

        # crude but effective near-duplicate filter
        key = text[:250].lower()

        if key in seen:
            continue

        seen.add(key)
        unique.append(chunk)

    return unique


def run_dedup(chunks):
    before = len(chunks)

    chunks = deduplicate_chunks(chunks)

    after = len(chunks)

    print(f"[DEDUP] {before} -> {after}")

    return chunks


def apply_post_processing(
    chunks,
    query,
    top_k,
    experiment_config,
):
    if experiment_config is None:
        return chunks

    if experiment_config.deduplication:
        chunks = run_dedup(chunks)

    if experiment_config.reranker:

        chunks = rerank_chunks(
            query=query,
            candidate_chunks=chunks,
            top_k=top_k,
            model_key=experiment_config.reranker_model,
        )

    return chunks


def retrieve(
    strategy: str,
    **kwargs,
):
    experiment_config = kwargs.pop(
        "experiment_config",
        None,
    )

    trace = kwargs.get("trace")

    kwargs.pop("trace", None)

    if strategy == "lexical":

        span = start_span(
            trace_id=trace.trace_id,
            name="retrieval",
        )

        trace.spans.append(span)
        query = kwargs.get("query")
        top_k = kwargs.get("top_k", 7)

        result = retrieve_chunks(**kwargs)

        result.retrieved_chunks = apply_post_processing(
            chunks=result.retrieved_chunks,
            query=query,
            top_k=top_k,
            experiment_config=experiment_config,
        )

        end_span(span)

        return result

    elif strategy == "vector":

        span = start_span(
            trace_id=trace.trace_id,
            name="vector_retrieval",
        )

        trace.spans.append(span)

        from pilotcore.retrieval.embeddings import get_embedding

        query = kwargs.pop("query")
        user_id = kwargs.pop("user_id", None)
        source = kwargs.pop("source", None)
        trace_id = kwargs.pop("trace_id")
        top_k = kwargs.pop("top_k", 7)
        query_embedding = get_embedding(query)

        result = search_vectors(
            user_id=user_id,
            query_embedding=query_embedding,
            source=source,
            trace_id=trace_id,
            top_k=top_k,
        )
        result.retrieved_chunks = apply_post_processing(
            chunks=result.retrieved_chunks,
            query=query,
            top_k=top_k,
            experiment_config=experiment_config,
        )
        end_span(span)

        return result

    elif strategy == "hybrid":

        span = start_span(
            trace_id=trace.trace_id,
            name="hybrid_retrieval",
        )

        trace.spans.append(span)

        from pilotcore.retrieval.embeddings import get_embedding

        from pilotcore.schemas.retrieval import RetrievalResult

        query = kwargs.pop("query")
        user_id = kwargs.pop("user_id", None)
        source = kwargs.pop("source", None)
        trace_id = kwargs.pop("trace_id")
        top_k = kwargs.pop("top_k", 7)

        query_variants = [query]

        if experiment_config and experiment_config.multi_query:

            query_variants = generate_queries(query)
            print("\n===== MULTI QUERY =====")

            for i, q in enumerate(query_variants, start=1):
                print(f"{i}. {q}")

            print("=======================\n")
            if trace:
                trace.generated_queries = query_variants

        all_vector_chunks = []
        all_lexical_chunks = []

        for query_variant in query_variants:

            query_embedding = get_embedding(query_variant)

            vector_result = search_vectors(
                user_id=user_id,
                query_embedding=query_embedding,
                source=source,
                trace_id=trace_id,
                top_k=top_k,
            )

            lexical_result = retrieve_chunks(
                user_id=user_id,
                query=query_variant,
                source=source,
                trace_id=trace_id,
                top_k=top_k,
            )

            all_vector_chunks.extend(vector_result.retrieved_chunks)

            all_lexical_chunks.extend(lexical_result.retrieved_chunks)

        vector_result.retrieved_chunks = deduplicate_chunks(all_vector_chunks)

        lexical_result.retrieved_chunks = deduplicate_chunks(all_lexical_chunks)

        # Reciprocal Rank Fusion (RRF)
        # -----------------------------------------
        # Instead of naïvely concatenating vector and BM25 results,
        # we fuse rankings from both retrievers.
        #
        # Why RRF?
        # - robust across retrievers with different score scales
        # - boosts chunks retrieved by BOTH systems
        # - improves hybrid retrieval quality significantly
        #
        # Formula:
        # score += 1 / (RRF_K + rank)

        from pilotcore.schemas.retrieval import RetrievedChunk

        RRF_K = 60

        rrf_scores = {}
        chunk_map = {}

        # Vector retrieval ranks
        for rank, chunk in enumerate(vector_result.retrieved_chunks, start=1):

            chunk_key = (
                chunk.chunk.document_id,
                chunk.chunk.chunk_id,
            )

            if chunk_key not in rrf_scores:

                rrf_scores[chunk_key] = 0.0

                chunk_map[chunk_key] = chunk

            else:

                existing = chunk_map[chunk_key]

                # Preserve dense lineage if newly available
                if chunk.dense_score is not None:
                    existing.dense_score = chunk.dense_score

                if chunk.dense_rank is not None:
                    existing.dense_rank = chunk.dense_rank

                # Preserve BM25 lineage if newly available
                if chunk.bm25_score is not None:
                    existing.bm25_score = chunk.bm25_score

                if chunk.bm25_rank is not None:
                    existing.bm25_rank = chunk.bm25_rank

                # Merge provenance safely
                existing.retrieval_sources = list(
                    set(existing.retrieval_sources + chunk.retrieval_sources)
                )

            rrf_scores[chunk_key] += 1.0 / (RRF_K + rank)

        # BM25 retrieval ranks
        for rank, chunk in enumerate(lexical_result.retrieved_chunks, start=1):

            chunk_key = (
                chunk.chunk.document_id,
                chunk.chunk.chunk_id,
            )

            if chunk_key not in rrf_scores:

                rrf_scores[chunk_key] = 0.0

                chunk_map[chunk_key] = chunk

            else:

                existing = chunk_map[chunk_key]

                # Preserve dense lineage if newly available
                if chunk.dense_score is not None:
                    existing.dense_score = chunk.dense_score

                if chunk.dense_rank is not None:
                    existing.dense_rank = chunk.dense_rank

                # Preserve BM25 lineage if newly available
                if chunk.bm25_score is not None:
                    existing.bm25_score = chunk.bm25_score

                if chunk.bm25_rank is not None:
                    existing.bm25_rank = chunk.bm25_rank

                # Merge provenance safely
                existing.retrieval_sources = list(
                    set(existing.retrieval_sources + chunk.retrieval_sources)
                )

            rrf_scores[chunk_key] += 1.0 / (RRF_K + rank)

        # Build fused chunk list
        fused_chunks = []

        for chunk_key, fused_score in rrf_scores.items():

            original_chunk = chunk_map[chunk_key]

            fused_chunks.append(
                RetrievedChunk(
                    chunk=original_chunk.chunk,
                    # temporary compatibility
                    score=float(fused_score),
                    # RRF lineage
                    rrf_score=float(fused_score),
                    # preserve upstream lineage
                    dense_score=original_chunk.dense_score,
                    dense_rank=original_chunk.dense_rank,
                    bm25_score=original_chunk.bm25_score,
                    bm25_rank=original_chunk.bm25_rank,
                    # provenance
                    retrieval_sources=original_chunk.retrieval_sources,
                )
            )

        # Global ranking by fused RRF score
        # Global ranking by fused RRF score
        fused_chunks.sort(
            key=lambda chunk: chunk.score,
            reverse=True,
        )
        fused_chunks = apply_post_processing(
            chunks=fused_chunks,
            query=query,
            top_k=top_k,
            experiment_config=experiment_config,
        )
        print("\n===== HYBRID DEBUG =====")

        for idx, chunk in enumerate(fused_chunks, start=1):

            print(f"\nRANK {idx}")
            print(chunk.chunk.text[:400])
            print("SCORE:", chunk.score)

        print("========================\n")
        result = RetrievalResult(
            trace_id=trace_id,
            query=query,
            retrieved_chunks=fused_chunks[:top_k],
            latency_ms=(vector_result.latency_ms + lexical_result.latency_ms),
            retriever_version="hybrid_rrf_v1",
        )
        end_span(span)

        return result

    raise ValueError(f"Unknown retrieval strategy: {strategy}")