CohereRerankOperator

Use CohereRerankOperator to reorder documents by their relevance to a query with Cohere’s Rerank API.

Before you begin

Configure a Cohere connection. The operator uses the cohere_default connection unless another conn_id is provided.

The operator requires:

  • query: The search query used to evaluate relevance.

  • documents: A list of text documents to rank.

The operator uses rerank-v3.5 by default. Set model to use another Cohere model or an endpoint-specific deployment name. Use top_n to limit the number of results and max_tokens_per_doc to control how much of each document Cohere processes. The query, documents, and both limits are templated fields; rendered limit values are converted to integers before they are sent to Cohere.

Using the operator

tests/system/cohere/example_cohere_rerank_operator.py[source]

CohereRerankOperator(
    task_id="rerank_documents",
    query="What is the capital of the United States?",
    documents=[
        "Carson City is the capital city of Nevada.",
        "Washington, D.C. is the capital of the United States.",
        "The capital city of France is Paris.",
    ],
    top_n=2,
)

Output

The operator converts the Cohere response to an XCom-serializable dictionary. Its results list is ordered from most to least relevant. Each result contains the original document’s zero-based index and its relevance_score. Use the index to associate a result with the corresponding item in the input documents list.

An abbreviated response looks like this:

{
  "id": "rerank-request-id",
  "results": [
    {"index": 1, "relevance_score": 0.99},
    {"index": 0, "relevance_score": 0.12}
  ]
}

Was this entry helpful?