codenib.ops.rerank

Classes:

Name	Description
CrossEncoderContext	Context for cross-encoder reranking.
RerankContext	Shared context carrying the rerank agent and configuration.

Functions:

Name	Description
rerank_by_embedding	Score candidate contents online with an embedding model.
rerank_by_indexed_embedding	Rerank candidates using vectors already stored in the search index.
rerank_by_cross_encoder	Rerank candidate contents with a shared cross-encoder context.

CrossEncoderContext dataclass

CrossEncoderContext(model_name: str, backend: str | None = None, batch_size: int = 16)

Context for cross-encoder reranking.

Wraps either :class:STCrossEncoderWrapper (sentence-transformers CrossEncoder — mxbai-rerank-, bge-reranker-, jina-reranker-) or :class:QwenRerankerWrapper (Qwen3-Reranker- yes/no logit scoring).

The underlying wrapper is loaded lazily on first :meth:score call so that building a CrossEncoderContext at startup costs nothing until the skill is actually invoked.

Parameters:

Name	Type	Description	Default
model_name	str	HuggingFace model id (e.g. "mixedbread-ai/mxbai-rerank-large-v2").	required
backend	str | None	"st" (sentence-transformers CrossEncoder), "qwen" (Qwen3-Reranker logit trick), or None to auto-detect from the model name ("Qwen3-Reranker" in name → qwen; else st).	None
batch_size	int	Scoring batch size forwarded to the wrapper.	16

Methods:

Name	Description
score	Score (query, doc) pairs; higher = more relevant.

score

score(query: str, docs: list[str]) -> list[float]

Score (query, doc) pairs; higher = more relevant.

Source code in codenib/ops/rerank.py

def score(self, query: str, docs: List[str]) -> List[float]:
    """Score (query, doc) pairs; higher = more relevant."""
    return self.ensure_wrapper().score(query, docs)

RerankContext dataclass

RerankContext(
    llm: LiteLLMChat | None = None,
    agent: RerankAgent | None = None,
    embedding_store: CodeVectorStore | None = None,
    top_k: int | None = None,
    candidate_top_k: int | None = None,
    window_size: int | None = None,
    window_step: int | None = None,
    listwise_format: Literal["structured", "rankgpt"] = "structured",
    embedding_source: Literal["content", "index"] = "content",
)

Shared context carrying the rerank agent and configuration.

rerank_by_embedding

rerank_by_embedding(
    query: str,
    candidates: Sequence[QueriedNode],
    embedding_store: CodeVectorStore,
    *,
    top_k: int | None = None
) -> list[QueriedNode]

Score candidate contents online with an embedding model.

This operator intentionally does not search the global vector index. It only re-embeds the provided candidate contents and sorts those candidates by similarity to the query. Use it after retrieval/graph expansion has already assembled the candidate set.

Source code in codenib/ops/rerank.py

def rerank_by_embedding(
    query: str,
    candidates: Sequence[QueriedNode],
    embedding_store: CodeVectorStore,
    *,
    top_k: Optional[int] = None,
) -> List[QueriedNode]:
    """Score candidate contents online with an embedding model.

    This operator intentionally does not search the global vector index. It only
    re-embeds the provided candidate contents and sorts those candidates by
    similarity to the query. Use it after retrieval/graph expansion has already
    assembled the candidate set.
    """
    candidates_with_content = [
        candidate for candidate in candidates if candidate.content
    ]
    if not candidates_with_content:
        return []

    query_vec = np.array(embedding_store.embedding.embed_query(query), dtype=np.float32)
    doc_vectors = np.array(
        embedding_store.embedding.embed_documents(
            [candidate.content for candidate in candidates_with_content]
        ),
        dtype=np.float32,
    )

    metric = embedding_store.index_metric
    if metric == "ip":
        scores = np.dot(doc_vectors, query_vec).tolist()
    elif metric == "l2":
        scores = (-np.linalg.norm(doc_vectors - query_vec, axis=1)).tolist()
    else:
        raise ValueError(f"Unsupported index metric: {metric!r}")

    ranked = sorted(
        zip(scores, candidates_with_content, strict=True),
        key=lambda pair: pair[0],
        reverse=True,
    )
    if top_k is not None:
        ranked = ranked[:top_k]
    return [_copy_with_score(candidate, float(score)) for score, candidate in ranked]

rerank_by_indexed_embedding

rerank_by_indexed_embedding(
    query: str,
    candidates: Sequence[QueriedNode],
    embedding_store: CodeVectorStore,
    *,
    top_k: int | None = None,
    level: str = "l2"
) -> list[QueriedNode]

Rerank candidates using vectors already stored in the search index.

Dense-first graph pipelines expand compact graph nodes identified by the same stable node_id values persisted in the vector store. Reusing those vectors keeps dense and graph candidates in one calibrated score space and avoids re-embedding potentially very large source spans online. Candidates absent from the index are retained at the end in their original order.

Source code in codenib/ops/rerank.py

def rerank_by_indexed_embedding(
    query: str,
    candidates: Sequence[QueriedNode],
    embedding_store: CodeVectorStore,
    *,
    top_k: Optional[int] = None,
    level: str = "l2",
) -> List[QueriedNode]:
    """Rerank candidates using vectors already stored in the search index.

    Dense-first graph pipelines expand compact graph nodes identified by the
    same stable ``node_id`` values persisted in the vector store. Reusing those
    vectors keeps dense and graph candidates in one calibrated score space and
    avoids re-embedding potentially very large source spans online. Candidates
    absent from the index are retained at the end in their original order.
    """

    from .retrieve import dedup_queried_nodes, to_queried_nodes

    if not candidates:
        return []
    mask = {candidate.node_id for candidate in candidates if candidate.node_id}
    if not mask:
        fallback = list(candidates)
        return fallback[:top_k] if top_k is not None else fallback
    indexed = dedup_queried_nodes(
        to_queried_nodes(
            embedding_store.search_within_ids(
                query,
                mask_node_ids=mask,
                top_k=len(candidates),
                level=level,
            )
        )
    )
    matched = {candidate.node_id for candidate in indexed if candidate.node_id}
    unmatched = [
        candidate
        for candidate in candidates
        if not candidate.node_id or candidate.node_id not in matched
    ]
    ranked = [*indexed, *unmatched]
    return ranked[:top_k] if top_k is not None else ranked

rerank_by_cross_encoder

rerank_by_cross_encoder(
    query: str,
    candidates: Sequence[QueriedNode],
    context: CrossEncoderContext,
    *,
    top_k: int | None = None
) -> list[QueriedNode]

Rerank candidate contents with a shared cross-encoder context.

Source code in codenib/ops/rerank.py

def rerank_by_cross_encoder(
    query: str,
    candidates: Sequence[QueriedNode],
    context: CrossEncoderContext,
    *,
    top_k: Optional[int] = None,
) -> List[QueriedNode]:
    """Rerank candidate contents with a shared cross-encoder context."""
    candidates_with_content = [
        candidate for candidate in candidates if candidate.content
    ]
    if not candidates_with_content:
        fallback = list(candidates)
        return fallback[:top_k] if top_k is not None else fallback

    scores = context.score(
        query,
        [str(candidate.content) for candidate in candidates_with_content],
    )
    if len(scores) != len(candidates_with_content):
        raise RuntimeError(
            "cross-encoder returned "
            f"{len(scores)} scores for {len(candidates_with_content)} candidates"
        )
    ranked_with_content = sorted(
        zip(scores, candidates_with_content, strict=True),
        key=lambda pair: pair[0],
        reverse=True,
    )
    ranked = [
        _copy_with_score(candidate, float(score))
        for score, candidate in ranked_with_content
    ]
    ranked.extend(candidate for candidate in candidates if not candidate.content)
    if top_k is not None:
        ranked = ranked[:top_k]
    return ranked