Semantic Search & Reranking Skill

Role

You are an elite retrieval engineering specialist. You know that first-pass vector

retrieval is recall-optimised (fast, broad) but not precision-optimised. Reranking

is the step that transforms good retrieval into excellent retrieval — the difference

between a RAG system that frustrates and one that impresses.

The Two-Stage Retrieval Architecture

Stage 1: Recall      → Fast ANN vector search (top-50 candidates, ~10ms)
Stage 2: Precision   → Slow cross-encoder rerank (top-5 from 50, ~100ms)

Why two stages?
- Cross-encoders are 10x more accurate than bi-encoders but 100x slower
- We use the fast bi-encoder to narrow from millions → 50
- We use the accurate cross-encoder to find the best 5 from those 50
- Net result: production speed + research-grade accuracy

Part 1: Cohere Rerank (Production Standard)

Used by Perplexity, Elastic, and enterprise RAG systems.

# pip install cohere
import cohere
import os

co = cohere.Client(os.environ.get("COHERE_API_KEY"))


def cohere_rerank(
    query: str,
    documents: list[str | dict],
    top_n: int = 5,
    model: str = "rerank-english-v3.0",
    return_documents: bool = True
) -> list[dict]:
    """
    Rerank documents using Cohere's cross-encoder.
    Dramatically improves RAG precision vs vector search alone.

    documents: list of strings, or list of dicts with a "text" key
    """
    # Normalise to strings
    if documents and isinstance(documents[0], dict):
        doc_texts = [d.get("text", d.get("content", str(d))) for d in documents]
        original_docs = documents
    else:
        doc_texts = documents
        original_docs = [{"text": d} for d in documents]

    response = co.rerank(
        query=query,
        documents=doc_texts,
        top_n=top_n,
        model=model,
        return_documents=return_documents
    )

    results = []
    for item in response.results:
        result = {
            "index": item.index,
            "relevance_score": item.relevance_score,
            "text": doc_texts[item.index]
        }
        # Merge original metadata back in
        if isinstance(original_docs[item.index], dict):
            result.update({
                k: v for k, v in original_docs[item.index].items()
                if k != "text"
            })
        results.append(result)

    return results


def two_stage_retrieval(
    query: str,
    vector_store,
    embed_fn: callable,
    recall_top_k: int = 50,
    precision_top_n: int = 5,
    metadata_filter: dict = None
) -> list[dict]:
    """
    Production two-stage retrieval: vector recall → Cohere rerank.
    """
    # Stage 1: Vector recall (broad)
    candidates = vector_store.query(query, n_results=recall_top_k,
                                     where=metadata_filter)

    if not candidates:
        return []

    # Stage 2: Rerank (precise)
    reranked = cohere_rerank(
        query=query,
        documents=candidates,
        top_n=precision_top_n
    )

    return reranked

Part 2: Local CrossEncoder Reranking (Free, Private)

# pip install sentence-transformers
from sentence_transformers import CrossEncoder
import numpy as np

class LocalReranker:
    """
    Local cross-encoder reranking using sentence-transformers.
    No API costs, fully private.

    Best models:
    - cross-encoder/ms-marco-MiniLM-L-6-v2  (fast, good quality)
    - BAAI/bge-reranker-large                (best quality, slower)
    - cross-encoder/ms-marco-electra-base   (production balance)
    """

    _instances = {}

    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        if model_name not in self.__class__._instances:
            self.__class__._instances[model_name] = CrossEncoder(model_name)
        self.model = self.__class__._instances[model_name]
        self.model_name = model_name

    def rerank(self, query: str, documents: list[str | dict],
                top_n: int = 5) -> list[dict]:
        """Rerank documents against query using cross-encoder."""
        if not documents:
            return []

        # Normalise
        if isinstance(documents[0], dict):
            doc_texts = [d.get("text", d.get("content", "")) for d in documents]
            originals = documents
        else:
            doc_texts = documents
            originals = [{"text": d} for d in documents]

        # Score all query-doc pairs
        pairs = [(query, doc) for doc in doc_texts]
        scores = self.model.predict(pairs)

        # Sort by score descending
        indexed_scores = sorted(
            enumerate(scores), key=lambda x: x[1], reverse=True
        )[:top_n]

        return [
            {
                **originals[idx],
                "text": doc_texts[idx],
                "relevance_score": float(score),
                "rank": rank
            }
            for rank, (idx, score) in enumerate(indexed_scores)
        ]

    def score_pair(self, query: str, document: str) -> float:
        """Score a single query-document pair."""
        return float(self.model.predict([(query, document)])[0])

Part 3: FlashRank (Microsecond Local Reranking)

Used when reranking latency must be under 10ms.

# pip install flashrank
from flashrank import Ranker, RerankRequest

class FlashReranker:
    """
    Ultra-fast local reranker using FlashRank.
    10-50x faster than CrossEncoder for same quality tier.
    """

    def __init__(self, model_name: str = "ms-marco-MiniLM-L-12-v2",
                  cache_dir: str = "/tmp/flashrank_cache"):
        self.ranker = Ranker(model_name=model_name, cache_dir=cache_dir)

    def rerank(self, query: str, passages: list[str | dict],
                top_k: int = 5) -> list[dict]:
        """Rerank passages — microsecond inference."""
        if isinstance(passages[0], str):
            passage_dicts = [{"id": i, "text": p} for i, p in enumerate(passages)]
        else:
            passage_dicts = passages

        request = RerankRequest(query=query, passages=passage_dicts)
        results = self.ranker.rerank(request)

        return [
            {
                "text": r.text,
                "score": r.score,
                "id": r.id
            }
            for r in results[:top_k]
        ]

Part 4: Contextual Compression

Anthropic's RAG technique: after retrieval, extract only the relevant sentences

rather than returning full chunks. Reduces noise and token usage.

def contextual_compress(
    query: str,
    documents: list[str],
    llm_fn: callable,
    compression_ratio: float = 0.4
) -> list[str]:
    """
    Compress retrieved documents to only the relevant parts.
    Used by Claude internally and in production RAG pipelines.
    Reduces hallucination by removing irrelevant context.
    """
    compressed = []

    for doc in documents:
        prompt = f"""Extract only the sentences from this passage that are relevant to answering the question.
If no sentences are relevant, respond with: "NOT RELEVANT"
If relevant sentences are found, return only those sentences verbatim.

Question: {query}

Passage:
{doc}

Relevant sentences:"""

        result = llm_fn(prompt)

        if result.strip() != "NOT RELEVANT" and result.strip():
            compressed.append(result.strip())

    return compressed


def rerank_and_compress_pipeline(
    query: str,
    candidates: list[dict],
    llm_fn: callable,
    reranker: LocalReranker,
    final_top_n: int = 3
) -> list[dict]:
    """
    Full pipeline: rerank → compress → return.
    This is the production-grade RAG context preparation pattern.
    """
    # Rerank
    reranked = reranker.rerank(query, candidates, top_n=final_top_n * 2)

    # Compress
    texts = [r["text"] for r in reranked]
    compressed_texts = contextual_compress(query, texts, llm_fn)

    # Reassemble with metadata
    result = []
    for i, (doc, compressed) in enumerate(zip(reranked, compressed_texts)):
        if compressed:
            result.append({
                **doc,
                "text": compressed,
                "original_text": doc["text"],
                "compressed": True,
                "final_rank": i
            })

    return result[:final_top_n]

Part 5: Maximal Marginal Relevance (Diversity Search)

Prevents returning 5 nearly-identical chunks from the same passage.

import numpy as np

def maximal_marginal_relevance(
    query_embedding: list[float],
    doc_embeddings: list[list[float]],
    documents: list[dict],
    top_k: int = 5,
    lambda_mult: float = 0.5
) -> list[dict]:
    """
    MMR selection: balance relevance to query with diversity among results.
    lambda_mult: 1.0 = pure relevance, 0.0 = pure diversity, 0.5 = balanced

    Used by LangChain, Claude, and production RAG to avoid redundant context.
    """
    q = np.array(query_embedding)
    docs = np.array(doc_embeddings)

    # Normalise
    q = q / (np.linalg.norm(q) + 1e-10)
    docs = docs / (np.linalg.norm(docs, axis=1, keepdims=True) + 1e-10)

    # Initial relevance scores
    relevance = docs @ q
    selected_indices = []
    remaining = list(range(len(docs)))

    for _ in range(min(top_k, len(docs))):
        if not selected_indices:
            # First: pick most relevant
            best = remaining[np.argmax(relevance[remaining])]
        else:
            # Subsequent: balance relevance vs diversity
            selected_embs = docs[selected_indices]
            scores = []
            for idx in remaining:
                rel = relevance[idx]
                # Redundancy = max similarity to already selected
                redundancy = np.max(docs[idx] @ selected_embs.T)
                mmr_score = lambda_mult * rel - (1 - lambda_mult) * redundancy
                scores.append(mmr_score)
            best = remaining[np.argmax(scores)]

        selected_indices.append(best)
        remaining.remove(best)

    return [documents[i] for i in selected_indices]

Output Standards

• Always use two-stage retrieval in production: recall (top-50) → rerank (top-5)
• Use Cohere rerank for highest quality; local CrossEncoder for private data
• Apply contextual compression when chunks are long (>500 tokens)
• Use MMR when diversity of information matters more than single-topic depth
• Measure reranking impact: compare end-to-end answer quality before/after — expect 15-40% improvement
• For IMI: rerank research findings by brand + query relevance before building prompts

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)