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