orchestrata_ultimate.semantic_graph.rerankers

Reference › orchestrata_ultimate › semantic_graph › rerankers

The reranking backend for semantic graph retrieval. OpenRouterReranker calls OpenRouter's rerank endpoint with a Cohere reranker model (default cohere/rerank-4-pro, 60-second timeout) to reorder candidate Chunk objects by relevance to a query. rerank is async and returns the top-k chunks (default 5). _request_rerank makes the real provider call over httpx; its docstring says it is reached only when the recorded-cassette contract allows a live request. Failures raise VectorSearchError from exceptions. Candidates come from hybrid search in ingestion, and the ingestion and evaluation scripts import this module. It does not read or write the graph.

At a glance

Class diagram

Dependency diagram

API

Classes

OpenRouterReranker

Line 21

Rerank candidate chunks via OpenRouter's rerank endpoint.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/semantic_graph/rerankers.py

"""Reranking backend for semantic graph retrieval.

Uses OpenRouter's rerank endpoint with Cohere rerankers to reorder candidate
chunks by relevance to a query.
"""

from __future__ import annotations

import logging
from collections.abc import Mapping
from typing import Any

import httpx

from orchestrata_ultimate.semantic_graph.exceptions import VectorSearchError
from orchestrata_ultimate.semantic_graph.models import Chunk

logger = logging.getLogger(__name__)


class OpenRouterReranker:
    """Rerank candidate chunks via OpenRouter's rerank endpoint.

    Defaults to ``cohere/rerank-4-pro``. The caller must supply an API key;
    the class never reads environment variables directly.
    """

    def __init__(
        self,
        api_key: str,
        model: str = "cohere/rerank-4-pro",
        base_url: str = "https://openrouter.ai/api/v1",
        timeout: float = 60.0,
        env: Mapping[str, str] | None = None,
    ) -> None:
        """Initialize the reranker.

        Args:
            api_key: OpenRouter API key.
            model: Model identifier on OpenRouter.
            base_url: OpenRouter API base URL.
            timeout: HTTP request timeout in seconds.
        """
        if not api_key:
            raise VectorSearchError("OpenRouter API key is required")
        self._api_key = api_key
        self._model = model
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        # See OpenRouterEmbedder.__init__ for why this seam exists.
        self._env = env

    async def rerank(
        self,
        query: str,
        chunks: list[Chunk],
        top_k: int = 5,
    ) -> list[Chunk]:
        """Return the most relevant chunks for ``query``.

        Args:
            query: The search query.
            chunks: Candidate chunks to rerank.
            top_k: Maximum number of chunks to return.

        Returns:
            Chunks ordered by descending relevance.
        """
        if not chunks:
            return []
        if len(chunks) == 1:
            return chunks[:top_k]

        documents = [{"id": chunk.id, "text": chunk.content} for chunk in chunks]
        top_n = min(top_k, len(chunks))
        payload = {
            "model": self._model,
            "query": query,
            "documents": documents,
            "top_n": top_n,
        }
        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }
        url = f"{self._base_url}/rerank"

        # Through the cassette contract: this path holds its own httpx client
        # and never reaches llm/client.py, so without it every rerank the suite
        # performed was a real billed OpenRouter request.
        from orchestrata_ultimate.llm._cassette import (  # noqa: PLC0415
            cassette_exchange,
        )

        def _synthetic() -> dict[str, Any]:
            # Input order, truncated to top_n: a rerank that reorders nothing
            # is the honest stand-in when no provider ranked anything.
            return {
                "results": [
                    {"id": document["id"], "relevance_score": 1.0 - index / top_n}
                    for index, document in enumerate(documents[:top_n])
                ]
            }

        data = await cassette_exchange(
            kind="rerank",
            payload=payload,
            synth=_synthetic,
            live=lambda: self._request_rerank(url, headers, payload),
            env=self._env,
        )

        results = data.get("results", [])
        chunk_by_id = {chunk.id: chunk for chunk in chunks}
        reranked: list[Chunk] = []
        for item in results:
            chunk_id = item.get("id") or item.get("index")
            if isinstance(chunk_id, int):
                chunk_id = documents[chunk_id]["id"]
            chunk = chunk_by_id.get(chunk_id)
            if chunk is not None:
                reranked.append(chunk)
        return reranked

    async def _request_rerank(
        self, url: str, headers: dict[str, str], payload: dict[str, Any]
    ) -> dict[str, Any]:
        """The real provider call. Reached only when the cassette contract
        says so — an unset LLM_CASSETTE_DIR, or an explicit re-record."""
        try:
            async with httpx.AsyncClient(timeout=self._timeout) as client:
                response = await client.post(url, headers=headers, json=payload)
                response.raise_for_status()
                data: dict[str, Any] = response.json()
        except httpx.HTTPStatusError as exc:
            body = exc.response.text if exc.response else ""
            logger.error("OpenRouter rerank request failed: %s", body)
            raise VectorSearchError(
                f"OpenRouter rerank HTTP error {exc.response.status_code if exc.response else 'unknown'}"
            ) from exc
        except httpx.RequestError as exc:
            logger.error("OpenRouter rerank request failed: %s", exc)
            raise VectorSearchError(f"OpenRouter rerank request failed: {exc}") from exc
        return data