test-search-engine/app/search_index.py

from __future__ import annotations

import math
from collections import Counter, defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING

from app.normalization import meaningful_tokens, normalize_text, tokenize, word_tokens
from app.typo import character_trigrams, trigram_dice, word_similarity

if TYPE_CHECKING:
    from app.catalog import CatalogItem


@dataclass(frozen=True, slots=True)
class LexicalScore:
    item_index: int
    score: float
    bm25: float
    trigram: float
    token_alignment: float
    exact_coverage: float


class LexicalIndex:
    """Small in-memory lexical index: BM25 + char trigrams + typo alignment."""

    def __init__(self, items: tuple[CatalogItem, ...]) -> None:
        self._items = items
        document_tokens: list[tuple[str, ...]] = []
        document_words: list[tuple[str, ...]] = []
        term_frequencies: list[Counter[str]] = []
        document_frequency: Counter[str] = Counter()
        trigram_postings: dict[str, set[int]] = defaultdict(set)

        for index, item in enumerate(items):
            text = f"{item.normalized_name} {item.product_type}"
            tokens = meaningful_tokens(tokenize(text, already_normalized=True))
            words = word_tokens(tokens)
            frequencies = Counter(tokens)
            document_tokens.append(tokens)
            document_words.append(words)
            term_frequencies.append(frequencies)
            document_frequency.update(frequencies.keys())
            for trigram in character_trigrams(" ".join(tokens)):
                trigram_postings[trigram].add(index)

        self._document_tokens = tuple(document_tokens)
        self._document_words = tuple(document_words)
        self._term_frequencies = tuple(term_frequencies)
        self._document_lengths = tuple(len(tokens) for tokens in document_tokens)
        self._average_document_length = max(
            1.0, sum(self._document_lengths) / max(1, len(self._document_lengths))
        )
        document_count = len(items)
        self._idf = {
            term: math.log(1.0 + (document_count - frequency + 0.5) / (frequency + 0.5))
            for term, frequency in document_frequency.items()
        }
        self._trigram_postings = {
            trigram: frozenset(indexes) for trigram, indexes in trigram_postings.items()
        }

    def trigram_candidates(self, query_text: str, limit: int = 80) -> tuple[int, ...]:
        normalized = normalize_text(query_text)
        query_tokens = meaningful_tokens(tokenize(normalized, already_normalized=True))
        counts: Counter[int] = Counter()
        for trigram in character_trigrams(" ".join(query_tokens)):
            counts.update(self._trigram_postings.get(trigram, ()))
        return tuple(index for index, _ in counts.most_common(limit))

    def rank(
        self,
        query_text: str,
        product_type: str,
        candidate_indices: tuple[int, ...],
    ) -> list[LexicalScore]:
        normalized = normalize_text(query_text)
        query_tokens = list(meaningful_tokens(tokenize(normalized, already_normalized=True)))
        # Alias normalization adds the canonical type as a lexical signal but
        # leaves the original wording intact for typo evidence.
        query_tokens.extend(meaningful_tokens(tokenize(product_type)))
        query_token_tuple = tuple(query_tokens)
        query_words = word_tokens(query_token_tuple)
        query_counter = Counter(query_token_tuple)
        query_trigram_text = " ".join(query_token_tuple)

        raw_bm25: dict[int, float] = {}
        for item_index in candidate_indices:
            raw_bm25[item_index] = self._bm25(item_index, query_counter)
        maximum_bm25 = max(raw_bm25.values(), default=0.0)

        scores: list[LexicalScore] = []
        for item_index in candidate_indices:
            item_tokens = self._document_tokens[item_index]
            item_words = self._document_words[item_index]
            bm25 = raw_bm25[item_index] / maximum_bm25 if maximum_bm25 > 0 else 0.0
            trigram = trigram_dice(query_trigram_text, " ".join(item_tokens))
            alignment = self._token_alignment(query_words, item_words)
            exact = self._exact_coverage(query_token_tuple, item_tokens)
            combined = bm25 * 0.42 + trigram * 0.23 + alignment * 0.25 + exact * 0.10
            scores.append(
                LexicalScore(
                    item_index=item_index,
                    score=max(0.0, min(1.0, combined)),
                    bm25=bm25,
                    trigram=trigram,
                    token_alignment=alignment,
                    exact_coverage=exact,
                )
            )

        scores.sort(key=lambda item: (-item.score, self._items[item.item_index].sku))
        return scores

    def _bm25(self, item_index: int, query: Counter[str]) -> float:
        k1 = 1.5
        b = 0.75
        frequencies = self._term_frequencies[item_index]
        document_length = self._document_lengths[item_index]
        score = 0.0
        for term, query_frequency in query.items():
            term_frequency = frequencies.get(term, 0)
            if term_frequency == 0:
                continue
            denominator = term_frequency + k1 * (
                1.0 - b + b * document_length / self._average_document_length
            )
            score += (
                self._idf.get(term, 0.0)
                * term_frequency
                * (k1 + 1.0)
                / denominator
                * min(query_frequency, 2)
            )
        return score

    @staticmethod
    def _token_alignment(
        query_words: tuple[str, ...], document_words: tuple[str, ...]
    ) -> float:
        if not query_words:
            return 0.0
        values: list[float] = []
        for query_word in query_words:
            if query_word in document_words:
                values.append(1.0)
                continue
            if len(query_word) < 4:
                values.append(0.0)
                continue
            best = max(
                (word_similarity(query_word, document_word) for document_word in document_words),
                default=0.0,
            )
            values.append(best if best >= 0.50 else 0.0)
        return sum(values) / len(values)

    @staticmethod
    def _exact_coverage(query_tokens: tuple[str, ...], document_tokens: tuple[str, ...]) -> float:
        if not query_tokens:
            return 0.0
        document_set = set(document_tokens)
        return sum(token in document_set for token in query_tokens) / len(query_tokens)