Skip to content

@reranker_span

Wraps a reranking function with an OpenInference RERANKER span. Captures the query, input documents, output documents, model name, and top-k parameter as span attributes.

Decorator / Wrapper

Signature

def reranker_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "query",
    docs_arg: str = "documents",
    top_k: int | None = None,
) -> Callable

interface RerankerSpanOptions {
    name?: string;
    modelName?: string;
    topK?: number;
}

function rerankerSpan<A extends unknown[], R>(
    options: RerankerSpanOptions,
    fn: (...args: A) => R | Promise<R>,
): (...args: A) => Promise<R>

Parameters

Parameter Type Default Description
name str \| None None Span name. Defaults to the decorated function's name.
model_name str \| None None Reranker model name to record (e.g., "rerank-v3", "cross-encoder/ms-marco").
query_arg str "query" Name of the keyword argument containing the query string. Falls back to the first positional str argument.
docs_arg str "documents" Name of the keyword argument containing the input documents list. Falls back to the first positional list argument.
top_k int \| None None Number of top documents to return. Recorded as a span attribute.

Span Attributes

Attribute Value
openinference.span.kind "RERANKER"
reranker.model_name Model name (set when model_name parameter is provided)
reranker.query The query string (from query_arg kwarg or first str positional argument)
reranker.top_k Top-K value (set when top_k parameter is provided)
reranker.input_documents JSON-encoded input documents (from docs_arg kwarg or first list positional argument)
reranker.output_documents JSON-encoded output documents (set when the function returns list[Document])

Example

from coalex.ext.reranker import reranker_span
from coalex.ext.retrieval import Document

@reranker_span(
    name="cohere_rerank",
    model_name="rerank-english-v3.0",
    query_arg="query",
    docs_arg="documents",
    top_k=3,
)
def rerank(query: str, documents: list[Document]) -> list[Document]:
    """Rerank documents using Cohere's reranker API."""
    response = cohere_client.rerank(
        model="rerank-english-v3.0",
        query=query,
        documents=[d.content for d in documents],
        top_n=3,
    )
    return [
        Document(
            content=documents[r.index].content,
            id=documents[r.index].id,
            score=r.relevance_score,
            metadata=documents[r.index].metadata,
        )
        for r in response.results
    ]

# Usage
docs = [
    Document(content="Metformin reduces blood sugar.", id="doc-1", score=0.85),
    Document(content="Exercise improves insulin sensitivity.", id="doc-2", score=0.82),
    Document(content="Metformin side effects include nausea.", id="doc-3", score=0.79),
]

ranked = rerank(query="metformin side effects", documents=docs)
# Span "cohere_rerank" is emitted with:
#   openinference.span.kind = "RERANKER"
#   reranker.model_name = "rerank-english-v3.0"
#   reranker.query = "metformin side effects"
#   reranker.top_k = 3
#   reranker.input_documents = '[{"document.content": "Metformin reduces...", ...}, ...]'
#   reranker.output_documents = '[{"document.content": "Metformin side effects...", ...}, ...]'

import { rerankerSpan } from "@coalex-ai/sdk/ext";
import type { Document } from "@coalex-ai/sdk/ext";

const rerank = rerankerSpan(
    { name: "cohere_rerank", modelName: "rerank-english-v3.0", topK: 3 },
    async (query: string, documents: Document[]): Promise<Document[]> => {
        const response = await cohereClient.rerank({
            model: "rerank-english-v3.0",
            query,
            documents: documents.map(d => d.content),
            topN: 3,
        });
        return response.results.map(r => ({
            content: documents[r.index].content,
            id: documents[r.index].id,
            score: r.relevanceScore,
            metadata: documents[r.index].metadata,
        }));
    },
);

// Usage
const docs: Document[] = [
    { content: "Metformin reduces blood sugar.", id: "doc-1", score: 0.85 },
    { content: "Exercise improves insulin sensitivity.", id: "doc-2", score: 0.82 },
    { content: "Metformin side effects include nausea.", id: "doc-3", score: 0.79 },
];

const ranked = await rerank("metformin side effects", docs);
// Span "cohere_rerank" is emitted with:
//   openinference.span.kind = "RERANKER"
//   reranker.model_name = "rerank-english-v3.0"
//   reranker.query = "metformin side effects"
//   reranker.top_k = 3
//   reranker.input_documents = '[{"document.content": "Metformin reduces...", ...}, ...]'
//   reranker.output_documents = '[{"document.content": "Metformin side effects...", ...}, ...]'

Async Support

@reranker_span(name="async_rerank", model_name="rerank-v3", top_k=5)
async def rerank_async(query: str, documents: list[Document]) -> list[Document]:
    response = await async_cohere_client.rerank(
        query=query,
        documents=[d.content for d in documents],
        top_n=5,
    )
    return [
        Document(content=documents[r.index].content, id=documents[r.index].id, score=r.relevance_score)
        for r in response.results
    ]

const rerankAsync = rerankerSpan(
    { name: "async_rerank", modelName: "rerank-v3", topK: 5 },
    async (query: string, documents: Document[]): Promise<Document[]> => {
        const response = await asyncCohereClient.rerank({
            query,
            documents: documents.map(d => d.content),
            topN: 5,
        });
        return response.results.map(r => ({
            content: documents[r.index].content,
            id: documents[r.index].id,
            score: r.relevanceScore,
        }));
    },
);

RerankerSpan Context Manager / Class

For fine-grained control over when attributes are set on the span, use the RerankerSpan context manager (Python) or class (TypeScript).

Signature

class RerankerSpan:
    def __init__(
        self,
        name: str = "reranker",
        model_name: str | None = None,
        top_k: int | None = None,
    ) -> None: ...

    def set_query(self, query: str) -> None: ...
    def set_input_documents(self, docs: list[Document]) -> None: ...
    def set_output_documents(self, docs: list[Document]) -> None: ...

class RerankerSpan {
    constructor(name?: string, modelName?: string, topK?: number);
    start(): void;
    setQuery(query: string): void;
    setInputDocuments(docs: Document[]): void;
    setOutputDocuments(docs: Document[]): void;
    end(): void;
}

Methods

Python TypeScript Description
set_query(query) setQuery(query) Set the reranker.query attribute.
set_input_documents(docs) setInputDocuments(docs) Set the reranker.input_documents attribute (JSON-encoded).
set_output_documents(docs) setOutputDocuments(docs) Set the reranker.output_documents attribute (JSON-encoded).
-- start() Start the span (TypeScript only; Python uses with statement).
-- end() End the span (TypeScript only; Python uses with statement).

Example

from coalex.ext.reranker import RerankerSpan
from coalex.ext.retrieval import Document

with RerankerSpan("multi_model_rerank", model_name="cross-encoder/ms-marco", top_k=5) as r:
    r.set_query("diabetes treatment options")
    r.set_input_documents(candidate_docs)

    # Run your reranking logic
    scored_docs = cross_encoder.predict(
        [(query, doc.content) for doc in candidate_docs]
    )

    # Build output documents with new scores
    output_docs = [
        Document(content=doc.content, id=doc.id, score=score)
        for doc, score in zip(candidate_docs, scored_docs)
    ]
    output_docs.sort(key=lambda d: d.score or 0, reverse=True)
    output_docs = output_docs[:5]

    r.set_output_documents(output_docs)

import { RerankerSpan } from "@coalex-ai/sdk/ext";
import type { Document } from "@coalex-ai/sdk/ext";

const r = new RerankerSpan("multi_model_rerank", "cross-encoder/ms-marco", 5);
r.start();
try {
    r.setQuery("diabetes treatment options");
    r.setInputDocuments(candidateDocs);

    // Run your reranking logic
    const scores = await crossEncoder.predict(
        candidateDocs.map(doc => [query, doc.content]),
    );

    // Build output documents with new scores
    const outputDocs: Document[] = candidateDocs
        .map((doc, i) => ({ ...doc, score: scores[i] }))
        .sort((a, b) => (b.score ?? 0) - (a.score ?? 0))
        .slice(0, 5);

    r.setOutputDocuments(outputDocs);
} finally {
    r.end();
}

Full Retrieve-Rerank Pipeline

import coalex
from coalex.ext.retrieval import retrieval_span, Document
from coalex.ext.reranker import reranker_span
from openai import OpenAI

coalex.register(api_key="your-key")
coalex.auto_instrument()

client = OpenAI()

@retrieval_span(name="bm25_search", query_arg="query")
def search_bm25(query: str) -> list[Document]:
    results = elasticsearch.search(query=query, size=20)
    return [
        Document(content=r["_source"]["text"], id=r["_id"], score=r["_score"])
        for r in results["hits"]["hits"]
    ]

@reranker_span(name="cross_encoder_rerank", model_name="cross-encoder/ms-marco-MiniLM-L-12-v2", top_k=5)
def rerank(query: str, documents: list[Document]) -> list[Document]:
    pairs = [(query, doc.content) for doc in documents]
    scores = cross_encoder.predict(pairs)
    for doc, score in zip(documents, scores):
        doc.score = float(score)
    return sorted(documents, key=lambda d: d.score or 0, reverse=True)[:5]

with coalex.coalex_context(agent_id="search-agent", request_id="req-001"):
    candidates = search_bm25(query="metformin dosage for elderly")
    top_docs = rerank(query="metformin dosage for elderly", documents=candidates)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": f"Answer using: {top_docs[0].content}"},
            {"role": "user", "content": "What is the metformin dosage for elderly?"},
        ],
    )
    # Trace:
    #   coalex.invocation
    #     ├── bm25_search (RETRIEVER) - 20 input docs
    #     ├── cross_encoder_rerank (RERANKER) - 20 in, 5 out
    #     └── openai.chat.completions (LLM)

API Reference

coalex.ext.reranker

Reranker span decorator and context manager.

Classes

RerankerSpan

Context manager for instrumenting a reranking step as a RERANKER span.

Example::

with RerankerSpan("rerank", model_name="rerank-v3", top_k=5) as r:
    r.set_query(query)
    r.set_input_documents(docs)
    ranked = reranker.rerank(query, docs, top_k=5)
    r.set_output_documents(ranked)
Source code in coalex/ext/reranker.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class RerankerSpan:
    """Context manager for instrumenting a reranking step as a RERANKER span.

    Example::

        with RerankerSpan("rerank", model_name="rerank-v3", top_k=5) as r:
            r.set_query(query)
            r.set_input_documents(docs)
            ranked = reranker.rerank(query, docs, top_k=5)
            r.set_output_documents(ranked)
    """

    def __init__(self, name: str = "reranker", model_name: str | None = None, top_k: int | None = None) -> None:
        self._name = name
        self._model_name = model_name
        self._top_k = top_k
        self._span: Any = None
        self._ctx: Any = None

    def __enter__(self) -> RerankerSpan:
        tracer = _get_tracer()
        parent_ctx = _get_parent_context()
        self._ctx = tracer.start_as_current_span(self._name, context=parent_ctx)
        self._span = self._ctx.__enter__()
        self._span.set_attribute("openinference.span.kind", "RERANKER")
        if self._model_name:
            self._span.set_attribute("reranker.model_name", self._model_name)
        if self._top_k is not None:
            self._span.set_attribute("reranker.top_k", self._top_k)
        return self

    def set_query(self, query: str) -> None:
        self._span.set_attribute("reranker.query", query)

    def set_input_documents(self, docs: list[Document]) -> None:
        self._span.set_attribute("reranker.input_documents", encode_documents(docs))

    def set_output_documents(self, docs: list[Document]) -> None:
        self._span.set_attribute("reranker.output_documents", encode_documents(docs))

    def __exit__(self, *exc_info: Any) -> bool | None:
        return self._ctx.__exit__(*exc_info)

Functions

reranker_span 

reranker_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "query",
    docs_arg: str = "documents",
    top_k: int | None = None,
) -> Callable[[F], F]

Decorator that wraps a reranker function with an OpenInference RERANKER span.

Supports both sync and async functions. The decorated function must accept a list[Document] as input and return list[Document] as output. The span captures: - openinference.span.kind = RERANKER - reranker.model_name — if provided - reranker.query — from query_arg kwarg or first str positional - reranker.top_k — if provided - reranker.input_documents — JSON-encoded input docs - reranker.output_documents — JSON-encoded output docs

Parameters:

Name Type Description Default
name str | None

Span name. Defaults to the function name.

 None
model_name str | None

Reranker model name to record.

 None
query_arg str

Name of the kwarg containing the query string.

 'query'
docs_arg str

Name of the kwarg containing the input documents list.

 'documents'
top_k int | None

Number of top documents to return (recorded as attribute).

 None
Source code in coalex/ext/reranker.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def reranker_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "query",
    docs_arg: str = "documents",
    top_k: int | None = None,
) -> Callable[[F], F]:
    """Decorator that wraps a reranker function with an OpenInference RERANKER span.

    Supports both sync and async functions. The decorated function must accept a
    ``list[Document]`` as input and return ``list[Document]`` as output. The span captures:
    - ``openinference.span.kind`` = ``RERANKER``
    - ``reranker.model_name`` — if provided
    - ``reranker.query`` — from ``query_arg`` kwarg or first ``str`` positional
    - ``reranker.top_k`` — if provided
    - ``reranker.input_documents`` — JSON-encoded input docs
    - ``reranker.output_documents`` — JSON-encoded output docs

    Args:
        name: Span name. Defaults to the function name.
        model_name: Reranker model name to record.
        query_arg: Name of the kwarg containing the query string.
        docs_arg: Name of the kwarg containing the input documents list.
        top_k: Number of top documents to return (recorded as attribute).
    """

    def _set_span_attrs(span: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
        span.set_attribute("openinference.span.kind", "RERANKER")
        if model_name:
            span.set_attribute("reranker.model_name", model_name)
        if top_k is not None:
            span.set_attribute("reranker.top_k", top_k)
        q = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
        if q:
            span.set_attribute("reranker.query", str(q))
        input_docs = kwargs.get(docs_arg) or next((a for a in args if isinstance(a, list)), None)
        if input_docs and input_docs and isinstance(input_docs[0], Document):
            span.set_attribute("reranker.input_documents", encode_documents(input_docs))

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _get_tracer()
                parent_ctx = _get_parent_context()
                with tracer.start_as_current_span(span_name, context=parent_ctx) as span:
                    _set_span_attrs(span, args, kwargs)
                    result = await fn(*args, **kwargs)
                    if isinstance(result, list) and result and isinstance(result[0], Document):
                        span.set_attribute("reranker.output_documents", encode_documents(result))
                    return result

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = _get_tracer()
            parent_ctx = _get_parent_context()
            with tracer.start_as_current_span(span_name, context=parent_ctx) as span:
                _set_span_attrs(span, args, kwargs)
                result = fn(*args, **kwargs)
                if isinstance(result, list) and result and isinstance(result[0], Document):
                    span.set_attribute("reranker.output_documents", encode_documents(result))
                return result

        return wrapper  # type: ignore[return-value]

    return decorator