Source code for lalandre_rag.retrieval.trace

"""
Retrieval trace utilities
Centralized definition of trace keys and extraction helper.
"""

from typing import Any, Dict, Optional, Tuple

TRACE_KEYS: Tuple[str, ...] = (
    "query_id",
    "search_method",
    "collection",
    "source_collection",
    "lexical_rank",
    "semantic_rank",
    "rrf_score",
    "rrf_score_raw",
    "lexical_score",
    "lexical_score_raw",
    "lexical_score_normalized",
    "semantic_score",
    "semantic_score_raw",
    "fusion_method",
    "rerank_score",
    "pre_rerank_score",
    "rerank_model",
    "chunk_id",
    "chunk_index",
    "char_start",
    "char_end",
    "parent_id",
    "content_length",
    "text_length",
    "token_count",
    "content_hash",
    "query_variant",
    "query_variant_strategy",
    "query_variant_weight",
    "query_variant_index",
    "query_variant_rank",
    "query_expansion_enabled",
    "graph_query_strategy",
    "row_index",
)

TRACE_KEYS_SET = set(TRACE_KEYS)




def extract_trace(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Extract a compact, traceable subset of metadata for sources.
    Avoids duplicating large payload fields while preserving retrieval provenance.
    """
    if not metadata:
        return {}
    return {key: metadata[key] for key in TRACE_KEYS if key in metadata}