hlamingo/hlamingo/search.py

"""Поиск предметов: по словам (с учётом русских окончаний) плюс семантика.

Порядок сигналов выбран по замерам на реальном эмбеддере, а не наугад.

1. СЛОВА — основной сигнал. Запрос и карточка предмета режутся на слова, слова
   грубо приводятся к основе («стабилизаторы» → «стабилизатор»), считается доля
   совпавших. Работает всегда, в том числе с выключенным Ollama.
2. СИНОНИМЫ от модели («изолента» → «изоляционная лента», «ПВХ-лента») лежат
   в самой карточке и расширяют словарь предмета — то есть усиливают пункт 1.
3. СЕМАНТИКА — дополнение, а не основа. Замер на nomic-embed-text показал, что
   на русском он почти не разделяет смыслы: «велосипедная камера» давала
   мультиметру 0.689, а верный ответ на «паяльник» — 0.691. Порогом такое не
   разделить, поэтому семантика подключается только когда по словам ничего не
   нашлось, с высоким порогом, и помечается в выдаче как догадка.

Хочется настоящего смыслового поиска — нужен многоязычный эмбеддер (bge-m3):
см. EMBED_MODEL в config.py и раздел про поиск в README.
"""

import re
import math
import logging

from . import ollama
from .config import TOP_K, MIN_SCORE

logger = logging.getLogger(__name__)

# Слова, которые есть в любом запросе и ничего не сообщают о предмете.
STOP = {
    "где", "что", "чем", "как", "куда", "мой", "моя", "моё", "мои", "мне",
    "в", "во", "на", "за", "под", "из", "от", "до", "по", "с", "со", "у", "к",
    "и", "а", "но", "или", "для", "же", "ли", "это", "тот", "та", "те",
    "лежит", "лежат", "находится", "хранится", "положил", "искать", "найти",
}

# Самые частые русские окончания, от длинных к коротким: обрезаем первое
# подходящее. Это не морфология, а дешёвая замена ей — но «стабилизаторы»,
# «стабилизатора» и «стабилизатором» после неё сходятся в одну основу.
_ENDINGS = (
    "ами", "ями", "ого", "ему", "ому", "ыми", "ими", "ах", "ях", "ов", "ев",
    "ам", "ям", "ой", "ей", "ые", "ий", "ый", "ая", "яя", "ое", "ее", "ии",
    "ом", "ем", "им", "ых", "их", "ую", "юю", "ть",
    "а", "я", "ы", "и", "о", "е", "у", "ю", "й", "ь",
)
_WORD_RE = re.compile(r"[0-9a-zA-Zа-яёА-ЯЁ]+")

MIN_WORD_SCORE = 0.5      # совпало меньше половины значимых слов — не показываем
SEMANTIC_MIN = 0.70       # порог догадки по смыслу (см. калибровку выше)


def stem(word):
    """Грубая основа слова. Короткие слова и коды («A3», «7805») не трогаем."""
    word = word.lower().replace("ё", "е")
    if len(word) <= 4 or word.isdigit():
        return word
    for ending in _ENDINGS:
        if word.endswith(ending) and len(word) - len(ending) >= 4:
            return word[:-len(ending)]
    return word


def tokens(text, keep_stop=False):
    """Значимые основы слов текста."""
    out = []
    for word in _WORD_RE.findall(text or ""):
        low = word.lower()
        if not keep_stop and low in STOP:
            continue
        out.append(stem(word))
    return out


def item_text(item):
    """Всё, по чему предмет можно узнать: название, синонимы, заметка."""
    parts = [item.get("name") or ""]
    parts += item.get("aliases") or []
    if item.get("note"):
        parts.append(item["note"])
    return ", ".join(p for p in parts if p)


def word_score(query, item):
    """Доля слов запроса, найденных в карточке предмета: 0…1."""
    q = tokens(query)
    if not q:
        return 0.0
    haystack = set(tokens(item_text(item)))
    if not haystack:
        return 0.0

    found = 0
    for word in q:
        if word in haystack:
            found += 1
            continue
        # Частичное совпадение: «стабилиз» ↔ «стабилизатор». Помогает там, где
        # обрезка окончаний сработала по-разному в запросе и в карточке.
        if len(word) >= 5 and any(h.startswith(word) or word.startswith(h)
                                  for h in haystack if len(h) >= 5):
            found += 0.8
    return min(1.0, found / len(q))


# ============================================================
# Семантика
# ============================================================
def cosine(a, b):
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


def ensure_embeddings(db):
    """Досчитывает векторы предметам без них. True, если базу надо сохранить."""
    changed = False
    for item in db["items"].values():
        if item.get("emb"):
            continue
        vector = ollama.embed(item_text(item))
        if not vector:
            return changed      # эмбеддер молчит — не долбим его на каждой записи
        item["emb"] = vector
        changed = True
    return changed


def reindex(db):
    """Полный пересчёт векторов — после смены модели эмбеддера.

    Векторы разных моделей несовместимы: смешанная база тихо испортит поиск,
    поэтому при смене EMBED_MODEL пересчитывать нужно всё сразу.
    """
    for item in db["items"].values():
        item["emb"] = None
    return ensure_embeddings(db)


# ============================================================
# Поиск
# ============================================================
def search(db, query, top_k=TOP_K, min_score=MIN_SCORE):
    """Список (item, score, why). why — «по названию» или «похоже по смыслу»."""
    if not (query or "").strip() or not db["items"]:
        return []

    scored = []
    for item in db["items"].values():
        score = word_score(query, item)
        if score >= MIN_WORD_SCORE:
            scored.append((item, score, "по названию"))
    scored.sort(key=lambda t: t[1], reverse=True)
    if scored:
        return scored[:top_k]

    # По словам не нашлось — пробуем догадаться по смыслу. Отдельно и с высоким
    # порогом: лучше честное «не нашёл», чем уверенно показанное не то.
    vector = ollama.embed(query, is_query=True)
    if not vector:
        return []
    guesses = []
    for item in db["items"].values():
        if not item.get("emb"):
            continue
        score = cosine(vector, item["emb"])
        if score >= max(SEMANTIC_MIN, min_score):
            guesses.append((item, score, "похоже по смыслу"))
    guesses.sort(key=lambda t: t[1], reverse=True)
    return guesses[:top_k]