hlamingo/hlamingo/brain.py

"""ИИ-слой: разбор свободной фразы в структуру и генерация синонимов.

Пользователь пишет одну строку — «положил стабилизаторы 7805, 12 штук, в коробку
A3 на второй полке стеллажа А в гараже». Модель превращает её в JSON, который
понимает store.py.

Ollama выключается вместе с компьютером, поэтому обе функции обязаны переживать
его отсутствие: разбор падает на регулярки, синонимы просто не генерируются.
"""

import re
import json
import logging

from . import ollama
from . import config

logger = logging.getLogger(__name__)

PARSE_SYSTEM = """Ты разбираешь фразы о том, куда человек положил вещи, в JSON.

Верни ТОЛЬКО JSON-объект, без пояснений и без markdown-ограждения:
{"action": "...", "items": ["..."], "qty": "...", "container": "...",
 "place_path": ["...", "..."], "note": "..."}

action:
  "put"            — человек кладёт или перекладывает вещь
  "find"           — спрашивает, где вещь лежит
  "list_container" — спрашивает, что лежит в конкретной коробке
  "list_place"     — спрашивает, что лежит в комнате/на полке/в шкафу
  "unknown"        — непонятно

items       — СПИСОК названий вещей в именительном падеже, единственном числе.
              Если названо несколько вещей — перечисли ВСЕ, ничего не теряй.
qty         — количество, если названо ("12 шт"), иначе null.
container   — код коробки: буква и цифра, например "A3". Если коробка не названа — null.
place_path  — путь от крупного к мелкому: ["Гараж","Стеллаж А","Полка 2"].
              Только реально названные места. Если мест нет — [].
note        — уточнение, если есть, иначе null.

ВАЖНО про перекладывание: во фразе вида «переложил X ИЗ A В B» нужно место
НАЗНАЧЕНИЯ — то есть B. Откуда вещь забрали — не записывай.

ВАЖНО про вопросы: фраза, начинающаяся с «где», «куда», «что», «не помню», —
это ВОПРОС, action="find", даже если внутри есть слово «положил».
«куда я положил зарядку» — человек СПРАШИВАЕТ, а не сообщает.
Признак записи: сказано, КУДА вещь легла. Нет места назначения — значит вопрос.

Примеры:
«положил стабилизаторы 7805, 12 штук, в коробку A3 на второй полке стеллажа А в гараже»
{"action":"put","items":["Стабилизатор 7805"],"qty":"12 шт","container":"A3","place_path":["Гараж","Стеллаж А","Полка 2"],"note":null}

«положил паспорт и права в верхний ящик стола в спальне»
{"action":"put","items":["паспорт","права"],"qty":null,"container":null,"place_path":["Спальня","Верхний ящик стола"],"note":null}

«переложил дрель из коробки A1 в коробку C2»
{"action":"put","items":["дрель"],"qty":null,"container":"C2","place_path":[],"note":null}

«где мультиметр»
{"action":"find","items":["мультиметр"],"qty":null,"container":null,"place_path":[],"note":null}

«куда я положил зарядку от шуруповёрта»
{"action":"find","items":["зарядка от шуруповёрта"],"qty":null,"container":null,"place_path":[],"note":null}

«что в коробке B7»
{"action":"list_container","items":[],"qty":null,"container":"B7","place_path":[],"note":null}

«лопата стоит в углу гаража»
{"action":"put","items":["лопата"],"qty":null,"container":null,"place_path":["Гараж"],"note":"в углу"}
"""

ALIAS_SYSTEM = """Ты придумываешь альтернативные названия вещи, чтобы её потом можно было найти.

Верни ТОЛЬКО JSON-массив строк, 3-7 штук, без пояснений.
Включай: обиходные названия, синонимы, маркировку и артикулы, расшифровку
сокращений, а для деталей — что это по функции.

«Стабилизатор 7805» → ["LM7805","стабилизатор 5 вольт","линейный регулятор напряжения","КР142ЕН5А","стабилизатор питания"]
«изолента» → ["изоляционная лента","синяя лента","ПВХ-лента","электроизоляционная лента"]
"""

_THINK_RE = re.compile(r"<think(?:ing)?>.*?</think(?:ing)?>", re.IGNORECASE | re.DOTALL)
_THINK_CLOSE_RE = re.compile(r"</think(?:ing)?>", re.IGNORECASE)
_THINK_OPEN_RE = re.compile(r"<think(?:ing)?>", re.IGNORECASE)


def _strip_thoughts(content):
    """Убирает <think>…</think>: модели-ризонеры пишут их прямо в content,
    и без чистки JSON не разбирается. Обрабатываем и «висячие» теги."""
    if not content:
        return content
    text = _THINK_RE.sub("", content)
    if _THINK_CLOSE_RE.search(text):
        text = _THINK_CLOSE_RE.split(text)[-1]
    if _THINK_OPEN_RE.search(text):
        text = _THINK_OPEN_RE.split(text)[0]
    return text.strip()


def _extract_json(content, opener="{", closer="}"):
    """Достаёт JSON из ответа модели, терпя обрамление и болтовню вокруг."""
    text = _strip_thoughts(content)
    if not text:
        return None
    fenced = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.DOTALL)
    if fenced:
        text = fenced.group(1).strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    match = re.search(re.escape(opener) + r".*" + re.escape(closer), text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except Exception:
        return None


# ============================================================
# Разбор фразы
# ============================================================
def parse(text, model=None):
    """Фраза → структура. Всегда возвращает словарь, никогда не бросает."""
    text = (text or "").strip()
    if not text:
        return _blank("unknown")
    model = model or config.CHAT_MODEL

    content = ollama.chat(model, [
        {"role": "system", "content": PARSE_SYSTEM},
        {"role": "user", "content": text},
    ])
    data = _extract_json(content) if content else None
    if isinstance(data, dict) and data.get("action") in (
            "put", "find", "list_container", "list_place", "unknown"):
        return _guard_question(_clean(data), text)

    if content:
        logger.info("Модель не дала разбираемый JSON, беру регулярки: %r", content[:200])
    return fallback_parse(text)


_QUESTION_RE = re.compile(r"^\s*(?:где|куда|что|чего|не\s+помню|подскажи|напомни)\b",
                          re.IGNORECASE)


def _guard_question(parsed, text):
    """Не даёт записать вопрос как факт.

    Модель цепляется за слово «положил» и превращает «куда я положил зарядку»
    в запись. Признак записи один: сказано, КУДА вещь легла. Нет ни коробки,
    ни места — значит человек спрашивает, и придумывать ему новое место нельзя.
    """
    if parsed.get("action") != "put":
        return parsed
    if parsed.get("container") or parsed.get("place_path"):
        return parsed
    if _QUESTION_RE.match(text or "") or (text or "").rstrip().endswith("?"):
        parsed["action"] = "find"
    return parsed


def _blank(action):
    return {"action": action, "items": [], "item": None, "qty": None,
            "container": None, "place_path": [], "note": None, "by": "regex"}


def _as_list(value):
    """Поле может прийти строкой или списком — приводим к списку строк."""
    if value is None:
        return []
    if isinstance(value, str):
        value = [value]
    if not isinstance(value, list):
        return []
    return [str(v).strip() for v in value if str(v).strip()]


def _clean(data):
    def s(key):
        value = data.get(key)
        if value is None:
            return None
        value = str(value).strip()
        return value or None

    # Старое поле item тоже принимаем: модель иногда сбивается на него,
    # и терять из-за этого запись нельзя.
    items = _as_list(data.get("items")) or _as_list(data.get("item"))
    return {
        "action": data.get("action"),
        "items": items,
        "item": items[0] if items else None,
        "qty": s("qty"),
        "container": s("container"),
        "place_path": _as_list(data.get("place_path")),
        "note": s("note"),
        "by": "model",
    }


# ============================================================
# Фолбэк без ИИ
# ============================================================
_PUT_RE = re.compile(
    r"^(?:я\s+)?(?:положил|положила|кладу|убрал|убрала|засунул|положу|переложил|"
    r"поставил|поставила|лежит|хранится)\s+(?P<item>.+?)"
    r"\s+(?:в|во|на|под|за)\s+(?P<place>.+)$", re.IGNORECASE)
_FIND_RE = re.compile(
    r"^(?:где|найди|найти|куда\s+я\s+положил|куда\s+положил|поищи|ищу)\s+"
    r"(?:мой|моя|моё|мои|the)?\s*(?P<item>.+?)[\?\.!]*$", re.IGNORECASE)
_LIST_BOX_RE = re.compile(
    r"^(?:что|чего|покажи)\s+(?:лежит\s+)?(?:в|во)\s+(?:коробке|ящике|боксе|)\s*"
    r"(?P<code>[A-ZА-Я]{1,2}\s?\d{1,3})[\?\.!]*$", re.IGNORECASE)
_CODE_RE = re.compile(r"\b(?:коробк\w*|ящик\w*|бокс\w*)\s*([A-ZА-Я]{1,2}\s?-?\s?\d{1,3})\b",
                      re.IGNORECASE)
_QTY_RE = re.compile(r"(\d+)\s*(шт\w*|штук\w*|метр\w*|м\b|кг\b|грамм\w*|пар\w*|упаковк\w*)",
                     re.IGNORECASE)


def fallback_parse(text):
    """Разбор регулярками — когда Ollama недоступен. Понимает простые фразы."""
    text = (text or "").strip()

    match = _LIST_BOX_RE.match(text)
    if match:
        result = _blank("list_container")
        result["container"] = match.group("code")
        return result

    match = _FIND_RE.match(text)
    if match:
        return _with_items(_blank("find"), [match.group("item").strip(" ?.!")])

    match = _PUT_RE.match(text)
    if match:
        result = _blank("put")
        item = match.group("item").strip()
        place = match.group("place").strip(" ?.!")

        # «переложил дрель ИЗ коробки A1 в коробку C2»: всё после «из» — это
        # место, откуда забрали. В название вещи оно попасть не должно.
        item = re.split(r"\s+из\s+", item, 1)[0].strip()

        qty = _QTY_RE.search(item)
        if qty:
            result["qty"] = qty.group(0).strip()
            item = _QTY_RE.sub("", item).strip(" ,")
        _with_items(result, re.split(r"\s*(?:,|\sи\s)\s*", item.strip(" ,")))

        code = _CODE_RE.search(place)
        if code:
            result["container"] = code.group(1)
            # Всё после кода коробки — описание места («A3 на полке 2 в гараже»).
            rest = place[code.end():].strip(" ,")
            result["place_path"] = _split_places(rest)
        else:
            result["place_path"] = _split_places(place)
        return result

    # Ничего не распознали: считаем, что это поиск — так безопаснее,
    # чем записать непонятую фразу в базу как предмет.
    return _with_items(_blank("find"), [text.strip(" ?.!")])


def _with_items(result, names):
    names = [n.strip(" ,.") for n in names if n and n.strip(" ,.")]
    result["items"] = names
    result["item"] = names[0] if names else None
    return result


def _split_places(text):
    """«на полке 2 в гараже» → ['гараже','полке 2'] (от крупного к мелкому).

    Предлоги идут от мелкого к крупному, поэтому части переворачиваем.
    """
    text = (text or "").strip(" ,.")
    if not text:
        return []
    parts = re.split(r"\s+(?:в|во|на|под|за)\s+", text)
    parts = [p.strip(" ,.") for p in parts if p.strip(" ,.")]
    parts.reverse()
    return parts


# ============================================================
# Синонимы
# ============================================================
def aliases(name, note=None, model=None):
    """3-7 альтернативных названий. Пустой список, если модель недоступна."""
    name = (name or "").strip()
    if not name:
        return []
    model = model or config.CHAT_MODEL
    prompt = name + (" (" + note + ")" if note else "")
    content = ollama.chat(model, [
        {"role": "system", "content": ALIAS_SYSTEM},
        {"role": "user", "content": prompt},
    ], num_predict=256)
    data = _extract_json(content, "[", "]") if content else None
    if not isinstance(data, list):
        return []
    out = []
    for value in data:
        value = str(value).strip()
        if value and value.lower() != name.lower() and value not in out:
            out.append(value)
    return out[:7]