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+")


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]


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()]