call-center/services/shared/kb_search.py

from __future__ import annotations

import json
import re
import unicodedata
from typing import Protocol, Sequence, TypeVar

_FIELD_WEIGHTS = {
    "tags": 3.0,
    "title": 2.0,
    "body": 1.0,
}
_EXACT_TOKEN_SCORE = 10.0
_SOFT_TOKEN_SCORE = 6.0
_PHRASE_BASE_SCORE = 12.0
_MIN_TOKEN_LENGTH = 2
_MIN_SOFT_MATCH_LENGTH = 4

_NON_WORD_RE = re.compile(r"[^\w]+", re.UNICODE)
_SPACE_RE = re.compile(r"\s+")

# Common RU/KZ greetings, confirmations, pronouns, and particles that carry no
# topical signal on their own. Without this, a caller utterance as thin as
# "да" or "хорошо" could still exact-token-match some unrelated KB article
# that happens to contain that word in its body, and get returned as the
# top/only search result — read back to the caller as if it were the answer
# to their question. Filtering these keeps _score_row's exact-token match
# meaningful: a match now requires an actual content word.
_STOPWORDS = frozenset(
    {
        # RU: greetings / confirmations / fillers
        "алло", "ага", "да", "неа", "нет", "ой", "ок", "окей", "угу", "ясно",
        "ладно", "хорошо", "понял", "поняла", "привет", "здравствуйте",
        "добрый", "день", "вечер", "утро", "слышу", "слышно", "спасибо",
        "пожалуйста", "извините", "простите", "алло",
        # RU: pronouns / conjunctions / particles with no topical content
        "я", "ты", "вы", "мы", "он", "она", "они", "это", "то", "и", "а",
        "но", "или", "что", "как", "где", "когда", "если", "чтобы", "для",
        "из", "по", "на", "в", "с", "у", "о", "же", "ли", "бы", "не", "ну",
        "вот", "просто", "есть", "быть", "можно", "нужно", "надо", "уже",
        # KZ: greetings / confirmations / fillers
        "иә", "ия", "жоқ", "жарайды", "түсінікті", "рахмет", "сәлем",
        "сәлеметсіз", "бе", "кешіріңіз",
        # KZ: pronouns / conjunctions / particles
        "мен", "сен", "сіз", "біз", "ол", "олар", "және", "бірақ", "немесе",
        "не", "қалай", "қайда", "қашан", "үшін", "туралы",
    }
)


class KBSearchRow(Protocol):
    id: int
    title: str
    body: str
    tags_json: str
    intent_code: str | None


T = TypeVar("T", bound=KBSearchRow)


def normalize_kb_text(value: str | None) -> str:
    text = unicodedata.normalize("NFKC", str(value or "")).lower().replace("\u0451", "\u0435")
    text = _NON_WORD_RE.sub(" ", text)
    return _SPACE_RE.sub(" ", text).strip()


def tokenize_kb_text(value: str | None) -> list[str]:
    return [
        token
        for token in normalize_kb_text(value).split(" ")
        if len(token) >= _MIN_TOKEN_LENGTH and token not in _STOPWORDS
    ]


def search_kb_rows(
    rows: Sequence[T],
    query: str | None,
    *,
    limit: int,
    empty_query_returns_all: bool = False,
) -> list[T]:
    ordered_rows = sorted(rows, key=lambda row: int(getattr(row, "id", 0) or 0), reverse=True)
    normalized_query = normalize_kb_text(query)
    if not normalized_query:
        return ordered_rows[:limit] if empty_query_returns_all else []

    query_tokens = list(dict.fromkeys(tokenize_kb_text(normalized_query)))
    if not query_tokens:
        return ordered_rows[:limit] if empty_query_returns_all else []

    ranked: list[tuple[float, int, T]] = []
    for row in ordered_rows:
        score = _score_row(row, normalized_query, query_tokens)
        if score <= 0:
            continue
        ranked.append((score, int(getattr(row, "id", 0) or 0), row))

    ranked.sort(key=lambda item: (-item[0], -item[1]))
    return [row for _, _, row in ranked[:limit]]


def _score_row(row: T, normalized_query: str, query_tokens: list[str]) -> float:
    tags_text = " ".join(_parse_tags(getattr(row, "tags_json", "[]")))
    fields = {
        "tags": {
            "text": normalize_kb_text(tags_text),
            "tokens": tokenize_kb_text(tags_text),
        },
        "title": {
            "text": normalize_kb_text(getattr(row, "title", "")),
            "tokens": tokenize_kb_text(getattr(row, "title", "")),
        },
        "body": {
            "text": normalize_kb_text(getattr(row, "body", "")),
            "tokens": tokenize_kb_text(getattr(row, "body", "")),
        },
    }

    score = 0.0
    for field_name, payload in fields.items():
        field_tokens = payload["tokens"]
        if not field_tokens:
            continue
        field_weight = _FIELD_WEIGHTS[field_name]
        score += _field_phrase_bonus(query_tokens, payload["text"], field_weight)
        score += _field_token_score(query_tokens, field_tokens, field_weight)
    return score


def _field_phrase_bonus(query_tokens: list[str], field_text: str, field_weight: float) -> float:
    max_window = min(4, len(query_tokens))
    for size in range(max_window, 1, -1):
        for start in range(0, len(query_tokens) - size + 1):
            phrase = " ".join(query_tokens[start : start + size])
            if phrase and phrase in field_text:
                return field_weight * (_PHRASE_BASE_SCORE + size)
    return 0.0


def _field_token_score(query_tokens: list[str], field_tokens: list[str], field_weight: float) -> float:
    score = 0.0
    field_token_set = set(field_tokens)
    for query_token in query_tokens:
        if query_token in field_token_set:
            score += field_weight * _EXACT_TOKEN_SCORE
            continue
        if _has_soft_token_match(query_token, field_tokens):
            score += field_weight * _SOFT_TOKEN_SCORE
    return score


def _has_soft_token_match(query_token: str, field_tokens: list[str]) -> bool:
    return any(_is_soft_token_match(query_token, field_token) for field_token in field_tokens)


def _is_soft_token_match(query_token: str, field_token: str) -> bool:
    min_length = min(len(query_token), len(field_token))
    if min_length < _MIN_SOFT_MATCH_LENGTH:
        return False
    common_prefix = _common_prefix_len(query_token, field_token)
    return common_prefix >= _MIN_SOFT_MATCH_LENGTH and common_prefix >= min_length - 1


def _common_prefix_len(left: str, right: str) -> int:
    index = 0
    for left_char, right_char in zip(left, right):
        if left_char != right_char:
            break
        index += 1
    return index


def _parse_tags(raw_tags: str | None) -> list[str]:
    try:
        data = json.loads(raw_tags or "[]")
    except json.JSONDecodeError:
        return []
    return [str(item) for item in data if str(item).strip()]