Search & Reranking Module Study Guide¶

class SearchResult:
    """Represents a search result with document and score"""

    def __init__(
        self,
        document: Document,
        score: float,
        highlights: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize search result

        Args:
            document: Document object containing content
            score: Relevance score (0.0-1.0)
            highlights: Optional highlighted snippets
        """

class SearchBackend(ABC):
    """Abstract base class for search backends"""

    @abstractmethod
    def initialize(self, config: Dict[str, Any]) -> None:
        """Initialize the search backend with configuration."""

    @abstractmethod
    def search(
        self,
        query_embeddings: Optional[np.ndarray],
        query_text: Optional[str],
        top_k: int = 10,
        filters: Optional[Dict[str, Any]] = None,
        ranking_strategy: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute a search query.

        Args:
            query_embeddings: Optional query embeddings for vector search
            query_text: Optional text query for keyword search
            top_k: Number of results to return
            filters: Optional filters to apply
            ranking_strategy: Optional ranking strategy to use

        Returns:
            List of search results with scores and metadata
        """

class QueryModality(Enum):
    """Content-modality taxonomy for queries and search results."""

    VIDEO = "video"
    IMAGE = "image"
    AUDIO = "audio"
    DOCUMENT = "document"
    TEXT = "text"
    MIXED = "mixed"


@dataclass
class RerankerSearchResult:
    """Search result shape consumed by rerankers for scoring and comparison.

    Distinct from `cogniverse_sdk.document.SearchResult` (Document-object based,
    used for API responses) — this dataclass is the reranker's internal view:
    flat fields for fast access during scoring loops.
    """

    id: str
    title: str
    content: str
    modality: str
    score: float
    metadata: Dict[str, Any]
    timestamp: Optional[datetime] = None

class SearchService:
    """Unified search service for video retrieval"""

    def __init__(
        self,
        config: Dict[str, Any],
        config_manager=None,
        schema_loader=None,
    ):
        """
        Initialize search service (profile-agnostic).

        Args:
            config: Configuration dictionary (full config.json content)
            config_manager: ConfigManager instance (REQUIRED - raises ValueError if None)
            schema_loader: SchemaLoader instance (REQUIRED - raises ValueError if None)
        """
        if config_manager is None:
            raise ValueError("config_manager is required")
        if schema_loader is None:
            raise ValueError("schema_loader is required")

        self.config = config
        self.config_manager = config_manager
        self.schema_loader = schema_loader

def search(
    self,
    query: str,
    profile: str,
    tenant_id: str,
    top_k: int = 10,
    filters: Optional[Dict[str, Any]] = None,
    ranking_strategy: Optional[str] = None,
) -> List[SearchResult]:
    """
    Search for videos matching the query.

    Args:
        query: Text query
        profile: Backend profile to use (e.g. "video_colpali_smol500_mv_frame")
        tenant_id: Tenant identifier
        top_k: Number of results to return
        filters: Optional filters (date range, etc.)
        ranking_strategy: Optional ranking strategy override

    Returns:
        List of SearchResult objects
    """

def get_document(
    self, document_id: str, tenant_id: str, profile: str
) -> Optional[Dict[str, Any]]:
    """
    Retrieve a specific document by ID.

    Returns a plain dict (document_id, source_id, content_type, metadata,
    temporal_info) rather than a Document/SearchResult object, or None if
    the backend has no matching document.
    """

class MultiModalReranker:
    """
    Heuristic reranker considering multiple modalities

    Scoring Components:
    - Cross-modal relevance (0.3 weight)
    - Temporal alignment (0.2 weight)
    - Result complementarity (0.2 weight)
    - Diversity bonus (0.15 weight)
    - Original score (0.15 weight)
    """

    def __init__(
        self,
        cross_modal_weight: float = 0.3,
        temporal_weight: float = 0.2,
        complementarity_weight: float = 0.2,
        diversity_weight: float = 0.15,
        original_score_weight: float = 0.15
    ):
        """Weights must sum to 1.0"""

async def rerank(
    self,
    query: str,
    results: List[RerankerSearchResult],
    modalities: Optional[List[QueryModality]] = None,
    context: Optional[Dict] = None,
) -> List[RerankerSearchResult]:
    """Uniform reranker entry point shared by all rerankers (MultiModalReranker,
    LearnedReranker, HybridReranker all expose this same signature). Delegates
    to `rerank_results`, defaulting modalities to []. This is what
    `rerank_service.build_reranker(...).rerank(...)` calls regardless of strategy.
    """

async def rerank_results(
    self,
    results: List[RerankerSearchResult],
    query: str,
    modalities: List[QueryModality],
    context: Optional[Dict] = None
) -> List[RerankerSearchResult]:
    """
    Rerank using heuristic multi-modal scoring

    Calculates:
    - Cross-modal alignment (modality compatibility matrix)
    - Temporal alignment (time range matching)
    - Complementarity (unique information contribution)
    - Diversity (modality distribution balance)

    Returns results with updated metadata:
    - reranking_score
    - score_components (breakdown)
    """

class LearnedReranker:
    """
    Unified learned reranker using LiteLLM

    Supported Models:
    - Cohere: rerank-english-v3.0
    - Together AI: Llama-Rank (ColBERT-style)
    - Jina AI: jina-reranker-v2-base-multilingual
    - Ollama: bge-reranker-v2-m3, mxbai-rerank-large-v2
    - Any LiteLLM-supported reranker
    """

    def __init__(
        self,
        model: Optional[str] = None,
        tenant_id: str = None,
        config_manager: "ConfigManager" = None
    ):
        """
        Initialize from config.json "reranking" section

        Args:
            model: Model name (e.g., "cohere/rerank-english-v3.0")
                   If None, loads from config.json
            tenant_id: Tenant identifier for config scoping
            config_manager: ConfigManager instance (required)

        Config structure:
        {
          "reranking": {
            "model": "cohere",
            "supported_models": {
              "cohere": "cohere/rerank-english-v3.0",
              "ollama_bge": "openai/bge-reranker-v2-m3"
            },
            "api_base": "http://localhost:11434/v1"  # For Ollama (OpenAI-compat endpoint)
          }
        }
        """

async def rerank(
    self,
    query: str,
    results: List[RerankerSearchResult],
    top_n: Optional[int] = None
) -> List[RerankerSearchResult]:
    """
    Rerank using LiteLLM neural model

    Process:
    1. Limit to max_results_to_rerank (default 100)
    2. Prepare documents as "title content" strings
    3. Call LiteLLM arerank API
    4. Map relevance scores back to RerankerSearchResult
    5. Add metadata: reranking_score, reranker_model, original_rank

    Returns:
        Reranked results with neural scores
    """

class HybridReranker:
    """
    Combines heuristic and learned reranking

    Strategies:
    - weighted_ensemble: Parallel scoring with weighted combination
    - cascade: Heuristic filter → learned rerank (efficient)
    - consensus: Borda count requiring agreement from both
    """

    def __init__(
        self,
        heuristic_reranker: Optional[MultiModalReranker] = None,
        learned_reranker: Optional[LearnedReranker] = None,
        strategy: Optional[str] = None,
        learned_weight: Optional[float] = None,
        heuristic_weight: Optional[float] = None
    ):
        """
        Loads from config.json if parameters are None

        Config example:
        {
          "reranking": {
            "hybrid_strategy": "weighted_ensemble",
            "learned_weight": 0.7,
            "heuristic_weight": 0.3
          }
        }
        """

async def rerank(
    self,
    query: str,
    results: List[RerankerSearchResult],
    modalities: Optional[List[QueryModality]] = None,
    context: Optional[Dict] = None,
) -> List[RerankerSearchResult]:
    """Uniform reranker entry point shared by all rerankers.

    Delegates to `rerank_hybrid`; callers without modality detection pass
    none and the heuristic/learned scoring degrades gracefully.
    """

# libs/agents/cogniverse_agents/search/rerank_service.py

def build_reranker(strategy: str, tenant_id: str, config_manager=None):
    if strategy == "learned":
        return LearnedReranker(tenant_id=tenant_id, config_manager=config_manager)
    if strategy == "hybrid":
        return HybridReranker(tenant_id=tenant_id, config_manager=config_manager)
    if strategy == "multi_modal":
        return MultiModalReranker()  # heuristic, no config required
    raise ValueError(f"Unknown strategy: {strategy}")

# The search router calls rerank_result_dicts(); ValueError from
# build_reranker surfaces as a 400 response.
reranked_dicts = await rerank_result_dicts(
    query=query, results=result_dicts, strategy=strategy,
    tenant_id=tenant_id, config_manager=config_manager
)

async def apply_reranking_to_traces(
    traces: List[Dict[str, Any]],
    strategy: str,
    config: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
    """Rerank each trace's ``results`` with the live ``strategy`` reranker.

    Args:
        traces: trace dicts, each with a ``query`` and a ``results`` list of dicts.
        strategy: live reranker name ("learned" / "hybrid" / "multi_modal").
        config: optional; reads ``tenant_id`` and ``config_manager`` (needed by
            the learned/hybrid rerankers to load their tenant-scoped settings).

    Returns:
        The same traces, with each ``results`` reordered by the reranker. A
        per-trace failure is logged and leaves that trace's results unchanged
        rather than crashing the whole evaluation run.
    """

from pathlib import Path

from cogniverse_agents.search.service import SearchService
from cogniverse_core.schemas.filesystem_loader import FilesystemSchemaLoader
from cogniverse_foundation.config.utils import create_default_config_manager

# `config` is the full config.json content; SearchService reads `search_backend`
# and `config["backend"]` directly, but profile/URL/port resolution happens
# through `config_manager` at query time (see _get_profile_config/_get_backend),
# so a minimal dict is enough here.
config = {"search_backend": "vespa"}
config_manager = create_default_config_manager()
# SchemaLoader is an abstract interface (cogniverse_sdk.interfaces.schema_loader);
# pass a concrete implementation.
schema_loader = FilesystemSchemaLoader(Path("configs/schemas"))

service = SearchService(
    config,
    config_manager=config_manager,
    schema_loader=schema_loader
)

# Perform search
results = service.search(
    query="Show me videos about quantum computing",
    profile="video_colpali_smol500_mv_frame",
    tenant_id="user_123",
    top_k=10,
)

# Access results
for result in results:
    print(f"Document: {result.document.id}")
    print(f"Score: {result.score}")
    print(f"Metadata: {result.document.metadata}")

from datetime import datetime
from cogniverse_agents.search.multi_modal_reranker import MultiModalReranker
from cogniverse_agents.search.types import QueryModality

# Initialize reranker
reranker = MultiModalReranker(
    cross_modal_weight=0.3,
    temporal_weight=0.2,
    complementarity_weight=0.2,
    diversity_weight=0.15,
    original_score_weight=0.15
)

# Detect query modalities
modalities = [QueryModality.VIDEO, QueryModality.TEXT]

# Add temporal context
context = {
    "temporal": {
        "time_range": (
            datetime(2024, 1, 1),
            datetime(2024, 12, 31)
        )
    }
}

# Rerank results
reranked = await reranker.rerank_results(
    results=search_results,
    query="quantum computing experiments",
    modalities=modalities,
    context=context
)

# Analyze score breakdown
for result in reranked[:5]:
    print(f"Final Score: {result.metadata['reranking_score']:.3f}")
    print(f"Components: {result.metadata['score_components']}")

from cogniverse_agents.search.learned_reranker import LearnedReranker
from cogniverse_foundation.config.utils import create_default_config_manager

# Initialize with Cohere model
config_manager = create_default_config_manager()
# Option 1: Load model from config.json
reranker = LearnedReranker(
    config_manager=config_manager,
    tenant_id="your_org:production"
)

# Option 2: Override with specific model
reranker = LearnedReranker(
    model="cohere/rerank-english-v3.0",
    config_manager=config_manager,
    tenant_id="your_org:production"
)

# Rerank with neural model
reranked = await reranker.rerank(
    query="quantum computing breakthroughs",
    results=search_results,
    top_n=10
)

# Check model info
info = reranker.get_model_info()  # Returns Dict[str, Any]
print(f"Model: {info['model']}")
print(f"Max rerank: {info['max_results_to_rerank']}")

from cogniverse_agents.search.hybrid_reranker import HybridReranker
from cogniverse_agents.search.multi_modal_reranker import MultiModalReranker
from cogniverse_agents.search.learned_reranker import LearnedReranker
from cogniverse_agents.search.types import QueryModality
from cogniverse_foundation.config.utils import create_default_config_manager

# Initialize hybrid reranker
config_manager = create_default_config_manager()
hybrid = HybridReranker(
    heuristic_reranker=MultiModalReranker(),
    learned_reranker=LearnedReranker(
        config_manager=config_manager,
        tenant_id="your_org:production"
    ),
    strategy="weighted_ensemble",
    learned_weight=0.7,
    heuristic_weight=0.3,
    tenant_id="your_org:production",
    config_manager=config_manager
)

# Rerank with hybrid approach
reranked = await hybrid.rerank_hybrid(
    query="machine learning tutorials",
    results=search_results,
    modalities=[QueryModality.VIDEO],
    context={}
)

# Examine fusion details
for result in reranked[:3]:
    metadata = result.metadata
    print(f"Final: {metadata['reranking_score']:.3f}")
    print(f"  Heuristic: {metadata['heuristic_score']:.3f}")
    print(f"  Learned: {metadata['learned_score']:.3f}")
    print(f"  Strategy: {metadata['fusion_strategy']}")

import asyncio

from cogniverse_foundation.config.utils import create_default_config_manager

# Initialize config manager (REQUIRED for learned/hybrid)
config_manager = create_default_config_manager()
tenant_id = "your_org:production"

# Mirror the search router: pick the reranker for the configured strategy
strategy = "hybrid"  # one of: "multi_modal", "learned", "hybrid"

if strategy == "learned":
    from cogniverse_agents.search.learned_reranker import LearnedReranker

    reranker = LearnedReranker(tenant_id=tenant_id, config_manager=config_manager)
elif strategy == "hybrid":
    from cogniverse_agents.search.hybrid_reranker import HybridReranker

    reranker = HybridReranker(tenant_id=tenant_id, config_manager=config_manager)
else:  # "multi_modal" — heuristic base, no config needed
    from cogniverse_agents.search.multi_modal_reranker import MultiModalReranker

    reranker = MultiModalReranker()

# rerank() is async and has the same signature on all three reranker classes
reranked = asyncio.run(
    reranker.rerank(query="deep learning frameworks", results=search_results)
)

from cogniverse_agents.search.hybrid_reranker import HybridReranker
from cogniverse_agents.search.types import QueryModality
from cogniverse_foundation.config.utils import create_default_config_manager

# Cascade: Fast heuristic filter → Expensive learned rerank
config_manager = create_default_config_manager()
hybrid = HybridReranker(
    strategy="cascade",
    tenant_id="your_org:production",
    config_manager=config_manager
)

# This processes 100 results → filters to 50 → reranks 50
# Much faster than reranking all 100 with learned model
reranked = await hybrid.rerank_hybrid(
    query="neural networks",
    results=large_result_set,  # 100 results
    modalities=[QueryModality.TEXT],
    context={}
)

# Only top 50 were sent to expensive learned model
print(f"Reranked {len(reranked)} results efficiently")

# Track reranking performance using search_span
from cogniverse_foundation.telemetry.context import search_span

with search_span(
    tenant_id="user_123",
    query=query,
    top_k=100,
    ranking_strategy="hybrid_weighted"
) as span:
    reranked = await reranker.rerank(query, results, modalities, context)

    # Add metrics to span
    span.set_attribute("reranked_count", len(reranked))
    span.set_attribute("top_score", reranked[0].metadata['reranking_score'])

try:
    reranked = await learned_reranker.rerank(query, results)
except Exception as e:
    logger.error(f"Learned reranking failed: {e}")
    # Fallback to heuristic
    reranked = await heuristic_reranker.rerank_results(
        results, query, modalities, context
    )

from cogniverse_agents.search.multi_modal_reranker import MultiModalReranker
from cogniverse_agents.search.types import QueryModality, RerankerSearchResult

# Note: There are TWO distinct result classes:
# 1. cogniverse_sdk.document.SearchResult (Document-object based, for API
#    responses; re-exported from cogniverse_agents.search)
# 2. cogniverse_agents.search.types.RerankerSearchResult (dataclass, for reranking)

def test_cross_modal_score():
    """Test cross-modal scoring logic (see tests/routing/unit/test_multi_modal_reranker.py)"""
    reranker = MultiModalReranker()

    # Create test result using RerankerSearchResult dataclass
    result = RerankerSearchResult(
        id="v1",
        title="Test",
        content="Content",
        modality="video",
        score=0.8,
        metadata={}
    )

    # Direct modality match
    score = reranker._calculate_cross_modal_score(
        result, "query", [QueryModality.VIDEO]
    )
    assert score == 1.0

    # Compatible match (image for video query)
    score = reranker._calculate_cross_modal_score(
        result, "query", [QueryModality.IMAGE]
    )
    assert score == 0.7

@pytest.mark.asyncio
async def test_hybrid_weighted_ensemble(sample_results, mock_config_manager):
    hybrid_reranker = HybridReranker(
        strategy="weighted_ensemble",
        tenant_id="test:unit",
        config_manager=mock_config_manager,
    )

    reranked = await hybrid_reranker.rerank_hybrid(
        query="quantum computing",
        results=sample_results,
        modalities=[QueryModality.VIDEO],
        context={}
    )

    # Verify
    assert len(reranked) == len(sample_results)
    assert all("reranking_score" in r.metadata for r in reranked)
    assert all("fusion_strategy" in r.metadata for r in reranked)

    # Check score monotonicity
    scores = [r.metadata['reranking_score'] for r in reranked]
    assert scores == sorted(scores, reverse=True)

import asyncio
import timeit

# Heuristic reranking benchmark
def bench_heuristic():
    reranker = MultiModalReranker()
    asyncio.run(reranker.rerank_results(results, query, modalities, {}))

time = timeit.timeit(bench_heuristic, number=100) / 100
print(f"Heuristic: {time*1000:.1f}ms")

# Learned reranking benchmark
def bench_learned():
    reranker = LearnedReranker(tenant_id=tenant_id, config_manager=config_manager)
    asyncio.run(reranker.rerank(query, results))

time = timeit.timeit(bench_learned, number=10) / 10
print(f"Learned: {time*1000:.1f}ms")  # dominated by the network round-trip to the reranking API/model