confidence.py

"""
Confidence scoring for retrieved text and image results.

The confidence score is a hand-tuned blend of several signals:
  - Vector similarity score from Chroma
  - Source hint match (does the top doc come from the expected product file?)
  - Token overlap ratio between query and top doc
  - Margin between #1 and #2 vector scores

A score below RAG_CONFIDENCE_THRESHOLD triggers fallback behaviour.
"""
import logging
from langchain_core.documents import Document

from retrieval.query import (
    normalize_for_match,
    infer_source_hints,
    source_matches_hints,
    path_matches_hints,
    rewrite_query_for_retrieval,
)
from retrieval.fusion import doc_identity_key

logger = logging.getLogger(__name__)


def build_vector_score_map(
    vector_results: list[tuple[Document, float]],
) -> dict[tuple, float]:
    """Build a {doc_key: normalized_score} map from similarity_search results."""
    score_map: dict[tuple, float] = {}
    for doc, score in vector_results or []:
        key = doc_identity_key(doc)
        normalized = min(max(float(score), 0.0), 1.0)
        if key not in score_map or normalized > score_map[key]:
            score_map[key] = normalized
    return score_map


def token_overlap_ratio(question: str, doc: Document) -> float:
    """Fraction of query tokens that appear in the doc source name + content."""
    q_tokens = set(normalize_for_match(question).split())
    if not q_tokens or doc is None:
        return 0.0
    source = normalize_for_match(doc.metadata.get("source", ""))
    content = normalize_for_match(doc.page_content[:800])
    doc_tokens = set((source + " " + content).split())
    if not doc_tokens:
        return 0.0
    overlap = len(q_tokens.intersection(doc_tokens))
    return min(overlap / max(1, len(q_tokens)), 1.0)


def compute_text_confidence(
    question: str,
    docs: list[Document],
    vector_score_map: dict[tuple, float],
) -> float:
    """
    Returns a confidence score in [0, 1].

    Weights:
      0.45 × vector similarity of top doc
      0.25 × source hint signal (1.0 if source matches inferred product, 0.6 if no hints)
      0.20 × token overlap ratio
      0.10 × score margin between top-1 and top-2
    """
    if not docs:
        return 0.0

    top_doc = docs[0]
    top_key = doc_identity_key(top_doc)
    top_score = vector_score_map.get(top_key, 0.0)
    top_overlap = token_overlap_ratio(question, top_doc)

    hints = infer_source_hints(question)
    if not hints:
        source_signal = 0.6
    else:
        source_signal = (
            1.0 if source_matches_hints(top_doc.metadata.get("source", ""), hints) else 0.0
        )

    ranked_scores = sorted(vector_score_map.values(), reverse=True)
    top1 = ranked_scores[0] if ranked_scores else top_score
    top2 = ranked_scores[1] if len(ranked_scores) > 1 else 0.0
    margin = min(max(top1 - top2, 0.0) / 0.3, 1.0)

    confidence = (
        0.45 * top_score
        + 0.25 * source_signal
        + 0.20 * top_overlap
        + 0.10 * margin
    )
    return min(max(confidence, 0.0), 1.0)


def compute_image_confidence(question: str, image_hits: list[str]) -> float:
    """
    Returns a confidence score in [0, 1] for an image retrieval result.

    Weights:
      0.70 × normalised path score (capped at 60 points → 1.0)
      0.20 × hint signal
      0.10 × margin between top-1 and top-2 path scores
    """
    from retrieval.images import score_image_path  # avoid circular at module level

    if not image_hits:
        return 0.0

    normalized_question = rewrite_query_for_retrieval(question)
    hints = infer_source_hints(normalized_question)
    scores = [score_image_path(normalized_question, p) for p in image_hits]
    top1 = scores[0] if scores else 0.0
    top2 = scores[1] if len(scores) > 1 else 0.0
    normalized_score = min(top1 / 60.0, 1.0)
    margin = min(max(top1 - top2, 0.0) / 20.0, 1.0)

    if not hints:
        hint_signal = 0.6
    else:
        hint_signal = 1.0 if path_matches_hints(image_hits[0], hints) else 0.0

    confidence = (
        0.70 * normalized_score
        + 0.20 * hint_signal
        + 0.10 * margin
    )
    return min(max(confidence, 0.0), 1.0)