How the LLM Reranker Integrates with ai-memory’s Retrieval Pipeline

The LLM reranker integrates as a final relevance-scoring layer after reciprocal rank fusion (RRF), operating on an over-fetched candidate pool to semantically reorder results before returning them to the client.

The ai-memory retrieval system employs a multi-stage pipeline to surface relevant pages from a knowledge graph. Understanding how the LLM reranker fits into this flow requires examining the precise handoff between candidate generation, fusion algorithms, and final result ordering in the akitaonrails/ai-memory codebase.

Multi-Stage Retrieval Architecture

The pipeline processes queries through distinct phases before reaching the optional reranking stage.

Candidate Generation Streams

Initially, the memory_query RPC gathers hits from four independent streams, as documented in docs/ARCHITECTURE.md at line 353:

  • Full-text search (FTS5) against indexed content
  • Entity-match tokens extracted from the query
  • Graph-expansion neighbours traversed from matched entities
  • Optional vector similarity when embeddings are configured

Each stream produces a ranked list of PageHit objects. These lists remain separate until the fusion stage.

RRF Fusion and Authority Adjustment

In crates/ai-memory-store/src/reader.rs (lines 305-320), the rerank_page_hits function merges the independent streams using reciprocal rank fusion (RRF) scores. The implementation adjusts ranks according to page authority signals—including pinned status, content tier, and tags—then truncates the fused list to the caller-requested limit (default ≤ 30 results).

LLM Reranker Integration Point

The LLM reranker operates as a post-processing pass after RRF fusion completes but before the final response is serialized.

Configuration and Activation

Integration is gated by the environment variable AI_MEMORY_RERANKER=llm. When set and a valid LLM provider is configured, the server injects one additional relevance pass over the fused candidates. Notably, global-scope queries (global=true) bypass the LLM reranker entirely, preserving their pure-FTS ranking according to the architecture documentation at docs/ARCHITECTURE.md line 353.

The Reranker Trait Interface

The contract is defined in crates/ai-memory-llm/src/reranker.rs (lines 53-66) by the Reranker trait:

pub trait Reranker: Send + Sync {
    fn name(&self) -> &str;
    fn model(&self) -> &str;
    async fn rerank(&self, query: &str, candidates: &[Candidate]) -> Result<Vec<RerankScore>>;
}

The built-in LlmReranker implementation (lines 71-89) packs the query and up to 30 candidate titles/snippets into a bounded JSON prompt (approximately 600 bytes per snippet) and transmits it to the configured LLM provider. Implementations return scores ∈ [0, 1] for each candidate.

Server-Side Implementation Details

In crates/ai-memory-mcp/src/server.rs (lines 1624-1641), the rerank_hits function orchestrates the integration:

  1. Over-fetching: The candidate set is enlarged by a configurable factor (default ≈ 3× the final limit) so the LLM can promote hits that fell just below the cutoff (lines 1606-1614).
  2. Concurrency control: A semaphore (RERANK_MAX_IN_FLIGHT) limits simultaneous LLM calls to prevent resource exhaustion.
  3. Timeout protection: A hard RERANK_TIMEOUT prevents slow LLM responses from blocking the query pipeline.

If the LLM call fails, times out, or returns malformed scores, the server falls back to the pre-rerank order and logs a diagnostic entry (lines 1669-1687). When successful, the server reorders hits by the returned rerank_score values and includes these scores in the explain output (lines 1740-1750).

Enabling and Using the LLM Reranker

Activate the integration via environment variables:


# .env or shell export

AI_MEMORY_RERANKER=llm
AI_MEMORY_LLM_PROVIDER=openai
AI_MEMORY_LLM_MODEL=gpt-4

Start the server with reranking enabled:

ai-memory serve

Query with explanation to observe rerank scores:

ai-memory query --query "consensus algorithm implementation" --limit 5 --explain

Programmatic usage via the MCP client:

use ai_memory_mcp::client::MemoryClient;

#[tokio::main]
async fn main() {
    let client = MemoryClient::new("http://127.0.0.1:49374");
    
    let resp = client
        .memory_query()
        .query("how does the consensus algorithm work?")
        .limit(10)
        .explain(true)
        .await
        .unwrap();

    for hit in resp.hits {
        println!("{} (score {:.2})", hit.path, hit.score);
        if let Some(rerank) = hit.explain.rerank_score {
            println!("  LLM rerank: {:.2}", rerank);
        }
    }
}

Summary

  • The LLM reranker sits after RRF fusion and authority scoring but before final result serialization.
  • It evaluates an over-fetched candidate pool (3× the result limit) to allow recovery of high-relevance items initially ranked lower.
  • Integration is controlled by the AI_MEMORY_RERANKER=llm environment variable and defined by the Reranker trait in ai-memory-llm.
  • The implementation in ai-memory-mcp includes concurrency limits (RERANK_MAX_IN_FLIGHT) and timeouts (RERANK_TIMEOUT) to maintain pipeline responsiveness.
  • Graceful degradation ensures that LLM failures do not break retrieval; the system falls back to the RRF-ordered results.

Frequently Asked Questions

Where exactly does the LLM reranker sit in the processing order?

The reranker executes after the four candidate streams undergo RRF fusion and authority adjustment in ai-memory-store, but before the final JSON response is constructed in ai-memory-mcp. It receives the truncated candidate list, expands it via over-fetching, and reorders based on semantic relevance scores returned by the LLM.

What happens if the LLM reranker fails or times out?

According to the server implementation in crates/ai-memory-mcp/src/server.rs (lines 1669-1687), any failure—including timeouts, network errors, or malformed JSON responses—triggers a fallback to the original RRF-ordered results. The system logs the failure for diagnostics but returns valid results to the client without interruption.

How many candidates does the LLM reranker evaluate?

By default, the reranker evaluates up to 30 candidates, though the exact number depends on the query's limit parameter multiplied by an over-fetch factor (approximately 3×). The LlmReranker implementation packs these into a prompt with roughly 600 bytes allocated per snippet to stay within context window constraints.

Can I use the LLM reranker with global scope queries?

No. As specified in docs/ARCHITECTURE.md line 353, queries executed with global=true bypass the LLM reranker entirely. Global queries rely solely on full-text search ranking without the secondary semantic reordering pass to ensure deterministic performance across large datasets.

Have a question about this repo?

These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:

Share the following with your agent to get started:
curl -s "https://instagit.com/install.md"

Works with
Claude Codex Cursor VS Code OpenClaw Any MCP Client

Maintain an open-source project? Get it listed too →