How LightRAG's Reranking Integration Works with Jina AI and Cohere APIs

LightRAG provides a thin, provider-agnostic wrapper in lightrag/rerank.py that normalizes Jina AI and Cohere rerank APIs behind unified async helpers (jina_rerank() and cohere_rerank()), with optional token-aware chunking for long documents.

The HKUDS/LightRAG repository implements reranking as a modular component that bridges external third-party rerank services with the retrieval pipeline. This integration allows developers to boost retrieval accuracy by sending candidate documents through Jina AI or Cohere's specialized ranking models without managing provider-specific request formats manually.

Architecture of the Reranking Pipeline

LightRAG's reranking logic centers on lightrag/rerank.py, which orchestrates three distinct stages: payload construction, optional document chunking, and response normalization. The design abstracts provider differences while exposing granular control over token limits and chunking behavior.

Payload Construction and Provider Formatting

The public entry points cohere_rerank() (lines 68-82) and jina_rerank() (lines 35-43) delegate to the internal generic_rerank_api() function. This helper builds request bodies tailored to each provider's specification:

• Standard format (Jina / Cohere): A flat JSON object containing model, query, and documents (lines 67-72)
• Aliyun format: A nested structure with input and parameters keys (lines 45-55)

Environment variables drive authentication: COHERE_API_KEY or JINA_API_KEY are read at runtime via load_dotenv (line 19), with a fallback to RERANK_BINDING_API_KEY for generic bindings.

Optional Document Chunking

When enable_chunking=True, documents exceeding the provider's token limit are automatically segmented before the API call. The chunk_documents_for_rerank() function (lines 22-41) uses a TiktokenTokenizer (or character-based fallback) to respect max_tokens and overlap_tokens constraints.

During chunking, the original top_n limit is temporarily disabled (lines 38-43) so every chunk receives a relevance score. After the API returns results, aggregate_chunk_scores() (lines 16-71) collapses chunk-level scores back to document-level scores using the default max strategy.

Response Normalization and Post-Processing

Provider responses are parsed and normalized to a uniform list of dictionaries: [{"index": int, "relevance_score": float}]. If chunking was enabled, the aggregated document scores are optionally trimmed back to the original top_n (lines 58-64) before returning to the caller.

Provider-Specific Implementations

Each supported rerank service follows a consistent pattern while handling provider quirks internally.

Cohere Rerank Integration

The cohere_rerank() helper retrieves credentials from COHERE_API_KEY or RERANK_BINDING_API_KEY, then invokes generic_rerank_api() with response_format="standard" and return_documents=None. This configuration sends a POST request to https://api.cohere.com/v2/rerank without requesting document text in the response, as Cohere does not support returning full document content in rerank calls.

Jina AI Rerank Integration

The jina_rerank() function pulls credentials from JINA_API_KEY or RERANK_BINDING_API_KEY. It calls the generic API with return_documents=False to minimize payload size, though Jina's API technically supports returning document text. This keeps the response lightweight and consistent with Cohere's behavior.

Aliyun Support

For Aliyun's rerank service, ali_rerank() (lines 75-87) uses the nested request/response format and a dedicated endpoint URL, automatically handling the structural differences from the standard flat JSON format.

Implementing Reranking in Your Pipeline

The following examples demonstrate practical usage patterns for integrating these rerank functions into async RAG workflows.

Basic Cohere Rerank Without Chunking

For standard retrieval scenarios with reasonably sized documents, disable chunking to minimize API latency:

import asyncio
import os
from lightrag.rerank import cohere_rerank

async def rerank_candidates():
    query = "What are the health benefits of green tea?"
    documents = [
        "Green tea contains antioxidants that may reduce inflammation.",
        "Coffee is a popular morning beverage.",
        "Green tea can improve brain function."
    ]
    
    # Requires COHERE_API_KEY in environment

    results = await cohere_rerank(
        query=query,
        documents=documents,
        top_n=2,
        enable_chunking=False
    )
    
    for result in results:
        idx = result["index"]
        score = result["relevance_score"]
        print(f"Doc {idx} (score {score:.4f}): {documents[idx]}")

# Run the async function

asyncio.run(rerank_candidates())

Jina Rerank with Document Chunking

For documents exceeding the provider's token limit (512 tokens for Jina), enable chunking to ensure the full document content is evaluated:

import asyncio
from lightrag.rerank import jina_rerank

async def rerank_long_document():
    query = "Explain the theory of relativity."
    
    # Simulate a very long document

    long_document = ("Relativity describes the relationship between space and time. " * 200)
    documents = [long_document, "Quantum mechanics is a different branch of physics."]
    
    results = await jina_rerank(
        query=query,
        documents=documents,
        top_n=1,
        enable_chunking=True,
        max_tokens_per_doc=480  # Jina's 512 limit with safety margin

    )
    
    print(f"Best match: Document {results[0]['index']}")
    print(f"Relevance score: {results[0]['relevance_score']:.4f}")

asyncio.run(rerank_long_document())

When chunking is active, chunk_documents_for_rerank() splits the long text into overlapping segments, the API scores each chunk independently, and aggregate_chunk_scores() merges them using the max strategy before returning the final ranked list.

Environment Configuration

Switch between providers without code changes by using environment variables:

import os
import asyncio
from lightrag.rerank import jina_rerank, cohere_rerank

async def provider_agnostic_rerank(query, documents):
    provider = os.getenv("RERANK_PROVIDER", "cohere")
    top_n = int(os.getenv("RERANK_TOP_N", "3"))
    
    if provider == "jina":
        return await jina_rerank(query, documents, top_n=top_n)
    else:
        return await cohere_rerank(query, documents, top_n=top_n)

# Usage

query = "What is the capital of Italy?"
docs = ["Rome is Italy's capital.", "Milan is a major city.", "Naples is a coastal city."]
results = asyncio.run(provider_agnostic_rerank(query, docs))

Key Files and Functions

Understanding the source structure helps when debugging or extending the rerank functionality:

• lightrag/rerank.py – Core implementation containing generic_rerank_api(), provider helpers (cohere_rerank, jina_rerank, ali_rerank), chunking logic (chunk_documents_for_rerank), and score aggregation (aggregate_chunk_scores)
• lightrag/utils.py – Provides TiktokenTokenizer class used for token-aware chunking and the logger instance referenced throughout rerank code
• .env.example – Template defining required variables: COHERE_API_KEY, JINA_API_KEY, RERANK_BINDING_HOST, and RERANK_BINDING_API_KEY
• tests/test_rerank_chunking.py – Unit tests verifying chunk boundaries, score aggregation strategies, and provider-specific payload formats

Summary

• LightRAG's reranking integration resides in lightrag/rerank.py and provides async, provider-agnostic wrappers for Jina AI, Cohere, and Aliyun rerank APIs.
• Three-stage pipeline: Payload construction handles provider-specific JSON formats, optional chunking splits long documents using token-aware segmentation, and response normalization returns uniform {index, relevance_score} objects.
• Automatic chunking triggers when enable_chunking=True, using chunk_documents_for_rerank() to respect token limits and aggregate_chunk_scores() to merge results using the max strategy.
• Environment-driven configuration allows switching between Jina and Cohere by changing COHERE_API_KEY or JINA_API_KEY without modifying application code.

Frequently Asked Questions

How does LightRAG handle documents longer than the rerank API's token limit?

LightRAG automatically segments oversized documents when enable_chunking=True. The chunk_documents_for_rerank() function in lightrag/rerank.py splits text into overlapping chunks respecting max_tokens and overlap_tokens parameters. Each chunk receives an individual relevance score from the API, and aggregate_chunk_scores() merges these using the default max strategy to produce a single document-level score.

Can I use both Jina and Cohere rerankers in the same LightRAG pipeline?

Yes. LightRAG exposes separate async helpers—jina_rerank() and cohere_rerank()—that can be called conditionally based on configuration. Both functions normalize responses to the same format, allowing you to swap providers by changing environment variables (JINA_API_KEY vs COHERE_API_KEY) without altering downstream processing logic.

What is the difference between the standard and Aliyun request formats in LightRAG?

The standard format used by Jina and Cohere sends a flat JSON payload with model, query, and documents keys. Aliyun requires a nested structure with input containing the query and documents, plus a separate parameters object. The generic_rerank_api() function detects the provider type and constructs the appropriate payload automatically when using the ali_rerank() helper.

Does LightRAG return the reranked document text or just indices and scores?

LightRAG returns only index and relevance_score by default. Both cohere_rerank() and jina_rerank() set return_documents=False (or None for Cohere) to minimize payload size and maintain consistency. The caller should use the returned indices to reference the original document list passed to the rerank function.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}